This commit is contained in:
Manav Rathi 2024-05-14 15:55:48 +05:30
parent 613324a4ae
commit 75814415ce
No known key found for this signature in database
2 changed files with 67 additions and 65 deletions

View file

@ -84,7 +84,10 @@ import {
getSectionSummaries,
} from "services/collectionService";
import downloadManager from "services/download";
import { syncEmbeddings, syncFileEmbeddings } from "services/embeddingService";
import {
syncCLIPEmbeddings,
syncFileEmbeddings,
} from "services/embeddingService";
import { syncEntities } from "services/entityService";
import { getLocalFiles, syncFiles } from "services/fileService";
import locationSearchService from "services/locationSearchService";
@ -698,7 +701,7 @@ export default function Gallery() {
await syncTrash(collections, setTrashedFiles);
await syncEntities();
await syncMapEnabled();
await syncEmbeddings();
await syncCLIPEmbeddings();
const electron = globalThis.electron;
if (electron) {
await syncFileEmbeddings();

View file

@ -69,8 +69,11 @@ const setModelEmbeddingSyncTime = async (
await localForage.setItem(`${model}-${embeddingSyncTimeLSKeySuffix}`, time);
};
export const syncEmbeddings = async () => {
const models: EmbeddingModel[] = ["onnx-clip"];
/**
* Sync our locally available CLIP embeddings with the server.
*/
export const syncCLIPEmbeddings = async () => {
const model: EmbeddingModel = "onnx-clip";
try {
let allEmbeddings = await storedCLIPEmbeddings();
const localFiles = await getAllLocalFiles();
@ -87,69 +90,65 @@ export const syncEmbeddings = async () => {
clipEmbeddingsLSKey,
);
log.info(`Syncing embeddings localCount: ${allEmbeddings.length}`);
for (const model of models) {
let modelLastSinceTime = await getModelEmbeddingSyncTime(model);
log.info(
`Syncing ${model} model's embeddings sinceTime: ${modelLastSinceTime}`,
);
let response: GetEmbeddingDiffResponse;
do {
response = await getEmbeddingsDiff(modelLastSinceTime, model);
if (!response.diff?.length) {
return;
}
const newEmbeddings = await Promise.all(
response.diff.map(async (embedding) => {
try {
const {
encryptedEmbedding,
decryptionHeader,
...rest
} = embedding;
const worker =
await ComlinkCryptoWorker.getInstance();
const fileKey = fileIdToKeyMap.get(
embedding.fileID,
);
if (!fileKey) {
throw Error(CustomError.FILE_NOT_FOUND);
}
const decryptedData = await worker.decryptEmbedding(
encryptedEmbedding,
decryptionHeader,
fileIdToKeyMap.get(embedding.fileID),
);
return {
...rest,
embedding: decryptedData,
} as Embedding;
} catch (e) {
let hasHiddenAlbums = false;
if (e.message === CustomError.FILE_NOT_FOUND) {
hasHiddenAlbums = hiddenAlbums?.length > 0;
}
log.error(
`decryptEmbedding failed for file (hasHiddenAlbums: ${hasHiddenAlbums})`,
e,
);
let modelLastSinceTime = await getModelEmbeddingSyncTime(model);
log.info(
`Syncing ${model} model's embeddings sinceTime: ${modelLastSinceTime}`,
);
let response: GetEmbeddingDiffResponse;
do {
response = await getEmbeddingsDiff(modelLastSinceTime, model);
if (!response.diff?.length) {
return;
}
const newEmbeddings = await Promise.all(
response.diff.map(async (embedding) => {
try {
const {
encryptedEmbedding,
decryptionHeader,
...rest
} = embedding;
const worker = await ComlinkCryptoWorker.getInstance();
const fileKey = fileIdToKeyMap.get(embedding.fileID);
if (!fileKey) {
throw Error(CustomError.FILE_NOT_FOUND);
}
}),
);
allEmbeddings = getLatestVersionEmbeddings([
...allEmbeddings,
...newEmbeddings,
]);
if (response.diff.length) {
modelLastSinceTime = response.diff.slice(-1)[0].updatedAt;
}
await localForage.setItem(clipEmbeddingsLSKey, allEmbeddings);
await setModelEmbeddingSyncTime(model, modelLastSinceTime);
log.info(
`Syncing embeddings syncedEmbeddingsCount: ${allEmbeddings.length}`,
);
} while (response.diff.length === DIFF_LIMIT);
}
const decryptedData = await worker.decryptEmbedding(
encryptedEmbedding,
decryptionHeader,
fileIdToKeyMap.get(embedding.fileID),
);
return {
...rest,
embedding: decryptedData,
} as Embedding;
} catch (e) {
let hasHiddenAlbums = false;
if (e.message === CustomError.FILE_NOT_FOUND) {
hasHiddenAlbums = hiddenAlbums?.length > 0;
}
log.error(
`decryptEmbedding failed for file (hasHiddenAlbums: ${hasHiddenAlbums})`,
e,
);
}
}),
);
allEmbeddings = getLatestVersionEmbeddings([
...allEmbeddings,
...newEmbeddings,
]);
if (response.diff.length) {
modelLastSinceTime = response.diff.slice(-1)[0].updatedAt;
}
await localForage.setItem(clipEmbeddingsLSKey, allEmbeddings);
await setModelEmbeddingSyncTime(model, modelLastSinceTime);
log.info(
`Syncing embeddings syncedEmbeddingsCount: ${allEmbeddings.length}`,
);
} while (response.diff.length === DIFF_LIMIT);
} catch (e) {
log.error("Sync embeddings failed", e);
}