From f2f7b483fdd471932bd787f35663a2c923bed632 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 09:40:08 +0530 Subject: [PATCH] comlink wip --- desktop/package.json | 1 + desktop/src/main/ipc.ts | 2 +- desktop/src/main/services/ml-worker.ts | 34 +++++++++++++------ desktop/yarn.lock | 5 +++ web/packages/new/photos/services/ml/index.ts | 11 ++---- web/packages/new/photos/services/ml/worker.ts | 11 ++++-- 6 files changed, 42 insertions(+), 22 deletions(-) diff --git a/desktop/package.json b/desktop/package.json index 453bb931fa..c3da2c3591 100644 --- a/desktop/package.json +++ b/desktop/package.json @@ -28,6 +28,7 @@ "auto-launch": "^5.0", "chokidar": "^3.6", "clip-bpe-js": "^0.0.6", + "comlink": "^4.4.1", "compare-versions": "^6.1", "electron-log": "^5.1", "electron-store": "^8.2", diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index caca5758b5..511460c6c7 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -49,7 +49,7 @@ import { computeCLIPTextEmbeddingIfAvailable, computeFaceEmbeddings, detectFaces, -} from "./services/ml-utility"; +} from "./services/ml-worker"; import { encryptionKey, lastShownChangelogVersion, diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index 255dbd3f5b..21553da3ec 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -9,6 +9,10 @@ // cannot import. import Tokenizer from "clip-bpe-js"; +import { expose } from "comlink"; +// eslint-disable-next-line @typescript-eslint/ban-ts-comment +// @ts-ignore +import nodeEndpoint from "comlink/dist/umd/node-adapter"; import { net } from "electron/main"; import { existsSync } from "fs"; import fs from "node:fs/promises"; @@ -60,12 +64,22 @@ process.parentPort.once("message", (e) => { parseInitData(e.data); const port = ensure(e.ports[0]); - port.on("message", (request) => { - void handleMessageFromRenderer(request.data).then((response) => - port.postMessage(response), - ); - }); - port.start(); + expose( + { + computeCLIPImageEmbedding, + computeCLIPTextEmbeddingIfAvailable, + detectFaces, + computeFaceEmbeddings, + }, + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument, @typescript-eslint/no-explicit-any + nodeEndpoint(port as unknown as any), + ); + // port.on("message", (request) => { + // void handleMessageFromRenderer(request.data).then((response) => + // port.postMessage(response), + // ); + // }); + // port.start(); }); /** @@ -109,15 +123,15 @@ const parseInitData = (data: unknown) => { * "result" (arbitrary result) property. Otherwise it will have a "error" * (string) property describing what went wrong. */ -const handleMessageFromRenderer = async (m: unknown) => { +export const handleMessageFromRenderer = (m: unknown) => { if (m && typeof m == "object" && "method" in m && "id" in m && "p" in m) { const id = m.id; - const p = m.p; + // const p = m.p; try { switch (m.method) { case "foo": - if (p && typeof p == "string") - return { id, result: await foo(p) }; + // if (p && typeof p == "string") + // return { id, result: await foo(p) }; break; } } catch (e) { diff --git a/desktop/yarn.lock b/desktop/yarn.lock index 5feaf65f6f..afbe850a91 100644 --- a/desktop/yarn.lock +++ b/desktop/yarn.lock @@ -968,6 +968,11 @@ combined-stream@^1.0.8: dependencies: delayed-stream "~1.0.0" +comlink@^4.4.1: + version "4.4.1" + resolved "https://registry.yarnpkg.com/comlink/-/comlink-4.4.1.tgz#e568b8e86410b809e8600eb2cf40c189371ef981" + integrity sha512-+1dlx0aY5Jo1vHy/tSsIGpSkN4tS9rZSW8FIhG0JH/crs9wwweswIo/POr451r7bZww3hFbPAKnTpimzL/mm4Q== + commander@^5.0.0: version "5.1.0" resolved "https://registry.yarnpkg.com/commander/-/commander-5.1.0.tgz#46abbd1652f8e059bddaef99bbdcb2ad9cf179ae" diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index c067006cf2..b903e0188c 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -12,7 +12,7 @@ import { FileType } from "@/media/file-type"; import type { EnteFile } from "@/new/photos/types/file"; import { ensure } from "@/utils/ensure"; import { throttled } from "@/utils/promise"; -import { proxy } from "comlink"; +import { proxy, transfer } from "comlink"; import { isInternalUser } from "../feature-flags"; import { getRemoteFlag, updateRemoteFlag } from "../remote-store"; import type { UploadItem } from "../upload/types"; @@ -59,11 +59,6 @@ const worker = async () => { const createComlinkWorker = async () => { const electron = ensureElectron(); - const mlWorkerElectron = { - detectFaces: electron.detectFaces, - computeFaceEmbeddings: electron.computeFaceEmbeddings, - computeCLIPImageEmbedding: electron.computeCLIPImageEmbedding, - }; const delegate = { workerDidProcessFile, }; @@ -78,9 +73,9 @@ const createComlinkWorker = async () => { await cw.remote.then((w) => { // Pass the message port to our web worker. - cw.worker.postMessage("createMLWorker/port", [messagePort]); + // cw.worker.postMessage("createMLWorker/port", [messagePort]); // Initialize it. - return w.init(proxy(mlWorkerElectron), proxy(delegate)); + return w.init(transfer(messagePort, [messagePort]), proxy(delegate)); }); return cw; diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 03b8adf2a2..d27f3021f4 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -8,7 +8,7 @@ import { fileLogID } from "@/new/photos/utils/file"; import { ensure } from "@/utils/ensure"; import { wait } from "@/utils/promise"; import { DOMParser } from "@xmldom/xmldom"; -import { expose } from "comlink"; +import { expose, wrap } from "comlink"; import downloadManager from "../download"; import { cmpNewLib2, extractRawExif } from "../exif"; import { getAllLocalFiles, getLocalTrashedFiles } from "../files"; @@ -95,8 +95,13 @@ export class MLWorker { * @param delegate The {@link MLWorkerDelegate} the worker can use to inform * the main thread of interesting events. */ - async init(electron: MLWorkerElectron, delegate?: MLWorkerDelegate) { - this.electron = electron; + async init(port: MessagePort, delegate: MLWorkerDelegate) { + // this.electron = electron; + this.electron = wrap(port); /* mlWorkerElectron = { + detectFaces: electron.detectFaces, + computeFaceEmbeddings: electron.computeFaceEmbeddings, + computeCLIPImageEmbedding: electron.computeCLIPImageEmbedding, + };*/ this.delegate = delegate; // Initialize the downloadManager running in the web worker with the // user's token. It'll be used to download files to index if needed.