comlink wip

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",

desktop/src/main/ipc.ts

@@ -49,7 +49,7 @@ import {
     computeCLIPTextEmbeddingIfAvailable,
     computeFaceEmbeddings,
     detectFaces,
-} from "./services/ml-utility";
+} from "./services/ml-worker";
 import {
     encryptionKey,
     lastShownChangelogVersion,

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) {

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"

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;

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<MLWorkerElectron>(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.