From 82f808e5334ca2cd2995ae0ba28c1bc4a2d1aba5 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 09:59:33 +0530 Subject: [PATCH 01/48] Outline --- desktop/src/main/services/ml.ts | 60 ++++++++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 4 deletions(-) diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 55bb8d79c2..dcda511df6 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -5,10 +5,6 @@ * * The ML runtime we use for inference is [ONNX](https://onnxruntime.ai). Models * for various tasks are not shipped with the app but are downloaded on demand. - * - * The primary reason for doing these tasks in the Node.js layer is so that we - * can use the binary ONNX runtime which is 10-20x faster than the WASM based - * web one. */ import { app, net } from "electron/main"; @@ -19,6 +15,62 @@ import * as ort from "onnxruntime-node"; import log from "../log"; import { writeStream } from "../stream"; +/** + * Create a new ML session. + * + * [Note: ML IPC] + * + * The primary reason for doing ML tasks in the Node.js layer is so that we can + * use the binary ONNX runtime, which is 10-20x faster than the WASM one that + * can be used directly on the web layer. + * + * For this to work, the main and renderer process need to communicate with each + * other. Further, in the web layer the ML indexing runs in a web worker (so as + * to not get in the way of the main thread). So the communication has 2 hops: + * + * Node.js main <-> Renderer main <-> Renderer web worker + * + * This naive way works, but has a problem. The Node.js main process is in the + * code path for delivering user events to the renderer process. The ML tasks we + * do take in the order of 100-300 ms (possibly more) for each individual + * inference. Thus, the Node.js main process is busy for those 100-300 ms, and + * does not forward events to the renderer, causing the UI to jitter. + * + * The solution for this is to spawn an Electron UtilityProcess, which we can + * think of a regular Node.js child process. This frees up the Node.js main + * process, and would remove the jitter. + * https://www.electronjs.org/docs/latest/tutorial/process-model + * + * It would seem that this introduces another hop in our IPC + * + * Node.js utility process <-> Node.js main <-> ... + * + * but here we can use the special bit about Electron utility processes that + * separates them from regular Node.js child processes: their support for + * message ports. https://www.electronjs.org/docs/latest/tutorial/message-ports + * + * As a brief summary, a MessagePort is a web feature that allows two contexts + * to communicate. A pair of message ports is called a message channel. The cool + * thing about these is that we can pass these ports themselves over IPC. + * + * So we + * + * 1. Spawn a utility process. + * 2. In the utility process create a message channel. + * 3. Keep one port of the pair with us, and send the other over IPC to the + * _web worker_ that is coordinating the ML indexing on the web layer. + * + * Thereafter, the utility process and web worker can directly talk to each + * other! + * + * Node.js utility process <-> Renderer web worker + * + */ +export const createMLSession = () => { + // }: Promise => { + throw new Error("Not implemented"); +}; + /** * Return a function that can be used to trigger a download of the specified * model, and the creating of an ONNX inference session initialized using it. From 65d2bfe1c1f79c6971c53d8a8a458db84feec4a9 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 10:19:03 +0530 Subject: [PATCH 02/48] Split on the main/utility axis --- desktop/src/main/ipc.ts | 5 +- desktop/src/main/services/ml-clip.ts | 68 ------- desktop/src/main/services/ml-face.ts | 53 ------ desktop/src/main/services/ml-utility.ts | 235 ++++++++++++++++++++++++ desktop/src/main/services/ml.ts | 136 +------------- 5 files changed, 248 insertions(+), 249 deletions(-) delete mode 100644 desktop/src/main/services/ml-clip.ts delete mode 100644 desktop/src/main/services/ml-face.ts create mode 100644 desktop/src/main/services/ml-utility.ts diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index 641ce9963d..37ee0478e5 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -45,8 +45,9 @@ import { logout } from "./services/logout"; import { computeCLIPImageEmbedding, computeCLIPTextEmbeddingIfAvailable, -} from "./services/ml-clip"; -import { computeFaceEmbeddings, detectFaces } from "./services/ml-face"; + computeFaceEmbeddings, + detectFaces, +} from "./services/ml-utility"; import { encryptionKey, lastShownChangelogVersion, diff --git a/desktop/src/main/services/ml-clip.ts b/desktop/src/main/services/ml-clip.ts deleted file mode 100644 index cea1d667b5..0000000000 --- a/desktop/src/main/services/ml-clip.ts +++ /dev/null @@ -1,68 +0,0 @@ -/** - * @file Compute CLIP embeddings for images and text. - * - * The embeddings are computed using ONNX runtime, with CLIP as the model. - */ - -import Tokenizer from "clip-bpe-js"; -import * as ort from "onnxruntime-node"; -import log from "../log"; -import { ensure, wait } from "../utils/common"; -import { makeCachedInferenceSession } from "./ml"; - -const cachedCLIPImageSession = makeCachedInferenceSession( - "clip-image-vit-32-float32.onnx", - 351468764 /* 335.2 MB */, -); - -export const computeCLIPImageEmbedding = async (input: Float32Array) => { - const session = await cachedCLIPImageSession(); - const t = Date.now(); - const feeds = { - input: new ort.Tensor("float32", input, [1, 3, 224, 224]), - }; - const results = await session.run(feeds); - log.debug(() => `ONNX/CLIP image embedding took ${Date.now() - t} ms`); - /* Need these model specific casts to type the result */ - return ensure(results.output).data as Float32Array; -}; - -const cachedCLIPTextSession = makeCachedInferenceSession( - "clip-text-vit-32-uint8.onnx", - 64173509 /* 61.2 MB */, -); - -let _tokenizer: Tokenizer | undefined; -const getTokenizer = () => { - if (!_tokenizer) _tokenizer = new Tokenizer(); - return _tokenizer; -}; - -export const computeCLIPTextEmbeddingIfAvailable = async (text: string) => { - const sessionOrSkip = await Promise.race([ - cachedCLIPTextSession(), - // Wait for a tick to get the session promise to resolved the first time - // this code runs on each app start (and the model has been downloaded). - wait(0).then(() => 1), - ]); - - // Don't wait for the download to complete. - if (typeof sessionOrSkip == "number") { - log.info( - "Ignoring CLIP text embedding request because model download is pending", - ); - return undefined; - } - - const session = sessionOrSkip; - const t = Date.now(); - const tokenizer = getTokenizer(); - const tokenizedText = Int32Array.from(tokenizer.encodeForCLIP(text)); - const feeds = { - input: new ort.Tensor("int32", tokenizedText, [1, 77]), - }; - - const results = await session.run(feeds); - log.debug(() => `ONNX/CLIP text embedding took ${Date.now() - t} ms`); - return ensure(results.output).data as Float32Array; -}; diff --git a/desktop/src/main/services/ml-face.ts b/desktop/src/main/services/ml-face.ts deleted file mode 100644 index 33c09efaa2..0000000000 --- a/desktop/src/main/services/ml-face.ts +++ /dev/null @@ -1,53 +0,0 @@ -/** - * @file Various face recognition related tasks. - * - * - Face detection with the YOLO model. - * - Face embedding with the MobileFaceNet model. - * - * The runtime used is ONNX. - */ - -import * as ort from "onnxruntime-node"; -import log from "../log"; -import { ensure } from "../utils/common"; -import { makeCachedInferenceSession } from "./ml"; - -const cachedFaceDetectionSession = makeCachedInferenceSession( - "yolov5s_face_640_640_dynamic.onnx", - 30762872 /* 29.3 MB */, -); - -export const detectFaces = async (input: Float32Array) => { - const session = await cachedFaceDetectionSession(); - const t = Date.now(); - const feeds = { - input: new ort.Tensor("float32", input, [1, 3, 640, 640]), - }; - const results = await session.run(feeds); - log.debug(() => `ONNX/YOLO face detection took ${Date.now() - t} ms`); - return ensure(results.output).data; -}; - -const cachedFaceEmbeddingSession = makeCachedInferenceSession( - "mobilefacenet_opset15.onnx", - 5286998 /* 5 MB */, -); - -export const computeFaceEmbeddings = async (input: Float32Array) => { - // Dimension of each face (alias) - const mobileFaceNetFaceSize = 112; - // Smaller alias - const z = mobileFaceNetFaceSize; - // Size of each face's data in the batch - const n = Math.round(input.length / (z * z * 3)); - const inputTensor = new ort.Tensor("float32", input, [n, z, z, 3]); - - const session = await cachedFaceEmbeddingSession(); - const t = Date.now(); - const feeds = { img_inputs: inputTensor }; - const results = await session.run(feeds); - log.debug(() => `ONNX/MFNT face embedding took ${Date.now() - t} ms`); - /* Need these model specific casts to extract and type the result */ - return (results.embeddings as unknown as Record) - .cpuData as Float32Array; -}; diff --git a/desktop/src/main/services/ml-utility.ts b/desktop/src/main/services/ml-utility.ts new file mode 100644 index 0000000000..79d39edea4 --- /dev/null +++ b/desktop/src/main/services/ml-utility.ts @@ -0,0 +1,235 @@ +/** + * @file ML related tasks. This code runs in a utility process. + * + * The ML runtime we use for inference is [ONNX](https://onnxruntime.ai). Models + * for various tasks are not shipped with the app but are downloaded on demand. + */ + +import Tokenizer from "clip-bpe-js"; +import { app, net } from "electron/main"; +import { existsSync } from "fs"; +import fs from "node:fs/promises"; +import path from "node:path"; +import * as ort from "onnxruntime-node"; +import log from "../log"; +import { writeStream } from "../stream"; +import { ensure, wait } from "../utils/common"; + +/** + * Return a function that can be used to trigger a download of the specified + * model, and the creating of an ONNX inference session initialized using it. + * + * Multiple parallel calls to the returned function are fine, it ensures that + * the the model will be downloaded and the session created using it only once. + * All pending calls to it meanwhile will just await on the same promise. + * + * And once the promise is resolved, the create ONNX inference session will be + * cached, so subsequent calls to the returned function will just reuse the same + * session. + * + * {@link makeCachedInferenceSession} can itself be called anytime, it doesn't + * actively trigger a download until the returned function is called. + * + * @param modelName The name of the model to download. + * + * @param modelByteSize The size in bytes that we expect the model to have. If + * the size of the downloaded model does not match the expected size, then we + * will redownload it. + * + * @returns a function. calling that function returns a promise to an ONNX + * session. + */ +const makeCachedInferenceSession = ( + modelName: string, + modelByteSize: number, +) => { + let session: Promise | undefined; + + const download = () => + modelPathDownloadingIfNeeded(modelName, modelByteSize); + + const createSession = (modelPath: string) => + createInferenceSession(modelPath); + + const cachedInferenceSession = () => { + if (!session) session = download().then(createSession); + return session; + }; + + return cachedInferenceSession; +}; + +/** + * Download the model named {@link modelName} if we don't already have it. + * + * Also verify that the size of the model we get matches {@expectedByteSize} (if + * not, redownload it). + * + * @returns the path to the model on the local machine. + */ +const modelPathDownloadingIfNeeded = async ( + modelName: string, + expectedByteSize: number, +) => { + const modelPath = modelSavePath(modelName); + + if (!existsSync(modelPath)) { + log.info("CLIP image model not found, downloading"); + await downloadModel(modelPath, modelName); + } else { + const size = (await fs.stat(modelPath)).size; + if (size !== expectedByteSize) { + log.error( + `The size ${size} of model ${modelName} does not match the expected size, downloading again`, + ); + await downloadModel(modelPath, modelName); + } + } + + return modelPath; +}; + +/** Return the path where the given {@link modelName} is meant to be saved */ +const modelSavePath = (modelName: string) => + path.join(app.getPath("userData"), "models", modelName); + +const downloadModel = async (saveLocation: string, name: string) => { + // `mkdir -p` the directory where we want to save the model. + const saveDir = path.dirname(saveLocation); + await fs.mkdir(saveDir, { recursive: true }); + // Download. + log.info(`Downloading ML model from ${name}`); + const url = `https://models.ente.io/${name}`; + const res = await net.fetch(url); + if (!res.ok) throw new Error(`Failed to fetch ${url}: HTTP ${res.status}`); + const body = res.body; + if (!body) throw new Error(`Received an null response for ${url}`); + // Save. + await writeStream(saveLocation, body); + log.info(`Downloaded CLIP model ${name}`); +}; + +/** + * Create an ONNX {@link InferenceSession} with some defaults. + */ +const createInferenceSession = async (modelPath: string) => { + return await ort.InferenceSession.create(modelPath, { + // Restrict the number of threads to 1. + intraOpNumThreads: 1, + // Be more conservative with RAM usage. + enableCpuMemArena: false, + }); +}; + +const cachedCLIPImageSession = makeCachedInferenceSession( + "clip-image-vit-32-float32.onnx", + 351468764 /* 335.2 MB */, +); + +/** + * Compute CLIP embeddings for an image. + * + * The embeddings are computed using ONNX runtime, with CLIP as the model. + */ +export const computeCLIPImageEmbedding = async (input: Float32Array) => { + const session = await cachedCLIPImageSession(); + const t = Date.now(); + const feeds = { + input: new ort.Tensor("float32", input, [1, 3, 224, 224]), + }; + const results = await session.run(feeds); + log.debug(() => `ONNX/CLIP image embedding took ${Date.now() - t} ms`); + /* Need these model specific casts to type the result */ + return ensure(results.output).data as Float32Array; +}; + +const cachedCLIPTextSession = makeCachedInferenceSession( + "clip-text-vit-32-uint8.onnx", + 64173509 /* 61.2 MB */, +); + +let _tokenizer: Tokenizer | undefined; +const getTokenizer = () => { + if (!_tokenizer) _tokenizer = new Tokenizer(); + return _tokenizer; +}; + +/** + * Compute CLIP embeddings for an text snippet. + * + * The embeddings are computed using ONNX runtime, with CLIP as the model. + */ +export const computeCLIPTextEmbeddingIfAvailable = async (text: string) => { + const sessionOrSkip = await Promise.race([ + cachedCLIPTextSession(), + // Wait for a tick to get the session promise to resolved the first time + // this code runs on each app start (and the model has been downloaded). + wait(0).then(() => 1), + ]); + + // Don't wait for the download to complete. + if (typeof sessionOrSkip == "number") { + log.info( + "Ignoring CLIP text embedding request because model download is pending", + ); + return undefined; + } + + const session = sessionOrSkip; + const t = Date.now(); + const tokenizer = getTokenizer(); + const tokenizedText = Int32Array.from(tokenizer.encodeForCLIP(text)); + const feeds = { + input: new ort.Tensor("int32", tokenizedText, [1, 77]), + }; + + const results = await session.run(feeds); + log.debug(() => `ONNX/CLIP text embedding took ${Date.now() - t} ms`); + return ensure(results.output).data as Float32Array; +}; + +const cachedFaceDetectionSession = makeCachedInferenceSession( + "yolov5s_face_640_640_dynamic.onnx", + 30762872 /* 29.3 MB */, +); + +/** + * Face detection with the YOLO model and ONNX runtime. + */ +export const detectFaces = async (input: Float32Array) => { + const session = await cachedFaceDetectionSession(); + const t = Date.now(); + const feeds = { + input: new ort.Tensor("float32", input, [1, 3, 640, 640]), + }; + const results = await session.run(feeds); + log.debug(() => `ONNX/YOLO face detection took ${Date.now() - t} ms`); + return ensure(results.output).data; +}; + +const cachedFaceEmbeddingSession = makeCachedInferenceSession( + "mobilefacenet_opset15.onnx", + 5286998 /* 5 MB */, +); + +/** + * Face embedding with the MobileFaceNet model and ONNX runtime. + */ +export const computeFaceEmbeddings = async (input: Float32Array) => { + // Dimension of each face (alias) + const mobileFaceNetFaceSize = 112; + // Smaller alias + const z = mobileFaceNetFaceSize; + // Size of each face's data in the batch + const n = Math.round(input.length / (z * z * 3)); + const inputTensor = new ort.Tensor("float32", input, [n, z, z, 3]); + + const session = await cachedFaceEmbeddingSession(); + const t = Date.now(); + const feeds = { img_inputs: inputTensor }; + const results = await session.run(feeds); + log.debug(() => `ONNX/MFNT face embedding took ${Date.now() - t} ms`); + /* Need these model specific casts to extract and type the result */ + return (results.embeddings as unknown as Record) + .cpuData as Float32Array; +}; diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index dcda511df6..5d8b0cf1e3 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -1,20 +1,7 @@ /** - * @file ML related functionality, generic layer. - * - * @see also `ml-clip.ts`, `ml-face.ts`. - * - * The ML runtime we use for inference is [ONNX](https://onnxruntime.ai). Models - * for various tasks are not shipped with the app but are downloaded on demand. + * @file ML related functionality. This code runs in the main process. */ -import { app, net } from "electron/main"; -import { existsSync } from "fs"; -import fs from "node:fs/promises"; -import path from "node:path"; -import * as ort from "onnxruntime-node"; -import log from "../log"; -import { writeStream } from "../stream"; - /** * Create a new ML session. * @@ -53,15 +40,18 @@ import { writeStream } from "../stream"; * to communicate. A pair of message ports is called a message channel. The cool * thing about these is that we can pass these ports themselves over IPC. * + * > One caveat here is that the message ports can only be passed using the + * > `postMessage` APIs, not the usual send/invoke APIs. + * * So we * - * 1. Spawn a utility process. - * 2. In the utility process create a message channel. - * 3. Keep one port of the pair with us, and send the other over IPC to the - * _web worker_ that is coordinating the ML indexing on the web layer. + * 1. In the utility process create a message channel. + * 2. Spawn a utility process, and send one port of the pair to it. + * 3. Send the other port of the pair to the renderer. * - * Thereafter, the utility process and web worker can directly talk to each - * other! + * The renderer will forward that port to the web worker that is coordinating + * the ML indexing on the web layer. Thereafter, the utility process and web + * worker can directly talk to each other! * * Node.js utility process <-> Renderer web worker * @@ -70,109 +60,3 @@ export const createMLSession = () => { // }: Promise => { throw new Error("Not implemented"); }; - -/** - * Return a function that can be used to trigger a download of the specified - * model, and the creating of an ONNX inference session initialized using it. - * - * Multiple parallel calls to the returned function are fine, it ensures that - * the the model will be downloaded and the session created using it only once. - * All pending calls to it meanwhile will just await on the same promise. - * - * And once the promise is resolved, the create ONNX inference session will be - * cached, so subsequent calls to the returned function will just reuse the same - * session. - * - * {@link makeCachedInferenceSession} can itself be called anytime, it doesn't - * actively trigger a download until the returned function is called. - * - * @param modelName The name of the model to download. - * - * @param modelByteSize The size in bytes that we expect the model to have. If - * the size of the downloaded model does not match the expected size, then we - * will redownload it. - * - * @returns a function. calling that function returns a promise to an ONNX - * session. - */ -export const makeCachedInferenceSession = ( - modelName: string, - modelByteSize: number, -) => { - let session: Promise | undefined; - - const download = () => - modelPathDownloadingIfNeeded(modelName, modelByteSize); - - const createSession = (modelPath: string) => - createInferenceSession(modelPath); - - const cachedInferenceSession = () => { - if (!session) session = download().then(createSession); - return session; - }; - - return cachedInferenceSession; -}; - -/** - * Download the model named {@link modelName} if we don't already have it. - * - * Also verify that the size of the model we get matches {@expectedByteSize} (if - * not, redownload it). - * - * @returns the path to the model on the local machine. - */ -const modelPathDownloadingIfNeeded = async ( - modelName: string, - expectedByteSize: number, -) => { - const modelPath = modelSavePath(modelName); - - if (!existsSync(modelPath)) { - log.info("CLIP image model not found, downloading"); - await downloadModel(modelPath, modelName); - } else { - const size = (await fs.stat(modelPath)).size; - if (size !== expectedByteSize) { - log.error( - `The size ${size} of model ${modelName} does not match the expected size, downloading again`, - ); - await downloadModel(modelPath, modelName); - } - } - - return modelPath; -}; - -/** Return the path where the given {@link modelName} is meant to be saved */ -const modelSavePath = (modelName: string) => - path.join(app.getPath("userData"), "models", modelName); - -const downloadModel = async (saveLocation: string, name: string) => { - // `mkdir -p` the directory where we want to save the model. - const saveDir = path.dirname(saveLocation); - await fs.mkdir(saveDir, { recursive: true }); - // Download. - log.info(`Downloading ML model from ${name}`); - const url = `https://models.ente.io/${name}`; - const res = await net.fetch(url); - if (!res.ok) throw new Error(`Failed to fetch ${url}: HTTP ${res.status}`); - const body = res.body; - if (!body) throw new Error(`Received an null response for ${url}`); - // Save. - await writeStream(saveLocation, body); - log.info(`Downloaded CLIP model ${name}`); -}; - -/** - * Crete an ONNX {@link InferenceSession} with some defaults. - */ -const createInferenceSession = async (modelPath: string) => { - return await ort.InferenceSession.create(modelPath, { - // Restrict the number of threads to 1. - intraOpNumThreads: 1, - // Be more conservative with RAM usage. - enableCpuMemArena: false, - }); -}; From 1a9170632e9f0fb85ac61198aac9bc0636801745 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 11:16:04 +0530 Subject: [PATCH 03/48] Take 1 --- desktop/src/main/services/ml.ts | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 5d8b0cf1e3..a6407e55c4 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -2,6 +2,9 @@ * @file ML related functionality. This code runs in the main process. */ +import { ipcRenderer, MessageChannelMain } from "electron"; +import { utilityProcess } from "electron/main"; + /** * Create a new ML session. * @@ -57,6 +60,10 @@ * */ export const createMLSession = () => { - // }: Promise => { - throw new Error("Not implemented"); + const { port1, port2 } = new MessageChannelMain(); + + const child = utilityProcess.fork("./ml-utility"); + child.postMessage(/* unused */ "", [port1]); + + ipcRenderer.postMessage("ml-session-port", /* unused */ "", [port2]); }; From 1e720b4b7d7010a35ae248db04714ad55dc8435d Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 11:23:32 +0530 Subject: [PATCH 04/48] Scaffold --- desktop/src/main/services/ml-util-test.ts | 10 ++++++++++ desktop/src/main/services/ml.ts | 2 +- web/packages/base/types/ipc.ts | 11 +++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) create mode 100644 desktop/src/main/services/ml-util-test.ts diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts new file mode 100644 index 0000000000..684b782c4b --- /dev/null +++ b/desktop/src/main/services/ml-util-test.ts @@ -0,0 +1,10 @@ +console.log("in utility process"); + +process.parentPort.once("message", (e) => { + console.log("got message in utility process", e); + const [port] = e.ports; + + port?.on("message", (e2) => { + console.log("got message on port in utility process", e2); + }); +}); diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index a6407e55c4..d7ea06caec 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -62,7 +62,7 @@ import { utilityProcess } from "electron/main"; export const createMLSession = () => { const { port1, port2 } = new MessageChannelMain(); - const child = utilityProcess.fork("./ml-utility"); + const child = utilityProcess.fork("./ml-util-test"); child.postMessage(/* unused */ "", [port1]); ipcRenderer.postMessage("ml-session-port", /* unused */ "", [port2]); diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index 7a11553835..ebf2670da2 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -334,6 +334,17 @@ export interface Electron { // - ML + /** + * Create a new ML session. + * + * This creates a new Node.js utility process, and sets things up so that we + * can communicate directly with that utility process using a + * {@link MessagePort} we get back. + * + * For more details about the IPC flow, see: [Note: ML IPC]. + */ + createMLSession: () => Promise; + /** * Return a CLIP embedding of the given image. * From 67a9417528354788d6fd23a73285454b16ad813d Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 11:35:53 +0530 Subject: [PATCH 05/48] Scaffold --- desktop/src/main/ipc.ts | 3 +++ desktop/src/main/services/ml.ts | 2 +- desktop/src/preload.ts | 16 ++++++++++++++++ web/packages/base/types/ipc.ts | 3 +++ web/packages/new/photos/services/ml/index.ts | 5 +++++ 5 files changed, 28 insertions(+), 1 deletion(-) diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index 37ee0478e5..f891fb390b 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -42,6 +42,7 @@ import { } from "./services/fs"; import { convertToJPEG, generateImageThumbnail } from "./services/image"; import { logout } from "./services/logout"; +import { createMLSession } from "./services/ml"; import { computeCLIPImageEmbedding, computeCLIPTextEmbeddingIfAvailable, @@ -187,6 +188,8 @@ export const attachIPCHandlers = () => { // - ML + ipcMain.on("createMLSession", () => createMLSession()); + ipcMain.handle("computeCLIPImageEmbedding", (_, input: Float32Array) => computeCLIPImageEmbedding(input), ); diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index d7ea06caec..4f69fd90b1 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -65,5 +65,5 @@ export const createMLSession = () => { const child = utilityProcess.fork("./ml-util-test"); child.postMessage(/* unused */ "", [port1]); - ipcRenderer.postMessage("ml-session-port", /* unused */ "", [port2]); + ipcRenderer.postMessage("createMLSession/port", /* unused */ "", [port2]); }; diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index f2366aa63d..16b271ce9f 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -163,6 +163,21 @@ const ffmpegExec = ( // - ML +const createMLSession = () => { + ipcRenderer.send("createMLSession"); + + // The main process will do its thing, and send back the port it created to + // us by sending an message on the "createMLSession/port" channel via the + // postMessage API. This roundabout way is needed because MessagePorts + // cannot be transferred via the usual send/invoke pattern. + + return new Promise((resolve) => { + ipcRenderer.on("createMLSession/port", (event) => { + resolve(event.ports[0]); + }); + }); +}; + const computeCLIPImageEmbedding = (input: Float32Array) => ipcRenderer.invoke("computeCLIPImageEmbedding", input); @@ -339,6 +354,7 @@ contextBridge.exposeInMainWorld("electron", { // - ML + createMLSession, computeCLIPImageEmbedding, computeCLIPTextEmbeddingIfAvailable, detectFaces, diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index ebf2670da2..7188019408 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -342,6 +342,9 @@ export interface Electron { * {@link MessagePort} we get back. * * For more details about the IPC flow, see: [Note: ML IPC]. + * + * Note: For simplicity of implementation, we assume that there is at most + * one outstanding call to {@link createMLSession}. */ createMLSession: () => Promise; diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index c3432b1023..493ac6612f 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -114,6 +114,11 @@ export const canEnableML = async () => */ export const initML = () => { _isMLEnabled = isMLEnabledLocal(); + void (async () => { + console.log("yyy", 1); + const port = await ensureElectron().createMLSession(); + console.log("yyy", port); + })(); }; export const logoutML = async () => { From 4087c6ef4e6f24c3d4d3e85ffda477e8cf77d15b Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 11:39:58 +0530 Subject: [PATCH 06/48] Fix path --- desktop/src/main/services/ml.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 4f69fd90b1..a8ab08ed16 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -4,6 +4,7 @@ import { ipcRenderer, MessageChannelMain } from "electron"; import { utilityProcess } from "electron/main"; +import path from "node:path"; /** * Create a new ML session. @@ -62,7 +63,7 @@ import { utilityProcess } from "electron/main"; export const createMLSession = () => { const { port1, port2 } = new MessageChannelMain(); - const child = utilityProcess.fork("./ml-util-test"); + const child = utilityProcess.fork(path.join(__dirname, "ml-util-test.js")); child.postMessage(/* unused */ "", [port1]); ipcRenderer.postMessage("createMLSession/port", /* unused */ "", [port2]); From 7d42f23abfea677fcd6597c976d8ed94f33b1586 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 11:48:00 +0530 Subject: [PATCH 07/48] Send to the right person --- desktop/src/main.ts | 2 ++ desktop/src/main/ipc.ts | 13 +++++++++++-- desktop/src/main/services/ml.ts | 8 +++++--- 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/desktop/src/main.ts b/desktop/src/main.ts index de969e3cf7..4ebe565bca 100644 --- a/desktop/src/main.ts +++ b/desktop/src/main.ts @@ -21,6 +21,7 @@ import { attachFSWatchIPCHandlers, attachIPCHandlers, attachLogoutIPCHandler, + attachMainWindowIPCHandlers, } from "./main/ipc"; import log, { initLogging } from "./main/log"; import { createApplicationMenu, createTrayContextMenu } from "./main/menu"; @@ -121,6 +122,7 @@ const main = () => { // Setup IPC and streams. const watcher = createWatcher(mainWindow); attachIPCHandlers(); + attachMainWindowIPCHandlers(mainWindow); attachFSWatchIPCHandlers(watcher); attachLogoutIPCHandler(watcher); registerStreamProtocol(); diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index f891fb390b..9bb5d4c001 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -9,6 +9,7 @@ */ import type { FSWatcher } from "chokidar"; +import type { BrowserWindow } from "electron"; import { ipcMain } from "electron/main"; import type { CollectionMapping, @@ -188,8 +189,6 @@ export const attachIPCHandlers = () => { // - ML - ipcMain.on("createMLSession", () => createMLSession()); - ipcMain.handle("computeCLIPImageEmbedding", (_, input: Float32Array) => computeCLIPImageEmbedding(input), ); @@ -235,6 +234,16 @@ export const attachIPCHandlers = () => { ipcMain.handle("clearPendingUploads", () => clearPendingUploads()); }; +/** + * A subset of {@link attachIPCHandlers} for functions that need a reference to + * the main window to do their thing. + */ +export const attachMainWindowIPCHandlers = (mainWindow: BrowserWindow) => { + // - ML + + ipcMain.on("createMLSession", () => createMLSession(mainWindow)); +}; + /** * Sibling of {@link attachIPCHandlers} that attaches handlers specific to the * watch folder functionality. diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index a8ab08ed16..316d70c346 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -2,7 +2,7 @@ * @file ML related functionality. This code runs in the main process. */ -import { ipcRenderer, MessageChannelMain } from "electron"; +import { MessageChannelMain, type BrowserWindow } from "electron"; import { utilityProcess } from "electron/main"; import path from "node:path"; @@ -60,11 +60,13 @@ import path from "node:path"; * Node.js utility process <-> Renderer web worker * */ -export const createMLSession = () => { +export const createMLSession = (window: BrowserWindow) => { const { port1, port2 } = new MessageChannelMain(); const child = utilityProcess.fork(path.join(__dirname, "ml-util-test.js")); child.postMessage(/* unused */ "", [port1]); - ipcRenderer.postMessage("createMLSession/port", /* unused */ "", [port2]); + window.webContents.postMessage("createMLSession/port", /* unused */ "", [ + port2, + ]); }; From 180389f3e25b1f3ce8e8f24df5557497a992c142 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 12:00:44 +0530 Subject: [PATCH 08/48] Can't circumvert that way --- desktop/src/preload.ts | 15 +------------- web/packages/base/types/ipc.ts | 4 ++-- web/packages/new/photos/services/ml/index.ts | 21 +++++++++++++++++++- 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index 16b271ce9f..699e1ebc52 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -163,20 +163,7 @@ const ffmpegExec = ( // - ML -const createMLSession = () => { - ipcRenderer.send("createMLSession"); - - // The main process will do its thing, and send back the port it created to - // us by sending an message on the "createMLSession/port" channel via the - // postMessage API. This roundabout way is needed because MessagePorts - // cannot be transferred via the usual send/invoke pattern. - - return new Promise((resolve) => { - ipcRenderer.on("createMLSession/port", (event) => { - resolve(event.ports[0]); - }); - }); -}; +const createMLSession = () => ipcRenderer.send("createMLSession"); const computeCLIPImageEmbedding = (input: Float32Array) => ipcRenderer.invoke("computeCLIPImageEmbedding", input); diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index 7188019408..122735e31a 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -339,14 +339,14 @@ export interface Electron { * * This creates a new Node.js utility process, and sets things up so that we * can communicate directly with that utility process using a - * {@link MessagePort} we get back. + * {@link MessagePort} that gets posted using "createMLSession/port". * * For more details about the IPC flow, see: [Note: ML IPC]. * * Note: For simplicity of implementation, we assume that there is at most * one outstanding call to {@link createMLSession}. */ - createMLSession: () => Promise; + createMLSession: () => void; /** * Return a CLIP embedding of the given image. diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 493ac6612f..1c5d5e9437 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -9,6 +9,7 @@ import log from "@/base/log"; import { ComlinkWorker } from "@/base/worker/comlink-worker"; 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 { isInternalUser } from "../feature-flags"; @@ -116,11 +117,29 @@ export const initML = () => { _isMLEnabled = isMLEnabledLocal(); void (async () => { console.log("yyy", 1); - const port = await ensureElectron().createMLSession(); + const port = await createMLSession(); console.log("yyy", port); })(); }; +const createMLSession = async () => { + ensureElectron().createMLSession(); + + // The main process will do its thing, and send back the port it created to + // us by sending an message on the "createMLSession/port" channel via the + // postMessage API. This roundabout way is needed because MessagePorts + // cannot be transferred via the usual send/invoke pattern. + + return new Promise((resolve) => { + window.onmessage = ({ source, data, ports }: MessageEvent) => { + // The source check verifies that the message is coming from the + // preload script. The data is used as an arbitrary identifying tag. + if (source == window && data == "createMLSession/port") + resolve(ensure(ports[0])); + }; + }); +}; + export const logoutML = async () => { // `terminateMLWorker` is conceptually also part of this sequence, but for // the reasons mentioned in [Note: Caching IDB instances in separate From ea8bb4529f85cef2d09ba70f329562ceebc38233 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 12:29:11 +0530 Subject: [PATCH 09/48] We need to go via the preload --- desktop/src/preload.ts | 36 +++++++++++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index 699e1ebc52..79792aa1c0 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -36,6 +36,20 @@ * - [main] desktop/src/main/ipc.ts contains impl */ +// This code runs in the (isolated) web layer. Contrary to the impression given +// by the Electron docs (as of 2024), the window object is actually available to +// the preload script, and it is necessary for legitimate uses too. +// +// > The isolated world is connected to the DOM just the same is the main world, +// > it is just the JS contexts that are separated. +// > +// > https://github.com/electron/electron/issues/27024#issuecomment-745618327 +// +// Adding this reference here tells TypeScript that DOM typings (in particular, +// window) should be introduced in the ambient scope. +// +/// + import { contextBridge, ipcRenderer, webUtils } from "electron/renderer"; // While we can't import other code, we can import types since they're just @@ -48,6 +62,19 @@ import type { ZipItem, } from "./types/ipc"; +// - Infrastructure + +// We need to wait until the renderer is ready before sending ports via +// postMessage, and this promise comes handy in such cases. We create the +// promise at the top level so that it is guaranteed to be registered before the +// load event is fired. +// +// See: https://www.electronjs.org/docs/latest/tutorial/message-ports + +const windowLoaded = new Promise((resolve) => { + window.onload = resolve; +}); + // - General const appVersion = () => ipcRenderer.invoke("appVersion"); @@ -163,7 +190,14 @@ const ffmpegExec = ( // - ML -const createMLSession = () => ipcRenderer.send("createMLSession"); +const createMLSession = () => { + ipcRenderer.send("createMLSession"); + ipcRenderer.on("createMLSession/port", (event) => { + void windowLoaded.then(() => { + window.postMessage("createMLSession/port", "", event.ports); + }); + }); +}; const computeCLIPImageEmbedding = (input: Float32Array) => ipcRenderer.invoke("computeCLIPImageEmbedding", input); From 0195a9b494c99a97f89ad9423715153eb932338a Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 12:41:01 +0530 Subject: [PATCH 10/48] Add workaround --- desktop/src/main/stream.ts | 8 ++++++-- desktop/src/preload.ts | 8 ++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/desktop/src/main/stream.ts b/desktop/src/main/stream.ts index 749c94f491..c86232fd64 100644 --- a/desktop/src/main/stream.ts +++ b/desktop/src/main/stream.ts @@ -142,6 +142,7 @@ const handleReadZip = async (zipPath: string, entryName: string) => { // https://github.com/antelle/node-stream-zip/blob/master/node_stream_zip.js const modifiedMs = entry.time; + // @ts-expect-error [Note: Node and web stream type mismatch] return new Response(webReadableStream, { headers: { // We don't know the exact type, but it doesn't really matter, just @@ -168,9 +169,12 @@ const handleWrite = async (path: string, request: Request) => { * * @param readableStream A web * [ReadableStream](https://developer.mozilla.org/en-US/docs/Web/API/ReadableStream). + * */ -export const writeStream = (filePath: string, readableStream: ReadableStream) => - writeNodeStream(filePath, Readable.fromWeb(readableStream)); +export const writeStream = ( + filePath: string, + readableStream: unknown /*ReadableStream*/, // @ts-expect-error [Note: Node and web stream type mismatch] +) => writeNodeStream(filePath, Readable.fromWeb(readableStream)); const writeNodeStream = async (filePath: string, fileStream: Readable) => { const writeable = createWriteStream(filePath); diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index 79792aa1c0..ad36052dc0 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -48,6 +48,14 @@ // Adding this reference here tells TypeScript that DOM typings (in particular, // window) should be introduced in the ambient scope. // +// [Note: Node and web stream type mismatch] +// +// Unfortunately, adding this reference causes the ReadableStream typings to +// break since lib.dom.d.ts adds its own incompatible definitions of +// ReadableStream to the global scope. +// +// https://github.com/DefinitelyTyped/DefinitelyTyped/discussions/68407 + /// import { contextBridge, ipcRenderer, webUtils } from "electron/renderer"; From e54910f8d09ad0e04025355e6868e456c43dcaad Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 12:43:49 +0530 Subject: [PATCH 11/48] Fix origin --- desktop/src/preload.ts | 3 ++- web/packages/new/photos/services/ml/index.ts | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index ad36052dc0..283a1fca16 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -202,7 +202,8 @@ const createMLSession = () => { ipcRenderer.send("createMLSession"); ipcRenderer.on("createMLSession/port", (event) => { void windowLoaded.then(() => { - window.postMessage("createMLSession/port", "", event.ports); + // "*"" is the origin + window.postMessage("createMLSession/port", "*", event.ports); }); }); }; diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 1c5d5e9437..ec583b75eb 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -133,7 +133,7 @@ const createMLSession = async () => { return new Promise((resolve) => { window.onmessage = ({ source, data, ports }: MessageEvent) => { // The source check verifies that the message is coming from the - // preload script. The data is used as an arbitrary identifying tag. + // preload script. The data is the message that was posted. if (source == window && data == "createMLSession/port") resolve(ensure(ports[0])); }; From 24bc175f1ce35662fb91407f3b156ebd837b32a6 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 13:21:31 +0530 Subject: [PATCH 12/48] Forward --- desktop/src/main/ipc.ts | 4 +- desktop/src/main/services/ml.ts | 10 ++-- desktop/src/preload.ts | 10 ++-- web/packages/base/types/ipc.ts | 6 +-- web/packages/base/worker/comlink-worker.ts | 2 +- web/packages/new/photos/services/ml/index.ts | 53 +++++++++++-------- web/packages/new/photos/services/ml/worker.ts | 4 ++ 7 files changed, 49 insertions(+), 40 deletions(-) diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index 9bb5d4c001..caca5758b5 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -43,7 +43,7 @@ import { } from "./services/fs"; import { convertToJPEG, generateImageThumbnail } from "./services/image"; import { logout } from "./services/logout"; -import { createMLSession } from "./services/ml"; +import { createMLWorker } from "./services/ml"; import { computeCLIPImageEmbedding, computeCLIPTextEmbeddingIfAvailable, @@ -241,7 +241,7 @@ export const attachIPCHandlers = () => { export const attachMainWindowIPCHandlers = (mainWindow: BrowserWindow) => { // - ML - ipcMain.on("createMLSession", () => createMLSession(mainWindow)); + ipcMain.on("createMLWorker", () => createMLWorker(mainWindow)); }; /** diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 316d70c346..05de34fc49 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -7,7 +7,7 @@ import { utilityProcess } from "electron/main"; import path from "node:path"; /** - * Create a new ML session. + * Create a new ML worker process. * * [Note: ML IPC] * @@ -60,13 +60,11 @@ import path from "node:path"; * Node.js utility process <-> Renderer web worker * */ -export const createMLSession = (window: BrowserWindow) => { +export const createMLWorker = (window: BrowserWindow) => { const { port1, port2 } = new MessageChannelMain(); const child = utilityProcess.fork(path.join(__dirname, "ml-util-test.js")); - child.postMessage(/* unused */ "", [port1]); + child.postMessage(undefined, [port1]); - window.webContents.postMessage("createMLSession/port", /* unused */ "", [ - port2, - ]); + window.webContents.postMessage("createMLWorker/port", undefined, [port2]); }; diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index 283a1fca16..5bd2f28987 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -198,12 +198,12 @@ const ffmpegExec = ( // - ML -const createMLSession = () => { - ipcRenderer.send("createMLSession"); - ipcRenderer.on("createMLSession/port", (event) => { +const createMLWorker = () => { + ipcRenderer.send("createMLWorker"); + ipcRenderer.on("createMLWorker/port", (event) => { void windowLoaded.then(() => { // "*"" is the origin - window.postMessage("createMLSession/port", "*", event.ports); + window.postMessage("createMLWorker/port", "*", event.ports); }); }); }; @@ -384,7 +384,7 @@ contextBridge.exposeInMainWorld("electron", { // - ML - createMLSession, + createMLWorker, computeCLIPImageEmbedding, computeCLIPTextEmbeddingIfAvailable, detectFaces, diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index 122735e31a..ed7877c966 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -339,14 +339,14 @@ export interface Electron { * * This creates a new Node.js utility process, and sets things up so that we * can communicate directly with that utility process using a - * {@link MessagePort} that gets posted using "createMLSession/port". + * {@link MessagePort} that gets posted using "createMLWorker/port". * * For more details about the IPC flow, see: [Note: ML IPC]. * * Note: For simplicity of implementation, we assume that there is at most - * one outstanding call to {@link createMLSession}. + * one outstanding call to {@link createMLWorker}. */ - createMLSession: () => void; + createMLWorker: () => void; /** * Return a CLIP embedding of the given image. diff --git a/web/packages/base/worker/comlink-worker.ts b/web/packages/base/worker/comlink-worker.ts index 4562805b3b..330c5637bd 100644 --- a/web/packages/base/worker/comlink-worker.ts +++ b/web/packages/base/worker/comlink-worker.ts @@ -28,7 +28,7 @@ export class ComlinkWorker InstanceType> { /** The class (T) exposed by the web worker */ public remote: Promise>>; /** The web worker */ - private worker: Worker; + public worker: Worker; /** An arbitrary name associated with this ComlinkWorker for debugging. */ private name: string; diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index ec583b75eb..697757da59 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -6,6 +6,7 @@ import { isDesktop } from "@/base/app"; import { blobCache } from "@/base/blob-cache"; import { ensureElectron } from "@/base/electron"; import log from "@/base/log"; +import type { Electron } from "@/base/types/ipc"; import { ComlinkWorker } from "@/base/worker/comlink-worker"; import { FileType } from "@/media/file-type"; import type { EnteFile } from "@/new/photos/types/file"; @@ -67,6 +68,9 @@ const createComlinkWorker = async () => { workerDidProcessFile, }; + // Obtain a message port from the Electron layer. + const messagePort = await createMLWorker(electron); + const cw = new ComlinkWorker( "ML", new Worker(new URL("worker.ts", import.meta.url)), @@ -74,6 +78,10 @@ const createComlinkWorker = async () => { await cw.remote.then((w) => w.init(proxy(mlWorkerElectron), proxy(delegate)), ); + + // Pass the message port to our web worker. + cw.worker.postMessage("createMLWorker/port", [messagePort]); + return cw; }; @@ -93,6 +101,28 @@ export const terminateMLWorker = () => { } }; +/** + * Obtain a port from the Node.js layer that can be used to communicate with the + * ML worker process. + */ +const createMLWorker = async (electron: Electron): Promise => { + electron.createMLWorker(); + + // The main process will do its thing, and send back the port it created to + // us by sending an message on the "createMLWorker/port" channel via the + // postMessage API. This roundabout way is needed because MessagePorts + // cannot be transferred via the usual send/invoke pattern. + + return new Promise((resolve) => { + window.onmessage = ({ source, data, ports }: MessageEvent) => { + // The source check verifies that the message is coming from the + // preload script. The data is the message that was posted. + if (source == window && data == "createMLWorker/port") + resolve(ensure(ports[0])); + }; + }); +}; + /** * Return true if the current client supports ML. * @@ -115,29 +145,6 @@ export const canEnableML = async () => */ export const initML = () => { _isMLEnabled = isMLEnabledLocal(); - void (async () => { - console.log("yyy", 1); - const port = await createMLSession(); - console.log("yyy", port); - })(); -}; - -const createMLSession = async () => { - ensureElectron().createMLSession(); - - // The main process will do its thing, and send back the port it created to - // us by sending an message on the "createMLSession/port" channel via the - // postMessage API. This roundabout way is needed because MessagePorts - // cannot be transferred via the usual send/invoke pattern. - - return new Promise((resolve) => { - window.onmessage = ({ source, data, ports }: MessageEvent) => { - // The source check verifies that the message is coming from the - // preload script. The data is the message that was posted. - if (source == window && data == "createMLSession/port") - resolve(ensure(ports[0])); - }; - }); }; export const logoutML = async () => { diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 3baf512a8c..043c339e3d 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -245,6 +245,10 @@ export class MLWorker { expose(MLWorker); +globalThis.onmessage = (event: MessageEvent) => { + console.log("worker onmessage", event); +}; + /** * Find out files which need to be indexed. Then index the next batch of them. * From b28e8c2fb45e46af548ee9472ff56e1691d577ec Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 13:41:58 +0530 Subject: [PATCH 13/48] IPC --- desktop/src/main/services/ml-util-test.ts | 35 +++++++++++++++++++---- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 684b782c4b..2dd68e2ffb 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -1,10 +1,33 @@ -console.log("in utility process"); +import log from "../log"; +import { ensure, wait } from "../utils/common"; + +log.debug(() => "Started ML worker process"); process.parentPort.once("message", (e) => { - console.log("got message in utility process", e); - const [port] = e.ports; - - port?.on("message", (e2) => { - console.log("got message on port in utility process", e2); + const port = ensure(e.ports[0]); + port.on("message", (event) => { + void handleMessage(event.data).then((response) => { + if (response) port.postMessage(response); + }); }); }); + +/** Our hand-rolled IPC handler and router */ +const handleMessage = async (m: unknown) => { + if (m && typeof m == "object" && "type" in m) { + switch (m.type) { + case "foo": + if ("a" in m && typeof m.a == "string") return await foo(m.a); + break; + } + } + + log.info("Ignoring unexpected message", m); + return undefined; +}; + +const foo = async (a: string) => { + console.log("got message foo with argument", a); + await wait(0); + return a.length; +}; From 3eaa9b449ac17123eeef2f20f76c20eee0bf0308 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 14:14:45 +0530 Subject: [PATCH 14/48] IPC --- desktop/src/main/services/ml-util-test.ts | 12 +++-- web/packages/new/photos/services/ml/worker.ts | 51 +++++++++++++++++-- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 2dd68e2ffb..26c6ad9096 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -12,12 +12,18 @@ process.parentPort.once("message", (e) => { }); }); -/** Our hand-rolled IPC handler and router */ +/** + * Our hand-rolled IPC handler and router - the Node.js utility process end. + * + * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. + */ const handleMessage = async (m: unknown) => { - if (m && typeof m == "object" && "type" in m) { + if (m && typeof m == "object" && "type" in m && "id" in m) { + const id = m.id; switch (m.type) { case "foo": - if ("a" in m && typeof m.a == "string") return await foo(m.a); + if ("a" in m && typeof m.a == "string") + return { id, data: await foo(m.a) }; break; } } diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 043c339e3d..5e00f47e88 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -9,6 +9,7 @@ import { ensure } from "@/utils/ensure"; import { wait } from "@/utils/promise"; import { DOMParser } from "@xmldom/xmldom"; import { expose } from "comlink"; +import { z } from "zod"; import downloadManager from "../download"; import { cmpNewLib2, extractRawExif } from "../exif"; import { getAllLocalFiles, getLocalTrashedFiles } from "../files"; @@ -46,6 +47,46 @@ interface IndexableItem { remoteDerivedData: RemoteDerivedData | undefined; } +/** + * The port used to communicate with the Node.js ML worker process + * + * See: [Note: ML IPC] + * */ +let _port: MessagePort | undefined; + +globalThis.onmessage = (event: MessageEvent) => { + if (event.data == "createMLWorker/port") { + _port = event.ports[0]; + _port?.start(); + } +}; + +const IPCResponse = z.object({ + id: z.number(), + data: z.any(), +}); + +/** + * Our hand-rolled IPC handler and router - the web worker end. + * + * Sibling of the handleMessage function (in `ml-worker.ts`) in the desktop app. + */ +const electronMLWorker = async (type: string, data: string) => { + const port = ensure(_port); + // Generate a unique nonce to identify this RPC interaction. + const id = Math.random(); + return new Promise((resolve) => { + const handleMessage = (event: MessageEvent) => { + const response = IPCResponse.parse(event.data); + if (response.id != id) return; + port.removeEventListener("message", handleMessage); + resolve(response.data); + }; + port.addEventListener("message", handleMessage); + port.postMessage({ type, id, data }); + }); +}; + /** * Run operations related to machine learning (e.g. indexing) in a Web Worker. * @@ -113,6 +154,12 @@ export class MLWorker { // need to monkey patch it (This also ensures that it is not tree // shaken). globalThis.DOMParser = DOMParser; + + void (async () => { + console.log("yyy calling foo with 3"); + const res = await electronMLWorker("foo", "3"); + console.log("yyy calling foo with 3 result", res); + })(); } /** @@ -245,10 +292,6 @@ export class MLWorker { expose(MLWorker); -globalThis.onmessage = (event: MessageEvent) => { - console.log("worker onmessage", event); -}; - /** * Find out files which need to be indexed. Then index the next batch of them. * From d92a31d8d8df09dfaf2c93a8471e662b190ddf2e Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 14:27:27 +0530 Subject: [PATCH 15/48] Indicate error --- web/packages/new/photos/services/ml/worker.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 5e00f47e88..68a1e81c21 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -72,7 +72,13 @@ const IPCResponse = z.object({ * Sibling of the handleMessage function (in `ml-worker.ts`) in the desktop app. */ const electronMLWorker = async (type: string, data: string) => { - const port = ensure(_port); + const port = _port; + if (!port) { + throw new Error( + "No MessagePort to communicate with Electron ML worker", + ); + } + // Generate a unique nonce to identify this RPC interaction. const id = Math.random(); return new Promise((resolve) => { From 3d83786f6c917138ea602cd63a4d0d174b4155ee Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 14:45:34 +0530 Subject: [PATCH 16/48] Workaround --- desktop/src/main/services/ml-util-test.ts | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 26c6ad9096..c7bec54356 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -1,6 +1,24 @@ -import log from "../log"; +/** + * [Note: Using Electron APIs in UtilityProcess] + * + * Only a small subset of the Electron APIs are available to a UtilityProcess. + * As of writing (Jul 2024, Electron 30), only the following are available: + * - net + * - systemPreferences + * + * In particular, `app` is not available. + * + * We structure our code so that it doesn't need anything apart from `net`. + */ + +// import log from "../log"; import { ensure, wait } from "../utils/common"; +const log = { + info: (...ms: unknown[]) => console.log(...ms), + debug: (fn: () => unknown) => console.log(fn()), +}; + log.debug(() => "Started ML worker process"); process.parentPort.once("message", (e) => { From c124cdff203f39c66f65a45c7c79895ac4714387 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 14:49:31 +0530 Subject: [PATCH 17/48] Fix ordering --- web/packages/new/photos/services/ml/index.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 697757da59..c067006cf2 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -75,12 +75,13 @@ const createComlinkWorker = async () => { "ML", new Worker(new URL("worker.ts", import.meta.url)), ); - await cw.remote.then((w) => - w.init(proxy(mlWorkerElectron), proxy(delegate)), - ); - // Pass the message port to our web worker. - cw.worker.postMessage("createMLWorker/port", [messagePort]); + await cw.remote.then((w) => { + // Pass the message port to our web worker. + cw.worker.postMessage("createMLWorker/port", [messagePort]); + // Initialize it. + return w.init(proxy(mlWorkerElectron), proxy(delegate)); + }); return cw; }; From 29877d119c80fcdc30a54ba212359c1343d344e2 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 14:50:25 +0530 Subject: [PATCH 18/48] Let it flow --- desktop/src/main/services/ml-util-test.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index c7bec54356..34bed7eec1 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -28,6 +28,7 @@ process.parentPort.once("message", (e) => { if (response) port.postMessage(response); }); }); + port.start(); }); /** @@ -40,8 +41,8 @@ const handleMessage = async (m: unknown) => { const id = m.id; switch (m.type) { case "foo": - if ("a" in m && typeof m.a == "string") - return { id, data: await foo(m.a) }; + if ("data" in m && typeof m.data == "string") + return { id, data: await foo(m.data) }; break; } } From 37367f72603e60178360193757c793d3f2afb74b Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 15:41:48 +0530 Subject: [PATCH 19/48] Logging 1 --- desktop/src/main/services/ml-util-test.ts | 43 +++++++++++------- desktop/src/main/services/ml.ts | 44 ++++++++++++++++++- web/packages/new/photos/services/ml/worker.ts | 3 +- 3 files changed, 70 insertions(+), 20 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 34bed7eec1..9d312a4d04 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -1,30 +1,39 @@ -/** - * [Note: Using Electron APIs in UtilityProcess] - * - * Only a small subset of the Electron APIs are available to a UtilityProcess. - * As of writing (Jul 2024, Electron 30), only the following are available: - * - net - * - systemPreferences - * - * In particular, `app` is not available. - * - * We structure our code so that it doesn't need anything apart from `net`. - */ - -// import log from "../log"; +// import { ensure, wait } from "../utils/common"; +/** + * We cannot do + * + * import log from "../log"; + * + * because that requires the Electron APIs that are not available to a utility + * process (See: [Note: Using Electron APIs in UtilityProcess]). But even if + * that were to work, logging will still be problematic since we'd try opening + * the log file from two different Node.js processes (this one, and the main + * one), and I didn't find any indication in the electron-log repository that + * the log file's integrity would be maintained in such cases. + * + * So instead we create this proxy log object that uses `process.parentPort` to + * transport the logs over to the main process. + */ const log = { - info: (...ms: unknown[]) => console.log(...ms), + info: (...ms: unknown[]) => mainProcess("log.info", ms), debug: (fn: () => unknown) => console.log(fn()), }; +/** + * Send a message to the main process using a barebones protocol. + */ +const mainProcess = (method: string, params: unknown[]) => { + process.parentPort.postMessage({ method, params }); +}; + log.debug(() => "Started ML worker process"); process.parentPort.once("message", (e) => { const port = ensure(e.ports[0]); port.on("message", (event) => { - void handleMessage(event.data).then((response) => { + void handleMessageFromRenderer(event.data).then((response) => { if (response) port.postMessage(response); }); }); @@ -36,7 +45,7 @@ process.parentPort.once("message", (e) => { * * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. */ -const handleMessage = async (m: unknown) => { +const handleMessageFromRenderer = async (m: unknown) => { if (m && typeof m == "object" && "type" in m && "id" in m) { const id = m.id; switch (m.type) { diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 05de34fc49..cb7ac17d68 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -2,9 +2,14 @@ * @file ML related functionality. This code runs in the main process. */ -import { MessageChannelMain, type BrowserWindow } from "electron"; +import { + MessageChannelMain, + type BrowserWindow, + type UtilityProcess, +} from "electron"; import { utilityProcess } from "electron/main"; import path from "node:path"; +import log from "../log"; /** * Create a new ML worker process. @@ -58,7 +63,6 @@ import path from "node:path"; * worker can directly talk to each other! * * Node.js utility process <-> Renderer web worker - * */ export const createMLWorker = (window: BrowserWindow) => { const { port1, port2 } = new MessageChannelMain(); @@ -67,4 +71,40 @@ export const createMLWorker = (window: BrowserWindow) => { child.postMessage(undefined, [port1]); window.webContents.postMessage("createMLWorker/port", undefined, [port2]); + + handleMLWorkerRequests(child); +}; + +/** + * Handle requests from the utility process. + * + * [Note: Using Electron APIs in UtilityProcess] + * + * Only a small subset of the Electron APIs are available to a UtilityProcess. + * As of writing (Jul 2024, Electron 30), only the following are available: + * + * - net + * - systemPreferences + * + * In particular, `app` is not available. + * + * We structure our code so that it doesn't need anything apart from `net`. + * + * For the other cases, + * + * - When we need to communicate from the utility process to the main process, + * we use the `parentPort` in the utility process. + */ +const handleMLWorkerRequests = (child: UtilityProcess) => { + child.on("message", (m: unknown) => { + if (m && typeof m == "object" && "method" in m) { + switch (m.method) { + default: + break; + } + } + + log.info("Ignoring unexpected message", m); + return undefined; + }); }; diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 68a1e81c21..540b15eafd 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -69,7 +69,8 @@ const IPCResponse = z.object({ /** * Our hand-rolled IPC handler and router - the web worker end. * - * Sibling of the handleMessage function (in `ml-worker.ts`) in the desktop app. + * Sibling of the handleMessageFromRenderer function (in `ml-worker.ts`) in the + * desktop code. */ const electronMLWorker = async (type: string, data: string) => { const port = _port; From e66e9251dbebcd2dedb995b8c749cf740cf500e4 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 15:48:44 +0530 Subject: [PATCH 20/48] Fancier --- desktop/src/main/services/ml.ts | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index cb7ac17d68..98eee821c6 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -97,14 +97,19 @@ export const createMLWorker = (window: BrowserWindow) => { */ const handleMLWorkerRequests = (child: UtilityProcess) => { child.on("message", (m: unknown) => { - if (m && typeof m == "object" && "method" in m) { + if (m && typeof m == "object" && "method" in m && "params" in m) { switch (m.method) { + case "log.info": + if (Array.isArray(m.params)) { + const params = m.params as unknown[]; + log.info("[ml-worker]", ...params); + return; + } + break; default: break; } } - - log.info("Ignoring unexpected message", m); - return undefined; + log.warn("Ignoring unexpected message from ML worker", m); }); }; From 81b52419a5b1ed128cbaa651530792df401a5456 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 16:01:17 +0530 Subject: [PATCH 21/48] debug strings --- desktop/src/main/services/ml-util-test.ts | 16 +++++++++++----- desktop/src/main/services/ml.ts | 19 ++++++++++++++----- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 9d312a4d04..18293c2a20 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -18,17 +18,23 @@ import { ensure, wait } from "../utils/common"; */ const log = { info: (...ms: unknown[]) => mainProcess("log.info", ms), - debug: (fn: () => unknown) => console.log(fn()), + /** + * Unlike the real {@link log.debug}, this is (a) eagerly evaluated, and (b) + * accepts only strings. + */ + debugString: (s: string) => mainProcess("log.debugString", s), }; /** * Send a message to the main process using a barebones protocol. */ -const mainProcess = (method: string, params: unknown[]) => { - process.parentPort.postMessage({ method, params }); +const mainProcess = (method: string, param: unknown) => { + process.parentPort.postMessage({ method, param }); }; -log.debug(() => "Started ML worker process"); +log.debugString( + `Started ML worker process with args ${process.argv.join(" ")}`, +); process.parentPort.once("message", (e) => { const port = ensure(e.ports[0]); @@ -61,7 +67,7 @@ const handleMessageFromRenderer = async (m: unknown) => { }; const foo = async (a: string) => { - console.log("got message foo with argument", a); + log.info("got message foo with argument", a); await wait(0); return a.length; }; diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 98eee821c6..5de35b4b28 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -96,20 +96,29 @@ export const createMLWorker = (window: BrowserWindow) => { * we use the `parentPort` in the utility process. */ const handleMLWorkerRequests = (child: UtilityProcess) => { + const logTag = "[ml-worker]"; child.on("message", (m: unknown) => { - if (m && typeof m == "object" && "method" in m && "params" in m) { + if (m && typeof m == "object" && "method" in m && "param" in m) { + const p = m.param; switch (m.method) { case "log.info": - if (Array.isArray(m.params)) { - const params = m.params as unknown[]; - log.info("[ml-worker]", ...params); + if (Array.isArray(p)) { + // Need to cast from any[] to unknown[] + log.info(logTag, ...(p as unknown[])); return; } break; + case "log.debugString": + if (typeof p == "string") { + log.debug(() => `${logTag} ${p}`); + return; + } + break; + default: break; } } - log.warn("Ignoring unexpected message from ML worker", m); + log.info("Ignoring unexpected message from ML worker", m); }); }; From 4ca40085c1e1b4da288884b26ce078cf73b0b9b3 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 19:13:54 +0530 Subject: [PATCH 22/48] init --- desktop/src/main/services/ml-util-test.ts | 36 +++++++++++++++++++++-- desktop/src/main/services/ml.ts | 15 ++++++++-- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 18293c2a20..9ed2f3a46d 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -1,4 +1,6 @@ -// +// See [Note: Using Electron APIs in UtilityProcess] about what we can and +// cannot import. + import { ensure, wait } from "../utils/common"; /** @@ -17,6 +19,11 @@ import { ensure, wait } from "../utils/common"; * transport the logs over to the main process. */ const log = { + /** + * Unlike the real {@link log.error}, this accepts only the first string + * argument, not the second optional error one. + */ + errorString: (s: string) => mainProcess("log.errorString", s), info: (...ms: unknown[]) => mainProcess("log.info", ms), /** * Unlike the real {@link log.debug}, this is (a) eagerly evaluated, and (b) @@ -37,6 +44,8 @@ log.debugString( ); process.parentPort.once("message", (e) => { + parseInitData(e.data); + const port = ensure(e.ports[0]); port.on("message", (event) => { void handleMessageFromRenderer(event.data).then((response) => { @@ -46,6 +55,29 @@ process.parentPort.once("message", (e) => { port.start(); }); +/** + * We cannot access Electron's {@link app} object within a utility process, so + * we pass the value of `app.getPath("userData")` during initialization, and it + * can be subsequently retrieved from here. + */ +let _userDataPath: string | undefined; + +/** Equivalent to app.getPath("userData") */ +const userDataPath = () => ensure(_userDataPath); + +const parseInitData = (data: unknown) => { + if ( + data && + typeof data == "object" && + "userDataPath" in data && + typeof data.userDataPath == "string" + ) { + _userDataPath = data.userDataPath; + } else { + log.errorString("Unparseable initialization data"); + } +}; + /** * Our hand-rolled IPC handler and router - the Node.js utility process end. * @@ -67,7 +99,7 @@ const handleMessageFromRenderer = async (m: unknown) => { }; const foo = async (a: string) => { - log.info("got message foo with argument", a); + log.info("got message foo with argument", a, userDataPath()); await wait(0); return a.length; }; diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 5de35b4b28..6e975d3910 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -7,7 +7,7 @@ import { type BrowserWindow, type UtilityProcess, } from "electron"; -import { utilityProcess } from "electron/main"; +import { app, utilityProcess } from "electron/main"; import path from "node:path"; import log from "../log"; @@ -68,7 +68,8 @@ export const createMLWorker = (window: BrowserWindow) => { const { port1, port2 } = new MessageChannelMain(); const child = utilityProcess.fork(path.join(__dirname, "ml-util-test.js")); - child.postMessage(undefined, [port1]); + const userDataPath = app.getPath("userData"); + child.postMessage({ userDataPath }, [port1]); window.webContents.postMessage("createMLWorker/port", undefined, [port2]); @@ -92,6 +93,9 @@ export const createMLWorker = (window: BrowserWindow) => { * * For the other cases, * + * - Additional parameters to the utility process are passed alongwith the + * initial message where we provide it the message port. + * * - When we need to communicate from the utility process to the main process, * we use the `parentPort` in the utility process. */ @@ -101,6 +105,12 @@ const handleMLWorkerRequests = (child: UtilityProcess) => { if (m && typeof m == "object" && "method" in m && "param" in m) { const p = m.param; switch (m.method) { + case "log.errorString": + if (typeof p == "string") { + log.error(`${logTag} ${p}`); + return; + } + break; case "log.info": if (Array.isArray(p)) { // Need to cast from any[] to unknown[] @@ -114,7 +124,6 @@ const handleMLWorkerRequests = (child: UtilityProcess) => { return; } break; - default: break; } From 18cb596d575e3e596d9913a583fa97429dc7da64 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 19:41:51 +0530 Subject: [PATCH 23/48] Error 1 --- desktop/src/main/services/ml-util-test.ts | 52 ++++++++++++++++------- desktop/src/main/services/ml.ts | 10 ++--- 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 9ed2f3a46d..69c0f33598 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -35,9 +35,8 @@ const log = { /** * Send a message to the main process using a barebones protocol. */ -const mainProcess = (method: string, param: unknown) => { - process.parentPort.postMessage({ method, param }); -}; +const mainProcess = (method: string, param: unknown) => + process.parentPort.postMessage({ method, p: param }); log.debugString( `Started ML worker process with args ${process.argv.join(" ")}`, @@ -47,10 +46,10 @@ process.parentPort.once("message", (e) => { parseInitData(e.data); const port = ensure(e.ports[0]); - port.on("message", (event) => { - void handleMessageFromRenderer(event.data).then((response) => { - if (response) port.postMessage(response); - }); + port.on("message", (request) => { + void handleMessageFromRenderer(request.data).then((response) => + port.postMessage(response), + ); }); port.start(); }); @@ -69,6 +68,7 @@ const parseInitData = (data: unknown) => { if ( data && typeof data == "object" && + "userDataPateh" in data && "userDataPath" in data && typeof data.userDataPath == "string" ) { @@ -79,23 +79,43 @@ const parseInitData = (data: unknown) => { }; /** - * Our hand-rolled IPC handler and router - the Node.js utility process end. + * Our hand-rolled RPC handler and router - the Node.js utility process end. * * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. + * + * [Note: Node.js ML worker RPC protocol] + * + * - Each RPC call (i.e. request message) has a "method" (string), "id" + * (number) and "p" (arbitrary param). + * + * - Each RPC result (i.e. response message) has an "id" (number) that is the + * same as the "id" for the request which it corresponds to. + * + * - If the RPC call was a success, then the response messege will have an + * "result" (arbitrary result) property. Otherwise it will have a "error" + * (string) property describing what went wrong. */ const handleMessageFromRenderer = async (m: unknown) => { - if (m && typeof m == "object" && "type" in m && "id" in m) { + if (m && typeof m == "object" && "method" in m && "id" in m && "p" in m) { const id = m.id; - switch (m.type) { - case "foo": - if ("data" in m && typeof m.data == "string") - return { id, data: await foo(m.data) }; - break; + const p = m.p; + try { + switch (m.method) { + case "foo": + if (p && typeof p == "string") + return { id, result: await foo(p) }; + break; + } + } catch (e) { + return { id, error: e instanceof Error ? e.message : String(e) }; } + return { id, error: "Unknown message" }; } - log.info("Ignoring unexpected message", m); - return undefined; + // We don't even have an "id", so at least log it lest the renderer also + // ignore the "id"-less response. + log.info("Ignoring unknown message", m); + return { error: "Unknown message" }; }; const foo = async (a: string) => { diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 6e975d3910..3f35d8573d 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -73,11 +73,11 @@ export const createMLWorker = (window: BrowserWindow) => { window.webContents.postMessage("createMLWorker/port", undefined, [port2]); - handleMLWorkerRequests(child); + handleMessagesFromUtilityProcess(child); }; /** - * Handle requests from the utility process. + * Handle messages posted from the utility process. * * [Note: Using Electron APIs in UtilityProcess] * @@ -99,11 +99,11 @@ export const createMLWorker = (window: BrowserWindow) => { * - When we need to communicate from the utility process to the main process, * we use the `parentPort` in the utility process. */ -const handleMLWorkerRequests = (child: UtilityProcess) => { +const handleMessagesFromUtilityProcess = (child: UtilityProcess) => { const logTag = "[ml-worker]"; child.on("message", (m: unknown) => { - if (m && typeof m == "object" && "method" in m && "param" in m) { - const p = m.param; + if (m && typeof m == "object" && "method" in m && "p" in m) { + const p = m.p; switch (m.method) { case "log.errorString": if (typeof p == "string") { From 3f3d10f57b6d8538fedc8ceb3a9d7daf448a0be3 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 19:53:39 +0530 Subject: [PATCH 24/48] Error 2 --- desktop/src/main/services/ml-util-test.ts | 2 +- desktop/src/main/services/ml.ts | 2 +- web/packages/new/photos/services/ml/worker.ts | 19 ++++++++++--------- 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts index 69c0f33598..54f4cd8e0f 100644 --- a/desktop/src/main/services/ml-util-test.ts +++ b/desktop/src/main/services/ml-util-test.ts @@ -33,7 +33,7 @@ const log = { }; /** - * Send a message to the main process using a barebones protocol. + * Send a message to the main process using a barebones RPC protocol. */ const mainProcess = (method: string, param: unknown) => process.parentPort.postMessage({ method, p: param }); diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 3f35d8573d..eef362ccc4 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -128,6 +128,6 @@ const handleMessagesFromUtilityProcess = (child: UtilityProcess) => { break; } } - log.info("Ignoring unexpected message from ML worker", m); + log.info("Ignoring unknown message from ML worker", m); }); }; diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 540b15eafd..2ce6315715 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -63,16 +63,15 @@ globalThis.onmessage = (event: MessageEvent) => { const IPCResponse = z.object({ id: z.number(), - data: z.any(), + result: z.any().optional(), + error: z.string().optional(), }); /** - * Our hand-rolled IPC handler and router - the web worker end. - * - * Sibling of the handleMessageFromRenderer function (in `ml-worker.ts`) in the - * desktop code. + * Make a call to the ML worker running in the Node.js layer using our + * hand-rolled RPC protocol. See: [Note: Node.js ML worker RPC protocol]. */ -const electronMLWorker = async (type: string, data: string) => { +const electronMLWorker = async (method: string, p: string) => { const port = _port; if (!port) { throw new Error( @@ -82,15 +81,17 @@ const electronMLWorker = async (type: string, data: string) => { // Generate a unique nonce to identify this RPC interaction. const id = Math.random(); - return new Promise((resolve) => { + return new Promise((resolve, reject) => { const handleMessage = (event: MessageEvent) => { const response = IPCResponse.parse(event.data); if (response.id != id) return; port.removeEventListener("message", handleMessage); - resolve(response.data); + const error = response.error; + if (error) reject(new Error(error)); + else resolve(response.result); }; port.addEventListener("message", handleMessage); - port.postMessage({ type, id, data }); + port.postMessage({ id, method, p }); }); }; From 65cfcc27a8db626c31738b13d1e66e4f62c004d8 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 20:00:21 +0530 Subject: [PATCH 25/48] Rearrange --- web/packages/new/photos/services/ml/worker.ts | 68 +++++++++---------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 2ce6315715..10e8f2fd7b 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -61,40 +61,6 @@ globalThis.onmessage = (event: MessageEvent) => { } }; -const IPCResponse = z.object({ - id: z.number(), - result: z.any().optional(), - error: z.string().optional(), -}); - -/** - * Make a call to the ML worker running in the Node.js layer using our - * hand-rolled RPC protocol. See: [Note: Node.js ML worker RPC protocol]. - */ -const electronMLWorker = async (method: string, p: string) => { - const port = _port; - if (!port) { - throw new Error( - "No MessagePort to communicate with Electron ML worker", - ); - } - - // Generate a unique nonce to identify this RPC interaction. - const id = Math.random(); - return new Promise((resolve, reject) => { - const handleMessage = (event: MessageEvent) => { - const response = IPCResponse.parse(event.data); - if (response.id != id) return; - port.removeEventListener("message", handleMessage); - const error = response.error; - if (error) reject(new Error(error)); - else resolve(response.result); - }; - port.addEventListener("message", handleMessage); - port.postMessage({ id, method, p }); - }); -}; - /** * Run operations related to machine learning (e.g. indexing) in a Web Worker. * @@ -300,6 +266,40 @@ export class MLWorker { expose(MLWorker); +/** + * Make a call to the ML worker running in the Node.js layer using our + * hand-rolled RPC protocol. See: [Note: Node.js ML worker RPC protocol]. + */ +const electronMLWorker = async (method: string, p: string) => { + const port = _port; + if (!port) { + throw new Error( + "No MessagePort to communicate with Electron ML worker", + ); + } + + // Generate a unique nonce to identify this RPC interaction. + const id = Math.random(); + return new Promise((resolve, reject) => { + const handleMessage = (event: MessageEvent) => { + const response = RPCResponse.parse(event.data); + if (response.id != id) return; + port.removeEventListener("message", handleMessage); + const error = response.error; + if (error) reject(new Error(error)); + else resolve(response.result); + }; + port.addEventListener("message", handleMessage); + port.postMessage({ id, method, p }); + }); +}; + +const RPCResponse = z.object({ + id: z.number(), + result: z.any().optional(), + error: z.string().optional(), +}); + /** * Find out files which need to be indexed. Then index the next batch of them. * From 7baacc6a778df9378e1fffca93e5bc90883339b6 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Tue, 30 Jul 2024 20:22:09 +0530 Subject: [PATCH 26/48] For real - 1 --- desktop/src/main/services/ml-util-test.ts | 125 ---------------- .../services/{ml-utility.ts => ml-worker.ts} | 135 ++++++++++++++++-- desktop/src/main/stream.ts | 38 +---- desktop/src/main/utils/stream.ts | 39 +++++ 4 files changed, 166 insertions(+), 171 deletions(-) delete mode 100644 desktop/src/main/services/ml-util-test.ts rename desktop/src/main/services/{ml-utility.ts => ml-worker.ts} (63%) create mode 100644 desktop/src/main/utils/stream.ts diff --git a/desktop/src/main/services/ml-util-test.ts b/desktop/src/main/services/ml-util-test.ts deleted file mode 100644 index 54f4cd8e0f..0000000000 --- a/desktop/src/main/services/ml-util-test.ts +++ /dev/null @@ -1,125 +0,0 @@ -// See [Note: Using Electron APIs in UtilityProcess] about what we can and -// cannot import. - -import { ensure, wait } from "../utils/common"; - -/** - * We cannot do - * - * import log from "../log"; - * - * because that requires the Electron APIs that are not available to a utility - * process (See: [Note: Using Electron APIs in UtilityProcess]). But even if - * that were to work, logging will still be problematic since we'd try opening - * the log file from two different Node.js processes (this one, and the main - * one), and I didn't find any indication in the electron-log repository that - * the log file's integrity would be maintained in such cases. - * - * So instead we create this proxy log object that uses `process.parentPort` to - * transport the logs over to the main process. - */ -const log = { - /** - * Unlike the real {@link log.error}, this accepts only the first string - * argument, not the second optional error one. - */ - errorString: (s: string) => mainProcess("log.errorString", s), - info: (...ms: unknown[]) => mainProcess("log.info", ms), - /** - * Unlike the real {@link log.debug}, this is (a) eagerly evaluated, and (b) - * accepts only strings. - */ - debugString: (s: string) => mainProcess("log.debugString", s), -}; - -/** - * Send a message to the main process using a barebones RPC protocol. - */ -const mainProcess = (method: string, param: unknown) => - process.parentPort.postMessage({ method, p: param }); - -log.debugString( - `Started ML worker process with args ${process.argv.join(" ")}`, -); - -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(); -}); - -/** - * We cannot access Electron's {@link app} object within a utility process, so - * we pass the value of `app.getPath("userData")` during initialization, and it - * can be subsequently retrieved from here. - */ -let _userDataPath: string | undefined; - -/** Equivalent to app.getPath("userData") */ -const userDataPath = () => ensure(_userDataPath); - -const parseInitData = (data: unknown) => { - if ( - data && - typeof data == "object" && - "userDataPateh" in data && - "userDataPath" in data && - typeof data.userDataPath == "string" - ) { - _userDataPath = data.userDataPath; - } else { - log.errorString("Unparseable initialization data"); - } -}; - -/** - * Our hand-rolled RPC handler and router - the Node.js utility process end. - * - * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. - * - * [Note: Node.js ML worker RPC protocol] - * - * - Each RPC call (i.e. request message) has a "method" (string), "id" - * (number) and "p" (arbitrary param). - * - * - Each RPC result (i.e. response message) has an "id" (number) that is the - * same as the "id" for the request which it corresponds to. - * - * - If the RPC call was a success, then the response messege will have an - * "result" (arbitrary result) property. Otherwise it will have a "error" - * (string) property describing what went wrong. - */ -const handleMessageFromRenderer = async (m: unknown) => { - if (m && typeof m == "object" && "method" in m && "id" in m && "p" in m) { - const id = m.id; - const p = m.p; - try { - switch (m.method) { - case "foo": - if (p && typeof p == "string") - return { id, result: await foo(p) }; - break; - } - } catch (e) { - return { id, error: e instanceof Error ? e.message : String(e) }; - } - return { id, error: "Unknown message" }; - } - - // We don't even have an "id", so at least log it lest the renderer also - // ignore the "id"-less response. - log.info("Ignoring unknown message", m); - return { error: "Unknown message" }; -}; - -const foo = async (a: string) => { - log.info("got message foo with argument", a, userDataPath()); - await wait(0); - return a.length; -}; diff --git a/desktop/src/main/services/ml-utility.ts b/desktop/src/main/services/ml-worker.ts similarity index 63% rename from desktop/src/main/services/ml-utility.ts rename to desktop/src/main/services/ml-worker.ts index 79d39edea4..255dbd3f5b 100644 --- a/desktop/src/main/services/ml-utility.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -5,15 +5,132 @@ * for various tasks are not shipped with the app but are downloaded on demand. */ +// See [Note: Using Electron APIs in UtilityProcess] about what we can and +// cannot import. + import Tokenizer from "clip-bpe-js"; -import { app, net } from "electron/main"; +import { net } from "electron/main"; import { existsSync } from "fs"; import fs from "node:fs/promises"; import path from "node:path"; import * as ort from "onnxruntime-node"; -import log from "../log"; -import { writeStream } from "../stream"; import { ensure, wait } from "../utils/common"; +import { writeStream } from "../utils/stream"; + +/** + * We cannot do + * + * import log from "../log"; + * + * because that requires the Electron APIs that are not available to a utility + * process (See: [Note: Using Electron APIs in UtilityProcess]). But even if + * that were to work, logging will still be problematic since we'd try opening + * the log file from two different Node.js processes (this one, and the main + * one), and I didn't find any indication in the electron-log repository that + * the log file's integrity would be maintained in such cases. + * + * So instead we create this proxy log object that uses `process.parentPort` to + * transport the logs over to the main process. + */ +const log = { + /** + * Unlike the real {@link log.error}, this accepts only the first string + * argument, not the second optional error one. + */ + errorString: (s: string) => mainProcess("log.errorString", s), + info: (...ms: unknown[]) => mainProcess("log.info", ms), + /** + * Unlike the real {@link log.debug}, this is (a) eagerly evaluated, and (b) + * accepts only strings. + */ + debugString: (s: string) => mainProcess("log.debugString", s), +}; + +/** + * Send a message to the main process using a barebones RPC protocol. + */ +const mainProcess = (method: string, param: unknown) => + process.parentPort.postMessage({ method, p: param }); + +log.debugString( + `Started ML worker process with args ${process.argv.join(" ")}`, +); + +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(); +}); + +/** + * We cannot access Electron's {@link app} object within a utility process, so + * we pass the value of `app.getPath("userData")` during initialization, and it + * can be subsequently retrieved from here. + */ +let _userDataPath: string | undefined; + +/** Equivalent to app.getPath("userData") */ +const userDataPath = () => ensure(_userDataPath); + +const parseInitData = (data: unknown) => { + if ( + data && + typeof data == "object" && + "userDataPateh" in data && + "userDataPath" in data && + typeof data.userDataPath == "string" + ) { + _userDataPath = data.userDataPath; + } else { + log.errorString("Unparseable initialization data"); + } +}; + +/** + * Our hand-rolled RPC handler and router - the Node.js utility process end. + * + * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. + * + * [Note: Node.js ML worker RPC protocol] + * + * - Each RPC call (i.e. request message) has a "method" (string), "id" + * (number) and "p" (arbitrary param). + * + * - Each RPC result (i.e. response message) has an "id" (number) that is the + * same as the "id" for the request which it corresponds to. + * + * - If the RPC call was a success, then the response messege will have an + * "result" (arbitrary result) property. Otherwise it will have a "error" + * (string) property describing what went wrong. + */ +const handleMessageFromRenderer = async (m: unknown) => { + if (m && typeof m == "object" && "method" in m && "id" in m && "p" in m) { + const id = m.id; + const p = m.p; + try { + switch (m.method) { + case "foo": + if (p && typeof p == "string") + return { id, result: await foo(p) }; + break; + } + } catch (e) { + return { id, error: e instanceof Error ? e.message : String(e) }; + } + return { id, error: "Unknown message" }; + } + + // We don't even have an "id", so at least log it lest the renderer also + // ignore the "id"-less response. + log.info("Ignoring unknown message", m); + return { error: "Unknown message" }; +}; /** * Return a function that can be used to trigger a download of the specified @@ -79,7 +196,7 @@ const modelPathDownloadingIfNeeded = async ( } else { const size = (await fs.stat(modelPath)).size; if (size !== expectedByteSize) { - log.error( + log.errorString( `The size ${size} of model ${modelName} does not match the expected size, downloading again`, ); await downloadModel(modelPath, modelName); @@ -91,7 +208,7 @@ const modelPathDownloadingIfNeeded = async ( /** Return the path where the given {@link modelName} is meant to be saved */ const modelSavePath = (modelName: string) => - path.join(app.getPath("userData"), "models", modelName); + path.join(userDataPath(), "models", modelName); const downloadModel = async (saveLocation: string, name: string) => { // `mkdir -p` the directory where we want to save the model. @@ -138,7 +255,7 @@ export const computeCLIPImageEmbedding = async (input: Float32Array) => { input: new ort.Tensor("float32", input, [1, 3, 224, 224]), }; const results = await session.run(feeds); - log.debug(() => `ONNX/CLIP image embedding took ${Date.now() - t} ms`); + log.debugString(`ONNX/CLIP image embedding took ${Date.now() - t} ms`); /* Need these model specific casts to type the result */ return ensure(results.output).data as Float32Array; }; @@ -184,7 +301,7 @@ export const computeCLIPTextEmbeddingIfAvailable = async (text: string) => { }; const results = await session.run(feeds); - log.debug(() => `ONNX/CLIP text embedding took ${Date.now() - t} ms`); + log.debugString(`ONNX/CLIP text embedding took ${Date.now() - t} ms`); return ensure(results.output).data as Float32Array; }; @@ -203,7 +320,7 @@ export const detectFaces = async (input: Float32Array) => { input: new ort.Tensor("float32", input, [1, 3, 640, 640]), }; const results = await session.run(feeds); - log.debug(() => `ONNX/YOLO face detection took ${Date.now() - t} ms`); + log.debugString(`ONNX/YOLO face detection took ${Date.now() - t} ms`); return ensure(results.output).data; }; @@ -228,7 +345,7 @@ export const computeFaceEmbeddings = async (input: Float32Array) => { const t = Date.now(); const feeds = { img_inputs: inputTensor }; const results = await session.run(feeds); - log.debug(() => `ONNX/MFNT face embedding took ${Date.now() - t} ms`); + log.debugString(`ONNX/MFNT face embedding took ${Date.now() - t} ms`); /* Need these model specific casts to extract and type the result */ return (results.embeddings as unknown as Record) .cpuData as Float32Array; diff --git a/desktop/src/main/stream.ts b/desktop/src/main/stream.ts index c86232fd64..d32eecc627 100644 --- a/desktop/src/main/stream.ts +++ b/desktop/src/main/stream.ts @@ -3,7 +3,6 @@ */ import { net, protocol } from "electron/main"; import { randomUUID } from "node:crypto"; -import { createWriteStream, existsSync } from "node:fs"; import fs from "node:fs/promises"; import { Readable } from "node:stream"; import { ReadableStream } from "node:stream/web"; @@ -12,6 +11,7 @@ import log from "./log"; import { ffmpegConvertToMP4 } from "./services/ffmpeg"; import { markClosableZip, openZip } from "./services/zip"; import { ensure } from "./utils/common"; +import { writeStream } from "./utils/stream"; import { deleteTempFile, deleteTempFileIgnoringErrors, @@ -160,42 +160,6 @@ const handleWrite = async (path: string, request: Request) => { return new Response("", { status: 200 }); }; -/** - * Write a (web) ReadableStream to a file at the given {@link filePath}. - * - * The returned promise resolves when the write completes. - * - * @param filePath The local file system path where the file should be written. - * - * @param readableStream A web - * [ReadableStream](https://developer.mozilla.org/en-US/docs/Web/API/ReadableStream). - * - */ -export const writeStream = ( - filePath: string, - readableStream: unknown /*ReadableStream*/, // @ts-expect-error [Note: Node and web stream type mismatch] -) => writeNodeStream(filePath, Readable.fromWeb(readableStream)); - -const writeNodeStream = async (filePath: string, fileStream: Readable) => { - const writeable = createWriteStream(filePath); - - fileStream.on("error", (err) => { - writeable.destroy(err); // Close the writable stream with an error - }); - - fileStream.pipe(writeable); - - await new Promise((resolve, reject) => { - writeable.on("finish", resolve); - writeable.on("error", (err) => { - if (existsSync(filePath)) { - void fs.unlink(filePath); - } - reject(err); - }); - }); -}; - /** * A map from token to file paths for convert-to-mp4 requests that we have * received. diff --git a/desktop/src/main/utils/stream.ts b/desktop/src/main/utils/stream.ts new file mode 100644 index 0000000000..f5a98de0f7 --- /dev/null +++ b/desktop/src/main/utils/stream.ts @@ -0,0 +1,39 @@ +import { createWriteStream, existsSync } from "node:fs"; +import fs from "node:fs/promises"; +import { Readable } from "node:stream"; + +/** + * Write a (web) ReadableStream to a file at the given {@link filePath}. + * + * The returned promise resolves when the write completes. + * + * @param filePath The local file system path where the file should be written. + * + * @param readableStream A web + * [ReadableStream](https://developer.mozilla.org/en-US/docs/Web/API/ReadableStream). + * + */ +export const writeStream = ( + filePath: string, + readableStream: unknown /*ReadableStream*/, // @ts-expect-error [Note: Node and web stream type mismatch] +) => writeNodeStream(filePath, Readable.fromWeb(readableStream)); + +const writeNodeStream = async (filePath: string, fileStream: Readable) => { + const writeable = createWriteStream(filePath); + + fileStream.on("error", (err) => { + writeable.destroy(err); // Close the writable stream with an error + }); + + fileStream.pipe(writeable); + + await new Promise((resolve, reject) => { + writeable.on("finish", resolve); + writeable.on("error", (err) => { + if (existsSync(filePath)) { + void fs.unlink(filePath); + } + reject(err); + }); + }); +}; From 95facd60e02c5a470b67b8d8a9b719179af41e23 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 09:20:29 +0530 Subject: [PATCH 27/48] integrate wip 1 --- web/packages/base/types/ipc.ts | 53 ------- .../new/photos/services/ml/worker-rpc.ts | 131 ++++++++++++++++++ web/packages/new/photos/services/ml/worker.ts | 55 +------- 3 files changed, 134 insertions(+), 105 deletions(-) create mode 100644 web/packages/new/photos/services/ml/worker-rpc.ts diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index ed7877c966..4aca8b865d 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -348,59 +348,6 @@ export interface Electron { */ createMLWorker: () => void; - /** - * Return a CLIP embedding of the given image. - * - * See: [Note: Natural language search using CLIP] - * - * The input is a opaque float32 array representing the image. The layout - * and exact encoding of the input is specific to our implementation and the - * ML model (CLIP) we use. - * - * @returns A CLIP embedding (an array of 512 floating point values). - */ - computeCLIPImageEmbedding: (input: Float32Array) => Promise; - - /** - * Return a CLIP embedding of the given image if we already have the model - * downloaded and prepped. If the model is not available return `undefined`. - * - * This differs from the other sibling ML functions in that it doesn't wait - * for the model download to finish. It does trigger a model download, but - * then immediately returns `undefined`. At some future point, when the - * model downloaded finishes, calls to this function will start returning - * the result we seek. - * - * The reason for doing it in this asymmetric way is because CLIP text - * embeddings are used as part of deducing user initiated search results, - * and we don't want to block that interaction on a large network request. - * - * See: [Note: Natural language search using CLIP] - * - * @param text The string whose embedding we want to compute. - * - * @returns A CLIP embedding. - */ - computeCLIPTextEmbeddingIfAvailable: ( - text: string, - ) => Promise; - - /** - * Detect faces in the given image using YOLO. - * - * Both the input and output are opaque binary data whose internal structure - * is specific to our implementation and the model (YOLO) we use. - */ - detectFaces: (input: Float32Array) => Promise; - - /** - * Return a MobileFaceNet embeddings for the given faces. - * - * Both the input and output are opaque binary data whose internal structure - * is specific to our implementation and the model (MobileFaceNet) we use. - */ - computeFaceEmbeddings: (input: Float32Array) => Promise; - // - Watch /** diff --git a/web/packages/new/photos/services/ml/worker-rpc.ts b/web/packages/new/photos/services/ml/worker-rpc.ts new file mode 100644 index 0000000000..782c0a3bd3 --- /dev/null +++ b/web/packages/new/photos/services/ml/worker-rpc.ts @@ -0,0 +1,131 @@ +import { z } from "zod"; + +/** + * The port used to communicate with the Node.js ML worker process + * + * See: [Note: ML IPC] + * */ +let _port: MessagePort | undefined; + +/** + * Use the given {@link MessagePort} to communicate with the Node.js ML worker + * process. + */ +export const startUsingMessagePort = (port: MessagePort) => { + _port = port; + port.start(); +}; + +/** + * Return a CLIP embedding of the given image. + * + * See: [Note: Natural language search using CLIP] + * + * The input is a opaque float32 array representing the image. The layout + * and exact encoding of the input is specific to our implementation and the + * ML model (CLIP) we use. + * + * @returns A CLIP embedding (an array of 512 floating point values). + */ +export const computeCLIPImageEmbedding = ( + input: Float32Array, +): Promise => + ensureFloat32Array(electronMLWorker("computeCLIPImageEmbedding", input)); + +/** + * Return a CLIP embedding of the given image if we already have the model + * downloaded and prepped. If the model is not available return `undefined`. + * + * This differs from the other sibling ML functions in that it doesn't wait + * for the model download to finish. It does trigger a model download, but + * then immediately returns `undefined`. At some future point, when the + * model downloaded finishes, calls to this function will start returning + * the result we seek. + * + * The reason for doing it in this asymmetric way is because CLIP text + * embeddings are used as part of deducing user initiated search results, + * and we don't want to block that interaction on a large network request. + * + * See: [Note: Natural language search using CLIP] + * + * @param text The string whose embedding we want to compute. + * + * @returns A CLIP embedding. + */ +export const computeCLIPTextEmbeddingIfAvailable = async ( + text: string, +): Promise => + ensureOptionalFloat32Array( + electronMLWorker("computeCLIPTextEmbeddingIfAvailable", text), + ); + +/** + * Detect faces in the given image using YOLO. + * + * Both the input and output are opaque binary data whose internal structure + * is specific to our implementation and the model (YOLO) we use. + */ +export const detectFaces = (input: Float32Array): Promise => + ensureFloat32Array(electronMLWorker("detectFaces", input)); + +/** + * Return a MobileFaceNet embeddings for the given faces. + * + * Both the input and output are opaque binary data whose internal structure + * is specific to our implementation and the model (MobileFaceNet) we use. + */ +export const computeFaceEmbeddings = ( + input: Float32Array, +): Promise => + ensureFloat32Array(electronMLWorker("computeFaceEmbeddings", input)); + +const ensureFloat32Array = async ( + pu: Promise, +): Promise => { + const u = await pu; + if (u instanceof Float32Array) return u; + throw new Error(`Expected a Float32Array but instead got ${typeof u}`); +}; + +const ensureOptionalFloat32Array = async ( + pu: Promise, +): Promise => { + const u = await pu; + if (u === undefined) return u; + if (u instanceof Float32Array) return u; + throw new Error(`Expected a Float32Array but instead got ${typeof u}`); +}; + +/** + * Make a call to the ML worker running in the Node.js layer using our + * hand-rolled RPC protocol. See: [Note: Node.js ML worker RPC protocol]. + */ +const electronMLWorker = async (method: string, p: string | Float32Array) => { + const port = _port; + if (!port) { + throw new Error( + "No MessagePort to communicate with Electron ML worker", + ); + } + + // Generate a unique nonce to identify this RPC interaction. + const id = Math.random(); + return new Promise((resolve, reject) => { + const handleMessage = (event: MessageEvent) => { + const response = RPCResponse.parse(event.data); + if (response.id != id) return; + port.removeEventListener("message", handleMessage); + const error = response.error; + if (error) reject(new Error(error)); + else resolve(response.result); + }; + port.addEventListener("message", handleMessage); + port.postMessage({ id, method, p }); + }); +}; + +const RPCResponse = z.object({ + id: z.number(), + result: z.any().optional(), + error: z.string().optional(), +}); diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 10e8f2fd7b..03b8adf2a2 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -9,7 +9,6 @@ import { ensure } from "@/utils/ensure"; import { wait } from "@/utils/promise"; import { DOMParser } from "@xmldom/xmldom"; import { expose } from "comlink"; -import { z } from "zod"; import downloadManager from "../download"; import { cmpNewLib2, extractRawExif } from "../exif"; import { getAllLocalFiles, getLocalTrashedFiles } from "../files"; @@ -33,6 +32,7 @@ import { type RemoteDerivedData, } from "./embedding"; import { faceIndexingVersion, indexFaces, type FaceIndex } from "./face"; +import { startUsingMessagePort } from "./worker-rpc"; import type { MLWorkerDelegate, MLWorkerElectron } from "./worker-types"; const idleDurationStart = 5; /* 5 seconds */ @@ -47,18 +47,9 @@ interface IndexableItem { remoteDerivedData: RemoteDerivedData | undefined; } -/** - * The port used to communicate with the Node.js ML worker process - * - * See: [Note: ML IPC] - * */ -let _port: MessagePort | undefined; - globalThis.onmessage = (event: MessageEvent) => { - if (event.data == "createMLWorker/port") { - _port = event.ports[0]; - _port?.start(); - } + if (event.data == "createMLWorker/port") + startUsingMessagePort(ensure(event.ports[0])); }; /** @@ -128,12 +119,6 @@ export class MLWorker { // need to monkey patch it (This also ensures that it is not tree // shaken). globalThis.DOMParser = DOMParser; - - void (async () => { - console.log("yyy calling foo with 3"); - const res = await electronMLWorker("foo", "3"); - console.log("yyy calling foo with 3 result", res); - })(); } /** @@ -266,40 +251,6 @@ export class MLWorker { expose(MLWorker); -/** - * Make a call to the ML worker running in the Node.js layer using our - * hand-rolled RPC protocol. See: [Note: Node.js ML worker RPC protocol]. - */ -const electronMLWorker = async (method: string, p: string) => { - const port = _port; - if (!port) { - throw new Error( - "No MessagePort to communicate with Electron ML worker", - ); - } - - // Generate a unique nonce to identify this RPC interaction. - const id = Math.random(); - return new Promise((resolve, reject) => { - const handleMessage = (event: MessageEvent) => { - const response = RPCResponse.parse(event.data); - if (response.id != id) return; - port.removeEventListener("message", handleMessage); - const error = response.error; - if (error) reject(new Error(error)); - else resolve(response.result); - }; - port.addEventListener("message", handleMessage); - port.postMessage({ id, method, p }); - }); -}; - -const RPCResponse = z.object({ - id: z.number(), - result: z.any().optional(), - error: z.string().optional(), -}); - /** * Find out files which need to be indexed. Then index the next batch of them. * From f2f7b483fdd471932bd787f35663a2c923bed632 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 09:40:08 +0530 Subject: [PATCH 28/48] 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. From 1ae0f9723c04fb8f78042f402f01b6b3ac6b2245 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 09:54:44 +0530 Subject: [PATCH 29/48] Fix 1 --- desktop/src/main/ipc.ts | 32 +++++++++++--------------- desktop/src/main/services/ml-worker.ts | 5 ++-- desktop/src/main/services/ml.ts | 2 +- 3 files changed, 17 insertions(+), 22 deletions(-) diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index 511460c6c7..fd6325c5cc 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -44,12 +44,7 @@ import { import { convertToJPEG, generateImageThumbnail } from "./services/image"; import { logout } from "./services/logout"; import { createMLWorker } from "./services/ml"; -import { - computeCLIPImageEmbedding, - computeCLIPTextEmbeddingIfAvailable, - computeFaceEmbeddings, - detectFaces, -} from "./services/ml-worker"; + import { encryptionKey, lastShownChangelogVersion, @@ -189,21 +184,22 @@ export const attachIPCHandlers = () => { // - ML - ipcMain.handle("computeCLIPImageEmbedding", (_, input: Float32Array) => - computeCLIPImageEmbedding(input), - ); + // TODO: + // ipcMain.handle("computeCLIPImageEmbedding", (_, input: Float32Array) => + // computeCLIPImageEmbedding(input), + // ); - ipcMain.handle("computeCLIPTextEmbeddingIfAvailable", (_, text: string) => - computeCLIPTextEmbeddingIfAvailable(text), - ); + // ipcMain.handle("computeCLIPTextEmbeddingIfAvailable", (_, text: string) => + // computeCLIPTextEmbeddingIfAvailable(text), + // ); - ipcMain.handle("detectFaces", (_, input: Float32Array) => - detectFaces(input), - ); + // ipcMain.handle("detectFaces", (_, input: Float32Array) => + // detectFaces(input), + // ); - ipcMain.handle("computeFaceEmbeddings", (_, input: Float32Array) => - computeFaceEmbeddings(input), - ); + // ipcMain.handle("computeFaceEmbeddings", (_, input: Float32Array) => + // computeFaceEmbeddings(input), + // ); // - Upload diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index 21553da3ec..7112b5be78 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -10,9 +10,6 @@ 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"; @@ -21,6 +18,8 @@ import * as ort from "onnxruntime-node"; import { ensure, wait } from "../utils/common"; import { writeStream } from "../utils/stream"; +const nodeEndpoint = require("comlink/dist/umd/node-adapter"); + /** * We cannot do * diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index eef362ccc4..66c699b5d2 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -67,7 +67,7 @@ import log from "../log"; export const createMLWorker = (window: BrowserWindow) => { const { port1, port2 } = new MessageChannelMain(); - const child = utilityProcess.fork(path.join(__dirname, "ml-util-test.js")); + const child = utilityProcess.fork(path.join(__dirname, "ml-worker.js")); const userDataPath = app.getPath("userData"); child.postMessage({ userDataPath }, [port1]); From b69d23028b884f263fd74b0b48af8c2aa23417e6 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 09:56:45 +0530 Subject: [PATCH 30/48] Remove test code --- desktop/src/main/services/ml-worker.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index 7112b5be78..1a1f5cd0d7 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -95,7 +95,6 @@ const parseInitData = (data: unknown) => { if ( data && typeof data == "object" && - "userDataPateh" in data && "userDataPath" in data && typeof data.userDataPath == "string" ) { From daed8a72dad89118d7467ebb4f37afd5740d6e6a Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 10:17:08 +0530 Subject: [PATCH 31/48] Only once --- web/packages/new/photos/services/ml/index.ts | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index b903e0188c..511b71e5cc 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -35,7 +35,7 @@ import { MLWorker } from "./worker"; let _isMLEnabled = false; /** Cached instance of the {@link ComlinkWorker} that wraps our web worker. */ -let _comlinkWorker: ComlinkWorker | undefined; +let _comlinkWorker: Promise> | undefined; /** * Subscriptions to {@link MLStatus}. @@ -52,10 +52,8 @@ let _mlStatusListeners: (() => void)[] = []; let _mlStatusSnapshot: MLStatus | undefined; /** Lazily created, cached, instance of {@link MLWorker}. */ -const worker = async () => { - if (!_comlinkWorker) _comlinkWorker = await createComlinkWorker(); - return _comlinkWorker.remote; -}; +const worker = () => + (_comlinkWorker ??= createComlinkWorker()).then((cw) => cw.remote); const createComlinkWorker = async () => { const electron = ensureElectron(); From a97e01171a1a3068761b95bd1ff46280744354e9 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 11:27:15 +0530 Subject: [PATCH 32/48] Commit incorrect but original motivations --- desktop/src/main/services/ml-worker.ts | 1 + desktop/src/main/utils/comlink-endpoint.ts | 66 ++++++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 desktop/src/main/utils/comlink-endpoint.ts diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index 1a1f5cd0d7..af5d2add62 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -63,6 +63,7 @@ process.parentPort.once("message", (e) => { parseInitData(e.data); const port = ensure(e.ports[0]); + port.on("message", (me: Electron.MessageEvent) => {}); expose( { computeCLIPImageEmbedding, diff --git a/desktop/src/main/utils/comlink-endpoint.ts b/desktop/src/main/utils/comlink-endpoint.ts new file mode 100644 index 0000000000..ee87d4176e --- /dev/null +++ b/desktop/src/main/utils/comlink-endpoint.ts @@ -0,0 +1,66 @@ +import type { Endpoint } from "comlink"; +import type { MessagePortMain } from "electron"; + +/** + * An adaptation of the `nodeEndpoint` function from comlink suitable for use in + * TypeScript with an Electron utility process. + * + * This is an adaption of the + * + * Comlink provides a `nodeEndpoint` [function][1] to allow a Node worker_thread + * to be treated as an {@link Endpoint} and be used with comlink. + * + * The first issue we run into when using it is that this the function is not + * exported as part of the normal comlink.d.ts. Accessing it via this + * [workaround][2] doesn't work for us either since we cannot currently change + * our package type to "module". + * + * We could skirt around that by doing + * + * const nodeEndpoint = require("comlink/dist/umd/node-adapter"); + * + * and silencing tsc and eslint. However, we then run into a different issue: + * the comlink implementation of the adapter adds an extra layer of nesting. + * This line: + * + * eh({ data } as MessageEvent); + * + * Should be + * + * eh(data) + * + * I don't currently know if it is because of an impedance mismatch between + * Node's worker_threads and Electron's UtilityProcesses, or if it is something + * else that I'm doing wrong somewhere else causing this to happen. + * + * To solve both these issues, we create this variant. This also removes the + * need for us to type cast when passing MessagePortMain. + * + * References: + * 1. https://github.com/GoogleChromeLabs/comlink/blob/main/src/node-adapter.ts + * 2. https://github.com/GoogleChromeLabs/comlink/pull/542 + * 3. https://github.com/GoogleChromeLabs/comlink/issues/129 + */ +export const messagePortMainEndpoint = (mp: MessagePortMain): Endpoint => { + const listeners = new WeakMap(); + return { + postMessage: mp.postMessage.bind(mp), + addEventListener: (_, eh) => { + const l = (data: Electron.MessageEvent) => + "handleEvent" in eh + ? eh.handleEvent({ data } as MessageEvent) + : eh(data as unknown as MessageEvent); + mp.on("message", (data) => { + l(data); + }); + listeners.set(eh, l); + }, + removeEventListener: (_, eh) => { + const l = listeners.get(eh); + if (!l) return; + mp.off("message", l); + listeners.delete(eh); + }, + start: mp.start.bind(mp), + }; +}; From 62f723e50cf2ea29753a06b2f7799dec9dbe5539 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 11:33:18 +0530 Subject: [PATCH 33/48] Adapt --- desktop/src/main/services/ml-worker.ts | 7 +-- desktop/src/main/utils/comlink-endpoint.ts | 50 ++++++---------------- 2 files changed, 15 insertions(+), 42 deletions(-) diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index af5d2add62..e028dd01d5 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -15,11 +15,10 @@ import { existsSync } from "fs"; import fs from "node:fs/promises"; import path from "node:path"; import * as ort from "onnxruntime-node"; +import { messagePortMainEndpoint } from "../utils/comlink-endpoint"; import { ensure, wait } from "../utils/common"; import { writeStream } from "../utils/stream"; -const nodeEndpoint = require("comlink/dist/umd/node-adapter"); - /** * We cannot do * @@ -63,7 +62,6 @@ process.parentPort.once("message", (e) => { parseInitData(e.data); const port = ensure(e.ports[0]); - port.on("message", (me: Electron.MessageEvent) => {}); expose( { computeCLIPImageEmbedding, @@ -71,8 +69,7 @@ process.parentPort.once("message", (e) => { detectFaces, computeFaceEmbeddings, }, - // eslint-disable-next-line @typescript-eslint/no-unsafe-argument, @typescript-eslint/no-explicit-any - nodeEndpoint(port as unknown as any), + messagePortMainEndpoint(port), ); // port.on("message", (request) => { // void handleMessageFromRenderer(request.data).then((response) => diff --git a/desktop/src/main/utils/comlink-endpoint.ts b/desktop/src/main/utils/comlink-endpoint.ts index ee87d4176e..4572837306 100644 --- a/desktop/src/main/utils/comlink-endpoint.ts +++ b/desktop/src/main/utils/comlink-endpoint.ts @@ -5,48 +5,24 @@ import type { MessagePortMain } from "electron"; * An adaptation of the `nodeEndpoint` function from comlink suitable for use in * TypeScript with an Electron utility process. * - * This is an adaption of the + * This is an adaption of the following function from comlink: + * https://github.com/GoogleChromeLabs/comlink/blob/main/src/node-adapter.ts * - * Comlink provides a `nodeEndpoint` [function][1] to allow a Node worker_thread - * to be treated as an {@link Endpoint} and be used with comlink. - * - * The first issue we run into when using it is that this the function is not - * exported as part of the normal comlink.d.ts. Accessing it via this - * [workaround][2] doesn't work for us either since we cannot currently change - * our package type to "module". - * - * We could skirt around that by doing - * - * const nodeEndpoint = require("comlink/dist/umd/node-adapter"); - * - * and silencing tsc and eslint. However, we then run into a different issue: - * the comlink implementation of the adapter adds an extra layer of nesting. - * This line: - * - * eh({ data } as MessageEvent); - * - * Should be - * - * eh(data) - * - * I don't currently know if it is because of an impedance mismatch between - * Node's worker_threads and Electron's UtilityProcesses, or if it is something - * else that I'm doing wrong somewhere else causing this to happen. - * - * To solve both these issues, we create this variant. This also removes the - * need for us to type cast when passing MessagePortMain. - * - * References: - * 1. https://github.com/GoogleChromeLabs/comlink/blob/main/src/node-adapter.ts - * 2. https://github.com/GoogleChromeLabs/comlink/pull/542 - * 3. https://github.com/GoogleChromeLabs/comlink/issues/129 + * It has been modified (somewhat hackily) to be useful with an Electron + * MessagePortMain instead of a Node.js worker_thread. Only things that we + * currently need have been made to work as you can see by the abuntant type + * casts. Caveat emptor. */ export const messagePortMainEndpoint = (mp: MessagePortMain): Endpoint => { - const listeners = new WeakMap(); + type NL = EventListenerOrEventListenerObject; + type EL = (data: Electron.MessageEvent) => void; + const listeners = new WeakMap(); return { - postMessage: mp.postMessage.bind(mp), + postMessage: (message, transfer) => { + mp.postMessage(message, transfer as unknown as MessagePortMain[]); + }, addEventListener: (_, eh) => { - const l = (data: Electron.MessageEvent) => + const l: EL = (data) => "handleEvent" in eh ? eh.handleEvent({ data } as MessageEvent) : eh(data as unknown as MessageEvent); From 6ad27a2d42d60e08284423be33cda799805a8751 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 11:58:22 +0530 Subject: [PATCH 34/48] Cleanup --- desktop/src/main/services/ml-worker.ts | 18 +-- desktop/src/main/services/ml.ts | 3 + .../utils/{comlink-endpoint.ts => comlink.ts} | 2 +- web/packages/base/types/ipc.ts | 64 ++++++++- web/packages/new/photos/services/ml/clip.ts | 9 +- web/packages/new/photos/services/ml/face.ts | 12 +- .../new/photos/services/ml/worker-rpc.ts | 131 ------------------ .../new/photos/services/ml/worker-types.ts | 20 +-- web/packages/new/photos/services/ml/worker.ts | 29 ++-- 9 files changed, 95 insertions(+), 193 deletions(-) rename desktop/src/main/utils/{comlink-endpoint.ts => comlink.ts} (99%) delete mode 100644 web/packages/new/photos/services/ml/worker-rpc.ts diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index e028dd01d5..379a1b1f4a 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -15,7 +15,7 @@ import { existsSync } from "fs"; import fs from "node:fs/promises"; import path from "node:path"; import * as ort from "onnxruntime-node"; -import { messagePortMainEndpoint } from "../utils/comlink-endpoint"; +import { messagePortMainEndpoint } from "../utils/comlink"; import { ensure, wait } from "../utils/common"; import { writeStream } from "../utils/stream"; @@ -54,14 +54,12 @@ const log = { const mainProcess = (method: string, param: unknown) => process.parentPort.postMessage({ method, p: param }); -log.debugString( - `Started ML worker process with args ${process.argv.join(" ")}`, -); +log.debugString(`Started ML worker process`); process.parentPort.once("message", (e) => { + // Initialize ourselves with the data we got from our parent. parseInitData(e.data); - - const port = ensure(e.ports[0]); + // Expose an expose( { computeCLIPImageEmbedding, @@ -69,14 +67,8 @@ process.parentPort.once("message", (e) => { detectFaces, computeFaceEmbeddings, }, - messagePortMainEndpoint(port), + messagePortMainEndpoint(ensure(e.ports[0])), ); - // port.on("message", (request) => { - // void handleMessageFromRenderer(request.data).then((response) => - // port.postMessage(response), - // ); - // }); - // port.start(); }); /** diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 66c699b5d2..05642fc95f 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -63,6 +63,9 @@ import log from "../log"; * worker can directly talk to each other! * * Node.js utility process <-> Renderer web worker + * + * The RPC protocol is handled using comlink on both ends. The port itself needs + * to be relayed using `postMessage`. */ export const createMLWorker = (window: BrowserWindow) => { const { port1, port2 } = new MessageChannelMain(); diff --git a/desktop/src/main/utils/comlink-endpoint.ts b/desktop/src/main/utils/comlink.ts similarity index 99% rename from desktop/src/main/utils/comlink-endpoint.ts rename to desktop/src/main/utils/comlink.ts index 4572837306..d2006e795b 100644 --- a/desktop/src/main/utils/comlink-endpoint.ts +++ b/desktop/src/main/utils/comlink.ts @@ -10,7 +10,7 @@ import type { MessagePortMain } from "electron"; * * It has been modified (somewhat hackily) to be useful with an Electron * MessagePortMain instead of a Node.js worker_thread. Only things that we - * currently need have been made to work as you can see by the abuntant type + * currently need have been made to work as you can see by the abundant type * casts. Caveat emptor. */ export const messagePortMainEndpoint = (mp: MessagePortMain): Endpoint => { diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index 4aca8b865d..748490ed58 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -335,12 +335,15 @@ export interface Electron { // - ML /** - * Create a new ML session. + * Create a new ML worker. * * This creates a new Node.js utility process, and sets things up so that we * can communicate directly with that utility process using a * {@link MessagePort} that gets posted using "createMLWorker/port". * + * At the other end of that port will be an object that conforms to the + * {@link ElectronMLWorker} interface. + * * For more details about the IPC flow, see: [Note: ML IPC]. * * Note: For simplicity of implementation, we assume that there is at most @@ -535,6 +538,65 @@ export interface Electron { clearPendingUploads: () => Promise; } +/** + * The shape of the object exposed by the Node.js ML worker process on the + * message port that the web layer obtains by doing {@link createMLWorker}. + */ +export interface ElectronMLWorker { + /** + * Return a CLIP embedding of the given image. + * + * See: [Note: Natural language search using CLIP] + * + * The input is a opaque float32 array representing the image. The layout + * and exact encoding of the input is specific to our implementation and the + * ML model (CLIP) we use. + * + * @returns A CLIP embedding (an array of 512 floating point values). + */ + computeCLIPImageEmbedding: (input: Float32Array) => Promise; + + /** + * Return a CLIP embedding of the given image if we already have the model + * downloaded and prepped. If the model is not available return `undefined`. + * + * This differs from the other sibling ML functions in that it doesn't wait + * for the model download to finish. It does trigger a model download, but + * then immediately returns `undefined`. At some future point, when the + * model downloaded finishes, calls to this function will start returning + * the result we seek. + * + * The reason for doing it in this asymmetric way is because CLIP text + * embeddings are used as part of deducing user initiated search results, + * and we don't want to block that interaction on a large network request. + * + * See: [Note: Natural language search using CLIP] + * + * @param text The string whose embedding we want to compute. + * + * @returns A CLIP embedding. + */ + computeCLIPTextEmbeddingIfAvailable: ( + text: string, + ) => Promise; + + /** + * Detect faces in the given image using YOLO. + * + * Both the input and output are opaque binary data whose internal structure + * is specific to our implementation and the model (YOLO) we use. + */ + detectFaces: (input: Float32Array) => Promise; + + /** + * Return a MobileFaceNet embeddings for the given faces. + * + * Both the input and output are opaque binary data whose internal structure + * is specific to our implementation and the model (MobileFaceNet) we use. + */ + computeFaceEmbeddings: (input: Float32Array) => Promise; +} + /** * Errors that have special semantics on the web side. * diff --git a/web/packages/new/photos/services/ml/clip.ts b/web/packages/new/photos/services/ml/clip.ts index eecf7e2209..f6230f2466 100644 --- a/web/packages/new/photos/services/ml/clip.ts +++ b/web/packages/new/photos/services/ml/clip.ts @@ -1,9 +1,8 @@ -import type { Electron } from "@/base/types/ipc"; +import type { Electron, ElectronMLWorker } from "@/base/types/ipc"; import type { ImageBitmapAndData } from "./blob"; import { clipIndexes } from "./db"; import { pixelRGBBicubic } from "./image"; import { dotProduct, norm } from "./math"; -import type { MLWorkerElectron } from "./worker-types"; /** * The version of the CLIP indexing pipeline implemented by the current client. @@ -98,19 +97,19 @@ export type LocalCLIPIndex = CLIPIndex & { * be set to the {@link UploadItem} that was uploaded. This way, we can directly * use the on-disk file instead of needing to download the original from remote. * - * @param electron The {@link MLWorkerElectron} instance that allows us to call + * @param electron The {@link ElectronMLWorker} instance that allows us to call * our Node.js layer to run the ONNX inference. */ export const indexCLIP = async ( image: ImageBitmapAndData, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => ({ embedding: await computeEmbedding(image.data, electron), }); const computeEmbedding = async ( imageData: ImageData, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => { const clipInput = convertToCLIPInput(imageData); return normalized(await electron.computeCLIPImageEmbedding(clipInput)); diff --git a/web/packages/new/photos/services/ml/face.ts b/web/packages/new/photos/services/ml/face.ts index 910970d3b9..7ecbf06002 100644 --- a/web/packages/new/photos/services/ml/face.ts +++ b/web/packages/new/photos/services/ml/face.ts @@ -7,6 +7,7 @@ // /* eslint-disable @typescript-eslint/no-non-null-assertion */ +import type { ElectronMLWorker } from "@/base/types/ipc"; import type { EnteFile } from "@/new/photos/types/file"; import { Matrix } from "ml-matrix"; import { getSimilarityTransformation } from "similarity-transformation"; @@ -24,7 +25,6 @@ import { warpAffineFloat32List, } from "./image"; import { clamp } from "./math"; -import type { MLWorkerElectron } from "./worker-types"; /** * The version of the face indexing pipeline implemented by the current client. @@ -236,13 +236,13 @@ export interface Box { * * @param image The file's contents. * - * @param electron The {@link MLWorkerElectron} instance that allows us to call + * @param electron The {@link ElectronMLWorker} instance that allows us to call * our Node.js layer to run the ONNX inference. */ export const indexFaces = async ( enteFile: EnteFile, { data: imageData }: ImageBitmapAndData, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => ({ width: imageData.width, height: imageData.height, @@ -252,7 +252,7 @@ export const indexFaces = async ( const indexFaces_ = async ( fileID: number, imageData: ImageData, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => { const { width, height } = imageData; const imageDimensions = { width, height }; @@ -316,7 +316,7 @@ const indexFaces_ = async ( */ const detectFaces = async ( imageData: ImageData, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => { const rect = ({ width, height }: Dimensions) => ({ x: 0, @@ -878,7 +878,7 @@ const mobileFaceNetEmbeddingSize = 192; */ const computeEmbeddings = async ( faceData: Float32Array, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => { const outputData = await electron.computeFaceEmbeddings(faceData); diff --git a/web/packages/new/photos/services/ml/worker-rpc.ts b/web/packages/new/photos/services/ml/worker-rpc.ts deleted file mode 100644 index 782c0a3bd3..0000000000 --- a/web/packages/new/photos/services/ml/worker-rpc.ts +++ /dev/null @@ -1,131 +0,0 @@ -import { z } from "zod"; - -/** - * The port used to communicate with the Node.js ML worker process - * - * See: [Note: ML IPC] - * */ -let _port: MessagePort | undefined; - -/** - * Use the given {@link MessagePort} to communicate with the Node.js ML worker - * process. - */ -export const startUsingMessagePort = (port: MessagePort) => { - _port = port; - port.start(); -}; - -/** - * Return a CLIP embedding of the given image. - * - * See: [Note: Natural language search using CLIP] - * - * The input is a opaque float32 array representing the image. The layout - * and exact encoding of the input is specific to our implementation and the - * ML model (CLIP) we use. - * - * @returns A CLIP embedding (an array of 512 floating point values). - */ -export const computeCLIPImageEmbedding = ( - input: Float32Array, -): Promise => - ensureFloat32Array(electronMLWorker("computeCLIPImageEmbedding", input)); - -/** - * Return a CLIP embedding of the given image if we already have the model - * downloaded and prepped. If the model is not available return `undefined`. - * - * This differs from the other sibling ML functions in that it doesn't wait - * for the model download to finish. It does trigger a model download, but - * then immediately returns `undefined`. At some future point, when the - * model downloaded finishes, calls to this function will start returning - * the result we seek. - * - * The reason for doing it in this asymmetric way is because CLIP text - * embeddings are used as part of deducing user initiated search results, - * and we don't want to block that interaction on a large network request. - * - * See: [Note: Natural language search using CLIP] - * - * @param text The string whose embedding we want to compute. - * - * @returns A CLIP embedding. - */ -export const computeCLIPTextEmbeddingIfAvailable = async ( - text: string, -): Promise => - ensureOptionalFloat32Array( - electronMLWorker("computeCLIPTextEmbeddingIfAvailable", text), - ); - -/** - * Detect faces in the given image using YOLO. - * - * Both the input and output are opaque binary data whose internal structure - * is specific to our implementation and the model (YOLO) we use. - */ -export const detectFaces = (input: Float32Array): Promise => - ensureFloat32Array(electronMLWorker("detectFaces", input)); - -/** - * Return a MobileFaceNet embeddings for the given faces. - * - * Both the input and output are opaque binary data whose internal structure - * is specific to our implementation and the model (MobileFaceNet) we use. - */ -export const computeFaceEmbeddings = ( - input: Float32Array, -): Promise => - ensureFloat32Array(electronMLWorker("computeFaceEmbeddings", input)); - -const ensureFloat32Array = async ( - pu: Promise, -): Promise => { - const u = await pu; - if (u instanceof Float32Array) return u; - throw new Error(`Expected a Float32Array but instead got ${typeof u}`); -}; - -const ensureOptionalFloat32Array = async ( - pu: Promise, -): Promise => { - const u = await pu; - if (u === undefined) return u; - if (u instanceof Float32Array) return u; - throw new Error(`Expected a Float32Array but instead got ${typeof u}`); -}; - -/** - * Make a call to the ML worker running in the Node.js layer using our - * hand-rolled RPC protocol. See: [Note: Node.js ML worker RPC protocol]. - */ -const electronMLWorker = async (method: string, p: string | Float32Array) => { - const port = _port; - if (!port) { - throw new Error( - "No MessagePort to communicate with Electron ML worker", - ); - } - - // Generate a unique nonce to identify this RPC interaction. - const id = Math.random(); - return new Promise((resolve, reject) => { - const handleMessage = (event: MessageEvent) => { - const response = RPCResponse.parse(event.data); - if (response.id != id) return; - port.removeEventListener("message", handleMessage); - const error = response.error; - if (error) reject(new Error(error)); - else resolve(response.result); - }; - port.addEventListener("message", handleMessage); - port.postMessage({ id, method, p }); - }); -}; - -const RPCResponse = z.object({ - id: z.number(), - result: z.any().optional(), - error: z.string().optional(), -}); diff --git a/web/packages/new/photos/services/ml/worker-types.ts b/web/packages/new/photos/services/ml/worker-types.ts index 1eb43933a3..a83a215ea4 100644 --- a/web/packages/new/photos/services/ml/worker-types.ts +++ b/web/packages/new/photos/services/ml/worker-types.ts @@ -1,22 +1,8 @@ /** - * @file Type for the objects shared (as a Comlink proxy) by the main thread and - * the ML worker. + * @file Types for the objects shared (as a Comlink proxy) by the main thread + * and the ML worker. */ -/** - * A subset of {@link Electron} provided to the {@link MLWorker}. - * - * `globalThis.electron` does not exist in the execution context of web workers. - * So instead, we manually provide a proxy object of type - * {@link MLWorkerElectron} that exposes a subset of the functions from - * {@link Electron} that are needed by the code running in the ML web worker. - */ -export interface MLWorkerElectron { - detectFaces: (input: Float32Array) => Promise; - computeFaceEmbeddings: (input: Float32Array) => Promise; - computeCLIPImageEmbedding: (input: Float32Array) => Promise; -} - /** * Callbacks invoked by the worker at various points in the indexing pipeline to * notify the main thread of events it might be interested in. @@ -25,7 +11,7 @@ export interface MLWorkerDelegate { /** * Called whenever a file is processed during indexing. * - * It is called both when the indexing was successful or failed. + * It is called both when the indexing was successful or it failed. */ workerDidProcessFile: () => void; } diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index d27f3021f4..e137531f6d 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -3,6 +3,7 @@ import { isHTTP4xxError } from "@/base/http"; import { getKVN } from "@/base/kv"; import { ensureAuthToken } from "@/base/local-user"; import log from "@/base/log"; +import type { ElectronMLWorker } from "@/base/types/ipc"; import type { EnteFile } from "@/new/photos/types/file"; import { fileLogID } from "@/new/photos/utils/file"; import { ensure } from "@/utils/ensure"; @@ -32,8 +33,7 @@ import { type RemoteDerivedData, } from "./embedding"; import { faceIndexingVersion, indexFaces, type FaceIndex } from "./face"; -import { startUsingMessagePort } from "./worker-rpc"; -import type { MLWorkerDelegate, MLWorkerElectron } from "./worker-types"; +import type { MLWorkerDelegate } from "./worker-types"; const idleDurationStart = 5; /* 5 seconds */ const idleDurationMax = 16 * 60; /* 16 minutes */ @@ -47,11 +47,6 @@ interface IndexableItem { remoteDerivedData: RemoteDerivedData | undefined; } -globalThis.onmessage = (event: MessageEvent) => { - if (event.data == "createMLWorker/port") - startUsingMessagePort(ensure(event.ports[0])); -}; - /** * Run operations related to machine learning (e.g. indexing) in a Web Worker. * @@ -75,7 +70,7 @@ globalThis.onmessage = (event: MessageEvent) => { * - "idle": in between state transitions. */ export class MLWorker { - private electron: MLWorkerElectron | undefined; + private electron: ElectronMLWorker | undefined; private delegate: MLWorkerDelegate | undefined; private state: "idle" | "indexing" = "idle"; private liveQ: IndexableItem[] = []; @@ -88,20 +83,16 @@ export class MLWorker { * This is conceptually the constructor, however it is easier to have this * as a separate function to avoid complicating the comlink types further. * - * @param electron The {@link MLWorkerElectron} that allows the worker to - * use the functionality provided by our Node.js layer when running in the - * context of our desktop app. + * @param port A {@link MessagePort} that allows us to communicate with an + * Electron utility process running in the Node.js layer of our desktop app, + * exposing an object that conforms to the {@link ElectronMLWorker} + * interface. * * @param delegate The {@link MLWorkerDelegate} the worker can use to inform * the main thread of interesting events. */ async init(port: MessagePort, delegate: MLWorkerDelegate) { - // this.electron = electron; - this.electron = wrap(port); /* mlWorkerElectron = { - detectFaces: electron.detectFaces, - computeFaceEmbeddings: electron.computeFaceEmbeddings, - computeCLIPImageEmbedding: electron.computeCLIPImageEmbedding, - };*/ + this.electron = wrap(port); 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. @@ -267,7 +258,7 @@ expose(MLWorker); */ const indexNextBatch = async ( items: IndexableItem[], - electron: MLWorkerElectron, + electron: ElectronMLWorker, delegate: MLWorkerDelegate | undefined, ) => { // Don't try to index if we wouldn't be able to upload them anyway. The @@ -385,7 +376,7 @@ const syncWithLocalFilesAndGetFilesToIndex = async ( */ const index = async ( { enteFile, uploadItem, remoteDerivedData }: IndexableItem, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ) => { const f = fileLogID(enteFile); const fileID = enteFile.id; From e55a7facc37ac874ceaf019ba006046f2a1d0e67 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 12:11:53 +0530 Subject: [PATCH 35/48] Replace our homebrew RPC --- desktop/src/main/ipc.ts | 19 ------------------- desktop/src/main/services/ml-worker.ts | 3 ++- web/apps/photos/src/services/logout.ts | 2 +- web/packages/new/photos/services/ml/index.ts | 16 +++++++--------- 4 files changed, 10 insertions(+), 30 deletions(-) diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index fd6325c5cc..6b837d8b4d 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -182,25 +182,6 @@ export const attachIPCHandlers = () => { ) => ffmpegExec(command, dataOrPathOrZipItem, outputFileExtension), ); - // - ML - - // TODO: - // ipcMain.handle("computeCLIPImageEmbedding", (_, input: Float32Array) => - // computeCLIPImageEmbedding(input), - // ); - - // ipcMain.handle("computeCLIPTextEmbeddingIfAvailable", (_, text: string) => - // computeCLIPTextEmbeddingIfAvailable(text), - // ); - - // ipcMain.handle("detectFaces", (_, input: Float32Array) => - // detectFaces(input), - // ); - - // ipcMain.handle("computeFaceEmbeddings", (_, input: Float32Array) => - // computeFaceEmbeddings(input), - // ); - // - Upload ipcMain.handle("listZipItems", (_, zipPath: string) => diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index 379a1b1f4a..48f819c8bf 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -59,7 +59,8 @@ log.debugString(`Started ML worker process`); process.parentPort.once("message", (e) => { // Initialize ourselves with the data we got from our parent. parseInitData(e.data); - // Expose an + // Expose an instance of `ElectronMLWorker` on the port we got from our + // parent. expose( { computeCLIPImageEmbedding, diff --git a/web/apps/photos/src/services/logout.ts b/web/apps/photos/src/services/logout.ts index 9722757689..ab4b1aa6c5 100644 --- a/web/apps/photos/src/services/logout.ts +++ b/web/apps/photos/src/services/logout.ts @@ -22,7 +22,7 @@ export const photosLogout = async () => { // See: [Note: Caching IDB instances in separate execution contexts]. try { - terminateMLWorker(); + await terminateMLWorker(); } catch (e) { ignoreError("face", e); } diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 511b71e5cc..5c114c0eb1 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -69,12 +69,10 @@ const createComlinkWorker = async () => { new Worker(new URL("worker.ts", import.meta.url)), ); - await cw.remote.then((w) => { - // Pass the message port to our web worker. - // cw.worker.postMessage("createMLWorker/port", [messagePort]); - // Initialize it. - return w.init(transfer(messagePort, [messagePort]), proxy(delegate)); - }); + await cw.remote.then((w) => + // Forward the port to the web worker. + w.init(transfer(messagePort, [messagePort]), proxy(delegate)), + ); return cw; }; @@ -88,9 +86,9 @@ const createComlinkWorker = async () => { * * It is also called when the user pauses or disables ML. */ -export const terminateMLWorker = () => { +export const terminateMLWorker = async () => { if (_comlinkWorker) { - _comlinkWorker.terminate(); + await _comlinkWorker.then((cw) => cw.terminate()); _comlinkWorker = undefined; } }; @@ -188,7 +186,7 @@ export const disableML = async () => { await updateIsMLEnabledRemote(false); setIsMLEnabledLocal(false); _isMLEnabled = false; - terminateMLWorker(); + await terminateMLWorker(); triggerStatusUpdate(); }; From 5a3838be342ac2bb9656b67bb6a27b12c2c3b0ea Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 12:30:15 +0530 Subject: [PATCH 36/48] Route via workers --- web/apps/photos/src/services/searchService.ts | 5 ++-- web/packages/new/photos/services/ml/clip.ts | 20 ++++------------ web/packages/new/photos/services/ml/index.ts | 17 +++++++++++++ .../new/photos/services/ml/worker-types.ts | 14 +++++++++-- web/packages/new/photos/services/ml/worker.ts | 24 +++++++++++++++---- 5 files changed, 56 insertions(+), 24 deletions(-) diff --git a/web/apps/photos/src/services/searchService.ts b/web/apps/photos/src/services/searchService.ts index 8b90652e64..750a1fb186 100644 --- a/web/apps/photos/src/services/searchService.ts +++ b/web/apps/photos/src/services/searchService.ts @@ -1,13 +1,12 @@ import { isDesktop } from "@/base/app"; -import { ensureElectron } from "@/base/electron"; import log from "@/base/log"; import { FileType } from "@/media/file-type"; import { + clipMatches, isMLEnabled, isMLSupported, mlStatusSnapshot, } from "@/new/photos/services/ml"; -import { clipMatches } from "@/new/photos/services/ml/clip"; import type { Person } from "@/new/photos/services/ml/people"; import { EnteFile } from "@/new/photos/types/file"; import * as chrono from "chrono-node"; @@ -374,7 +373,7 @@ const searchClip = async ( searchPhrase: string, ): Promise => { if (!isMLEnabled()) return undefined; - const matches = await clipMatches(searchPhrase, ensureElectron()); + const matches = await clipMatches(searchPhrase); log.debug(() => ["clip/scores", matches]); return matches; }; diff --git a/web/packages/new/photos/services/ml/clip.ts b/web/packages/new/photos/services/ml/clip.ts index f6230f2466..78eff1c04d 100644 --- a/web/packages/new/photos/services/ml/clip.ts +++ b/web/packages/new/photos/services/ml/clip.ts @@ -1,8 +1,9 @@ -import type { Electron, ElectronMLWorker } from "@/base/types/ipc"; +import type { ElectronMLWorker } from "@/base/types/ipc"; import type { ImageBitmapAndData } from "./blob"; import { clipIndexes } from "./db"; import { pixelRGBBicubic } from "./image"; import { dotProduct, norm } from "./math"; +import type { CLIPMatches } from "./worker-types"; /** * The version of the CLIP indexing pipeline implemented by the current client. @@ -166,26 +167,15 @@ const normalized = (embedding: Float32Array) => { }; /** - * Use CLIP to perform a natural language search over image embeddings. - * - * @param searchPhrase The text entered by the user in the search box. - * - * @param electron The {@link Electron} instance to use to communicate with the - * native code running in our desktop app (the embedding happens in the native - * layer). - * - * It returns file (IDs) that should be shown in the search results. They're - * returned as a map from fileIDs to the scores they got (higher is better). - * This map will only contains entries whose score was above our minimum - * threshold. + * Find the files whose CLIP embedding "matches" the given {@link searchPhrase}. * * The result can also be `undefined`, which indicates that the download for the * ML model is still in progress (trying again later should succeed). */ export const clipMatches = async ( searchPhrase: string, - electron: Electron, -): Promise | undefined> => { + electron: ElectronMLWorker, +): Promise => { const t = await electron.computeCLIPTextEmbeddingIfAvailable(searchPhrase); if (!t) return undefined; diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 5c114c0eb1..ba8083ad74 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -19,6 +19,7 @@ import type { UploadItem } from "../upload/types"; import { regenerateFaceCrops } from "./crop"; import { clearMLDB, faceIndex, indexableAndIndexedCounts } from "./db"; import { MLWorker } from "./worker"; +import type { CLIPMatches } from "./worker-types"; /** * In-memory flag that tracks if ML is enabled. @@ -392,6 +393,22 @@ const setInterimScheduledStatus = () => { const workerDidProcessFile = throttled(updateMLStatusSnapshot, 2000); +/** + * Use CLIP to perform a natural language search over image embeddings. + * + * @param searchPhrase The text entered by the user in the search box. + * + * It returns file (IDs) that should be shown in the search results, along with + * their scores. + * + * The result can also be `undefined`, which indicates that the download for the + * ML model is still in progress (trying again later should succeed). + */ +export const clipMatches = ( + searchPhrase: string, +): Promise => + worker().then((w) => w.clipMatches(searchPhrase)); + /** * Return the IDs of all the faces in the given {@link enteFile} that are not * associated with a person cluster. diff --git a/web/packages/new/photos/services/ml/worker-types.ts b/web/packages/new/photos/services/ml/worker-types.ts index a83a215ea4..72d6bce61b 100644 --- a/web/packages/new/photos/services/ml/worker-types.ts +++ b/web/packages/new/photos/services/ml/worker-types.ts @@ -1,6 +1,5 @@ /** - * @file Types for the objects shared (as a Comlink proxy) by the main thread - * and the ML worker. + * @file Types for the objects shared between the main thread and the ML worker. */ /** @@ -15,3 +14,14 @@ export interface MLWorkerDelegate { */ workerDidProcessFile: () => void; } + +/** + * The result of file ids that should be considered as matches for a particular + * search phrase, each with their associated score. + * + * This is a map of file (IDs) that should be shown in the search results. + * They're returned as a map from fileIDs to the scores they got (higher is + * better). This map will only contains entries whose score was above our + * minimum threshold. + */ +export type CLIPMatches = Map; diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index e137531f6d..cbeba5f844 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -19,7 +19,12 @@ import { indexableBlobs, type ImageBitmapAndData, } from "./blob"; -import { clipIndexingVersion, indexCLIP, type CLIPIndex } from "./clip"; +import { + clipIndexingVersion, + clipMatches, + indexCLIP, + type CLIPIndex, +} from "./clip"; import { saveFaceCrops } from "./crop"; import { indexableFileIDs, @@ -33,7 +38,7 @@ import { type RemoteDerivedData, } from "./embedding"; import { faceIndexingVersion, indexFaces, type FaceIndex } from "./face"; -import type { MLWorkerDelegate } from "./worker-types"; +import type { CLIPMatches, MLWorkerDelegate } from "./worker-types"; const idleDurationStart = 5; /* 5 seconds */ const idleDurationMax = 16 * 60; /* 16 minutes */ @@ -68,6 +73,9 @@ interface IndexableItem { * - "backfillq": fetching remote embeddings of unindexed items, and then * indexing them if needed, * - "idle": in between state transitions. + * + * In addition, MLWorker can also be invoked for interactive tasks: in + * particular, for finding the closest CLIP match when the user does a search. */ export class MLWorker { private electron: ElectronMLWorker | undefined; @@ -178,6 +186,13 @@ export class MLWorker { return this.state == "indexing"; } + /** + * Find {@link CLIPMatches} for a given {@link searchPhrase}. + */ + async clipMatches(searchPhrase: string): Promise { + return clipMatches(searchPhrase, ensure(this.electron)); + } + private async tick() { log.debug(() => [ "ml/tick", @@ -226,7 +241,7 @@ export class MLWorker { } /** Return the next batch of items to backfill (if any). */ - async backfillQ() { + private async backfillQ() { const userID = ensure(await getKVN("userID")); // Find files that our local DB thinks need syncing. const filesByID = await syncWithLocalFilesAndGetFilesToIndex( @@ -278,7 +293,8 @@ const indexNextBatch = async ( try { await index(item, electron); delegate?.workerDidProcessFile(); - // Possibly unnecessary, but let us drain the microtask queue. + // Let us drain the microtask queue. This also gives a chance for other + // interactive tasks like `clipMatches` to run. await wait(0); } catch (e) { log.warn(`Skipping unindexable file ${item.enteFile.id}`, e); From 3a5843f5324c0813d6e2d9a47dd6d73563397bbc Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 12:34:30 +0530 Subject: [PATCH 37/48] tail --- web/packages/new/photos/services/ml/blob.ts | 13 +++++++------ web/packages/new/photos/utils/native-stream.ts | 7 +++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/web/packages/new/photos/services/ml/blob.ts b/web/packages/new/photos/services/ml/blob.ts index 015dc7462d..d52772b6a4 100644 --- a/web/packages/new/photos/services/ml/blob.ts +++ b/web/packages/new/photos/services/ml/blob.ts @@ -1,4 +1,5 @@ import { basename } from "@/base/file"; +import type { ElectronMLWorker } from "@/base/types/ipc"; import { FileType } from "@/media/file-type"; import { decodeLivePhoto } from "@/media/live-photo"; import { ensure } from "@/utils/ensure"; @@ -7,7 +8,6 @@ import { renderableImageBlob } from "../../utils/file"; import { readStream } from "../../utils/native-stream"; import DownloadManager from "../download"; import type { UploadItem } from "../upload/types"; -import type { MLWorkerElectron } from "./worker-types"; /** * A pair of blobs - the original, and a possibly converted "renderable" one - @@ -103,13 +103,14 @@ export const imageBitmapAndData = async ( * be set to the {@link UploadItem} that was uploaded. This way, we can directly * use the on-disk file instead of needing to download the original from remote. * - * @param electron The {@link MLWorkerElectron} instance that allows us to call - * our Node.js layer for various functionality. + * @param electron The {@link ElectronMLWorker} instance that stands as a + * witness that we're actually running in our desktop app (and thus can safely + * call our Node.js layer for various functionality). */ export const indexableBlobs = async ( enteFile: EnteFile, uploadItem: UploadItem | undefined, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => uploadItem ? await indexableUploadItemBlobs(enteFile, uploadItem, electron) @@ -118,7 +119,7 @@ export const indexableBlobs = async ( const indexableUploadItemBlobs = async ( enteFile: EnteFile, uploadItem: UploadItem, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ) => { const fileType = enteFile.metadata.fileType; let originalImageBlob: Blob | undefined; @@ -149,7 +150,7 @@ const indexableUploadItemBlobs = async ( */ const readNonVideoUploadItem = async ( uploadItem: UploadItem, - electron: MLWorkerElectron, + electron: ElectronMLWorker, ): Promise => { if (typeof uploadItem == "string" || Array.isArray(uploadItem)) { const { response, lastModifiedMs } = await readStream( diff --git a/web/packages/new/photos/utils/native-stream.ts b/web/packages/new/photos/utils/native-stream.ts index 6f61656597..6aee016c0c 100644 --- a/web/packages/new/photos/utils/native-stream.ts +++ b/web/packages/new/photos/utils/native-stream.ts @@ -6,8 +6,7 @@ * See: [Note: IPC streams]. */ -import type { Electron, ZipItem } from "@/base/types/ipc"; -import type { MLWorkerElectron } from "../services/ml/worker-types"; +import type { Electron, ElectronMLWorker, ZipItem } from "@/base/types/ipc"; /** * Stream the given file or zip entry from the user's local file system. @@ -18,7 +17,7 @@ import type { MLWorkerElectron } from "../services/ml/worker-types"; * * To avoid accidentally invoking it in a non-desktop app context, it requires * the {@link Electron} (or a functionally similar) object as a parameter (even - * though it doesn't use it). + * though it doesn't need or use it). * * @param pathOrZipItem Either the path on the file on the user's local file * system whose contents we want to stream. Or a tuple containing the path to a @@ -36,7 +35,7 @@ import type { MLWorkerElectron } from "../services/ml/worker-types"; * reading, expressed as epoch milliseconds. */ export const readStream = async ( - _: Electron | MLWorkerElectron, + _: Electron | ElectronMLWorker, pathOrZipItem: string | ZipItem, ): Promise<{ response: Response; size: number; lastModifiedMs: number }> => { let url: URL; From 192e491acbe8ca81cf3fdefd2aaf2bbdf2556b8b Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 12:46:02 +0530 Subject: [PATCH 38/48] Match the documented behaviour --- web/packages/new/photos/services/ml/worker.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index cbeba5f844..5a3719d983 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -292,14 +292,14 @@ const indexNextBatch = async ( for (const item of items) { try { await index(item, electron); - delegate?.workerDidProcessFile(); - // Let us drain the microtask queue. This also gives a chance for other - // interactive tasks like `clipMatches` to run. - await wait(0); } catch (e) { log.warn(`Skipping unindexable file ${item.enteFile.id}`, e); allSuccess = false; } + delegate?.workerDidProcessFile(); + // Let us drain the microtask queue. This also gives a chance for other + // interactive tasks like `clipMatches` to run. + await wait(0); } // Return true if nothing failed. From 4647f9fac21847e7e9c95e94abce1232ec62c460 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 12:47:52 +0530 Subject: [PATCH 39/48] Undup and scope --- web/packages/new/photos/services/ml/worker.ts | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 5a3719d983..8ccd190cbb 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -292,8 +292,7 @@ const indexNextBatch = async ( for (const item of items) { try { await index(item, electron); - } catch (e) { - log.warn(`Skipping unindexable file ${item.enteFile.id}`, e); + } catch { allSuccess = false; } delegate?.workerDidProcessFile(); @@ -493,8 +492,12 @@ const index = async ( throw e; } - if (originalImageBlob && exif) - await cmpNewLib2(enteFile, originalImageBlob, exif); + try { + if (originalImageBlob && exif) + await cmpNewLib2(enteFile, originalImageBlob, exif); + } catch (e) { + log.warn(`Skipping exif cmp for ${f}`, e); + } log.debug(() => { const ms = Date.now() - startTime; From ebbb9a61ee1e2bc03146416a8c260d0043a319ed Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 12:55:48 +0530 Subject: [PATCH 40/48] Don't fail on exif errors --- web/packages/new/photos/services/ml/worker.ts | 43 ++++++++++++++++--- 1 file changed, 36 insertions(+), 7 deletions(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 8ccd190cbb..5de8c08390 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -11,7 +11,7 @@ import { wait } from "@/utils/promise"; import { DOMParser } from "@xmldom/xmldom"; import { expose, wrap } from "comlink"; import downloadManager from "../download"; -import { cmpNewLib2, extractRawExif } from "../exif"; +import { cmpNewLib2, extractRawExif, type RawExifTags } from "../exif"; import { getAllLocalFiles, getLocalTrashedFiles } from "../files"; import type { UploadItem } from "../upload/types"; import { @@ -35,6 +35,7 @@ import { import { fetchDerivedData, putDerivedData, + type RawRemoteDerivedData, type RemoteDerivedData, } from "./embedding"; import { faceIndexingVersion, indexFaces, type FaceIndex } from "./face"; @@ -480,10 +481,7 @@ const index = async ( [faceIndex, clipIndex, exif] = await Promise.all([ existingFaceIndex ?? indexFaces(enteFile, image, electron), existingCLIPIndex ?? indexCLIP(image, electron), - existingExif ?? - (originalImageBlob - ? extractRawExif(originalImageBlob) - : undefined), + existingExif ?? tryExtractExif(originalImageBlob, f), ]); } catch (e) { // See: [Note: Transient and permanent indexing failures] @@ -525,11 +523,11 @@ const index = async ( // parts. See: [Note: Preserve unknown derived data fields]. const existingRawDerivedData = remoteDerivedData?.raw ?? {}; - const rawDerivedData = { + const rawDerivedData: RawRemoteDerivedData = { ...existingRawDerivedData, face: remoteFaceIndex, clip: remoteCLIPIndex, - exif, + ...(exif ? { exif } : {}), }; log.debug(() => ["Uploading derived data", rawDerivedData]); @@ -571,3 +569,34 @@ const index = async ( image.bitmap.close(); } }; + +/** + * A helper function that tries to extract the raw Exif, but returns `undefined` + * if something goes wrong (or it isn't possible) instead of throwing. + * + * Exif extraction is not a critical item, we don't want the actual indexing to + * fail because we were unable to extract Exif. This is not rare: one scenario + * is if we were trying to index a file in an exotic format. The ML indexing + * will succeed (because we convert it to a renderable blob), but the Exif + * extraction will fail (since it needs the original blob, but the original blob + * can be an arbitrary format). + * + * @param originalImageBlob A {@link Blob} containing the original data for the + * image (or the image component of a live photo) whose Exif we're trying to + * extract. If this is not available, we skip the extraction and return + * `undefined`. + * + * @param f The {@link fileLogID} for the file this blob corresponds to. + */ +export const tryExtractExif = async ( + originalImageBlob: Blob | undefined, + f: string, +): Promise => { + if (!originalImageBlob) return undefined; + try { + return await extractRawExif(originalImageBlob); + } catch (e) { + log.warn(`Ignoring error during Exif extraction for ${f}`, e); + return undefined; + } +}; From ef3231380786b172f14f9a3275087938775ea2ee Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 14:04:29 +0530 Subject: [PATCH 41/48] x4 --- web/apps/cast/src/services/pair.ts | 3 +- web/packages/new/photos/services/ml/worker.ts | 31 +++++++++++++++---- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/web/apps/cast/src/services/pair.ts b/web/apps/cast/src/services/pair.ts index b5646698cc..287122c456 100644 --- a/web/apps/cast/src/services/pair.ts +++ b/web/apps/cast/src/services/pair.ts @@ -82,7 +82,8 @@ export const register = async (): Promise => { // Register keypair with museum to get a pairing code. let pairingCode: string | undefined; - // TODO: eslint has fixed this spurious warning, but we're not on the latest + // [TODO: spurious while(true) eslint warning]. + // eslint has fixed this spurious warning, but we're not on the latest // version yet, so add a disable. // https://github.com/eslint/eslint/pull/18286 /* eslint-disable no-constant-condition */ diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 5de8c08390..ca5f112468 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -288,15 +288,34 @@ const indexNextBatch = async ( // Nothing to do. if (items.length == 0) return false; - // Index, keeping track if any of the items failed. + // Keep track if any of the items failed. let allSuccess = true; - for (const item of items) { - try { - await index(item, electron); - } catch { - allSuccess = false; + + // Index up to 4 items simultaneously. + const tasks = new Array | undefined>(4).fill(undefined); + + let i = 0; + while (i < items.length) { + for (let j = 0; j < tasks.length; j++) { + if (!tasks[j]) { + tasks[j] = index(ensure(items[i++]), electron) + .then(() => { + tasks[j] = undefined; + }) + .catch(() => { + allSuccess = false; + tasks[j] = undefined; + }); + } } + + // Wait for at least one to complete (the other runners continue running + // even if one promise reaches the finish line). + await Promise.race(tasks); + + // Let the main thread now we're doing something. delegate?.workerDidProcessFile(); + // Let us drain the microtask queue. This also gives a chance for other // interactive tasks like `clipMatches` to run. await wait(0); From bf6aa5f8406b976114afd54ed8cc1df74510dab2 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 14:08:16 +0530 Subject: [PATCH 42/48] Fix --- web/packages/new/photos/services/ml/worker.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index ca5f112468..12eeaa648a 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -297,7 +297,7 @@ const indexNextBatch = async ( let i = 0; while (i < items.length) { for (let j = 0; j < tasks.length; j++) { - if (!tasks[j]) { + if (i < items.length && !tasks[j]) { tasks[j] = index(ensure(items[i++]), electron) .then(() => { tasks[j] = undefined; From 1b0fe5fd4cca0636050fdadf1b53ecc4308a2483 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 14:19:25 +0530 Subject: [PATCH 43/48] Tighten timings --- desktop/src/main/services/ml-worker.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index 48f819c8bf..bc788a9fad 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -253,10 +253,10 @@ const cachedCLIPImageSession = makeCachedInferenceSession( */ export const computeCLIPImageEmbedding = async (input: Float32Array) => { const session = await cachedCLIPImageSession(); - const t = Date.now(); const feeds = { input: new ort.Tensor("float32", input, [1, 3, 224, 224]), }; + const t = Date.now(); const results = await session.run(feeds); log.debugString(`ONNX/CLIP image embedding took ${Date.now() - t} ms`); /* Need these model specific casts to type the result */ @@ -296,13 +296,13 @@ export const computeCLIPTextEmbeddingIfAvailable = async (text: string) => { } const session = sessionOrSkip; - const t = Date.now(); const tokenizer = getTokenizer(); const tokenizedText = Int32Array.from(tokenizer.encodeForCLIP(text)); const feeds = { input: new ort.Tensor("int32", tokenizedText, [1, 77]), }; + const t = Date.now(); const results = await session.run(feeds); log.debugString(`ONNX/CLIP text embedding took ${Date.now() - t} ms`); return ensure(results.output).data as Float32Array; @@ -318,10 +318,10 @@ const cachedFaceDetectionSession = makeCachedInferenceSession( */ export const detectFaces = async (input: Float32Array) => { const session = await cachedFaceDetectionSession(); - const t = Date.now(); const feeds = { input: new ort.Tensor("float32", input, [1, 3, 640, 640]), }; + const t = Date.now(); const results = await session.run(feeds); log.debugString(`ONNX/YOLO face detection took ${Date.now() - t} ms`); return ensure(results.output).data; @@ -345,8 +345,8 @@ export const computeFaceEmbeddings = async (input: Float32Array) => { const inputTensor = new ort.Tensor("float32", input, [n, z, z, 3]); const session = await cachedFaceEmbeddingSession(); - const t = Date.now(); const feeds = { img_inputs: inputTensor }; + const t = Date.now(); const results = await session.run(feeds); log.debugString(`ONNX/MFNT face embedding took ${Date.now() - t} ms`); /* Need these model specific casts to extract and type the result */ From 59cc01053ae2e16514cbc4d733282ff12596e592 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 14:35:47 +0530 Subject: [PATCH 44/48] Handle refresh --- desktop/src/main/services/ml.ts | 13 ++++++++++++- web/packages/base/types/ipc.ts | 2 +- web/packages/new/photos/services/ml/index.ts | 19 ++++++++++++------- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/desktop/src/main/services/ml.ts b/desktop/src/main/services/ml.ts index 05642fc95f..cc1ae5764c 100644 --- a/desktop/src/main/services/ml.ts +++ b/desktop/src/main/services/ml.ts @@ -11,8 +11,11 @@ import { app, utilityProcess } from "electron/main"; import path from "node:path"; import log from "../log"; +/** The active ML worker (utility) process, if any. */ +let _child: UtilityProcess | undefined; + /** - * Create a new ML worker process. + * Create a new ML worker process, terminating the older ones (if any). * * [Note: ML IPC] * @@ -68,6 +71,12 @@ import log from "../log"; * to be relayed using `postMessage`. */ export const createMLWorker = (window: BrowserWindow) => { + if (_child) { + log.debug(() => "Terminating previous ML worker process"); + _child.kill(); + _child = undefined; + } + const { port1, port2 } = new MessageChannelMain(); const child = utilityProcess.fork(path.join(__dirname, "ml-worker.js")); @@ -77,6 +86,8 @@ export const createMLWorker = (window: BrowserWindow) => { window.webContents.postMessage("createMLWorker/port", undefined, [port2]); handleMessagesFromUtilityProcess(child); + + _child = child; }; /** diff --git a/web/packages/base/types/ipc.ts b/web/packages/base/types/ipc.ts index 748490ed58..c0644760c0 100644 --- a/web/packages/base/types/ipc.ts +++ b/web/packages/base/types/ipc.ts @@ -335,7 +335,7 @@ export interface Electron { // - ML /** - * Create a new ML worker. + * Create a new ML worker, terminating the older ones (if any). * * This creates a new Node.js utility process, and sets things up so that we * can communicate directly with that utility process using a diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index ba8083ad74..5b57dade21 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -98,22 +98,27 @@ export const terminateMLWorker = async () => { * Obtain a port from the Node.js layer that can be used to communicate with the * ML worker process. */ -const createMLWorker = async (electron: Electron): Promise => { - electron.createMLWorker(); - +const createMLWorker = (electron: Electron): Promise => { // The main process will do its thing, and send back the port it created to // us by sending an message on the "createMLWorker/port" channel via the // postMessage API. This roundabout way is needed because MessagePorts // cannot be transferred via the usual send/invoke pattern. - return new Promise((resolve) => { - window.onmessage = ({ source, data, ports }: MessageEvent) => { - // The source check verifies that the message is coming from the + const port = new Promise((resolve) => { + const l = ({ source, data, ports }: MessageEvent) => { + // The source check verifies that the message is coming from our own // preload script. The data is the message that was posted. - if (source == window && data == "createMLWorker/port") + if (source == window && data == "createMLWorker/port") { + window.removeEventListener("message", l); resolve(ensure(ports[0])); + } }; + window.addEventListener("message", l); }); + + electron.createMLWorker(); + + return port; }; /** From 5e055b6039e27dd2b673b92a6905df334f115384 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 14:44:42 +0530 Subject: [PATCH 45/48] opt unnecessary uploads --- web/packages/new/photos/services/ml/worker.ts | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 12eeaa648a..37d6dc259e 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -549,15 +549,20 @@ const index = async ( ...(exif ? { exif } : {}), }; - log.debug(() => ["Uploading derived data", rawDerivedData]); + if (existingFaceIndex && existingCLIPIndex && !exif) { + // If we were indexing just for exif, but exif generation didn't + // happen, there is no need to upload. + } else { + log.debug(() => ["Uploading derived data", rawDerivedData]); - try { - await putDerivedData(enteFile, rawDerivedData); - } catch (e) { - // See: [Note: Transient and permanent indexing failures] - log.error(`Failed to put derived data for ${f}`, e); - if (isHTTP4xxError(e)) await markIndexingFailed(enteFile.id); - throw e; + try { + await putDerivedData(enteFile, rawDerivedData); + } catch (e) { + // See: [Note: Transient and permanent indexing failures] + log.error(`Failed to put derived data for ${f}`, e); + if (isHTTP4xxError(e)) await markIndexingFailed(enteFile.id); + throw e; + } } try { From 46cc696ccd957dc8178c3a9ac79ef48a1be84c65 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 14:52:08 +0530 Subject: [PATCH 46/48] Avoid jargon people might not understand --- .../photos/src/components/Sidebar/AdvancedSettings.tsx | 2 +- web/apps/photos/src/components/Sidebar/Preferences.tsx | 7 +------ web/packages/new/photos/components/MLSettings.tsx | 10 +++++----- web/packages/new/photos/components/MLSettingsBeta.tsx | 2 +- 4 files changed, 8 insertions(+), 13 deletions(-) diff --git a/web/apps/photos/src/components/Sidebar/AdvancedSettings.tsx b/web/apps/photos/src/components/Sidebar/AdvancedSettings.tsx index 5fdcd23852..8980c8ed5f 100644 --- a/web/apps/photos/src/components/Sidebar/AdvancedSettings.tsx +++ b/web/apps/photos/src/components/Sidebar/AdvancedSettings.tsx @@ -82,7 +82,7 @@ export default function AdvancedSettings({ open, onClose, onRootClose }) { } onClick={() => setOpenMLSettings(true)} - label={pt("ML search")} + label={pt("Face and magic search")} /> diff --git a/web/apps/photos/src/components/Sidebar/Preferences.tsx b/web/apps/photos/src/components/Sidebar/Preferences.tsx index 38cccf45de..000a1e44cc 100644 --- a/web/apps/photos/src/components/Sidebar/Preferences.tsx +++ b/web/apps/photos/src/components/Sidebar/Preferences.tsx @@ -85,14 +85,9 @@ export default function Preferences({ open, onClose, onRootClose }) { } onClick={() => setOpenMLSettings(true)} - label={pt("ML search")} + label={pt("Face and magic search")} /> - )} diff --git a/web/packages/new/photos/components/MLSettings.tsx b/web/packages/new/photos/components/MLSettings.tsx index 7e2d53f872..2e5a0c6def 100644 --- a/web/packages/new/photos/components/MLSettings.tsx +++ b/web/packages/new/photos/components/MLSettings.tsx @@ -124,7 +124,7 @@ export const MLSettings: React.FC = ({ {component} @@ -305,7 +305,7 @@ const ManageML: React.FC = ({ let status: string; switch (phase) { case "indexing": - status = pt("Indexing"); + status = pt("Running"); break; case "scheduled": status = pt("Scheduled"); @@ -319,9 +319,9 @@ const ManageML: React.FC = ({ const confirmDisableML = () => { setDialogBoxAttributesV2({ - title: pt("Disable ML search"), + title: pt("Disable face and magic search"), content: pt( - "Do you want to disable ML search on all your devices?", + "Do you want to disable face and magic search on all your devices?", ), close: { text: t("cancel") }, proceed: { @@ -356,7 +356,7 @@ const ManageML: React.FC = ({ justifyContent={"space-between"} > - {pt("Status")} + {pt("Indexing")} {status} diff --git a/web/packages/new/photos/components/MLSettingsBeta.tsx b/web/packages/new/photos/components/MLSettingsBeta.tsx index db1b83da11..2f9bae19f4 100644 --- a/web/packages/new/photos/components/MLSettingsBeta.tsx +++ b/web/packages/new/photos/components/MLSettingsBeta.tsx @@ -42,7 +42,7 @@ export const MLSettingsBeta: React.FC = ({ From 154fffd620226a144e364f8e4ccadbabf6437435 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 15:06:13 +0530 Subject: [PATCH 47/48] Clean unused --- desktop/src/main/ipc.ts | 1 - desktop/src/main/services/ml-worker.ts | 40 -------------------------- desktop/src/preload.ts | 18 +----------- 3 files changed, 1 insertion(+), 58 deletions(-) diff --git a/desktop/src/main/ipc.ts b/desktop/src/main/ipc.ts index 6b837d8b4d..6c4020d6ee 100644 --- a/desktop/src/main/ipc.ts +++ b/desktop/src/main/ipc.ts @@ -44,7 +44,6 @@ import { import { convertToJPEG, generateImageThumbnail } from "./services/image"; import { logout } from "./services/logout"; import { createMLWorker } from "./services/ml"; - import { encryptionKey, lastShownChangelogVersion, diff --git a/desktop/src/main/services/ml-worker.ts b/desktop/src/main/services/ml-worker.ts index bc788a9fad..f4b9221f64 100644 --- a/desktop/src/main/services/ml-worker.ts +++ b/desktop/src/main/services/ml-worker.ts @@ -95,46 +95,6 @@ const parseInitData = (data: unknown) => { } }; -/** - * Our hand-rolled RPC handler and router - the Node.js utility process end. - * - * Sibling of the electronMLWorker function (in `ml/worker.ts`) in the web code. - * - * [Note: Node.js ML worker RPC protocol] - * - * - Each RPC call (i.e. request message) has a "method" (string), "id" - * (number) and "p" (arbitrary param). - * - * - Each RPC result (i.e. response message) has an "id" (number) that is the - * same as the "id" for the request which it corresponds to. - * - * - If the RPC call was a success, then the response messege will have an - * "result" (arbitrary result) property. Otherwise it will have a "error" - * (string) property describing what went wrong. - */ -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; - try { - switch (m.method) { - case "foo": - // if (p && typeof p == "string") - // return { id, result: await foo(p) }; - break; - } - } catch (e) { - return { id, error: e instanceof Error ? e.message : String(e) }; - } - return { id, error: "Unknown message" }; - } - - // We don't even have an "id", so at least log it lest the renderer also - // ignore the "id"-less response. - log.info("Ignoring unknown message", m); - return { error: "Unknown message" }; -}; - /** * Return a function that can be used to trigger a download of the specified * model, and the creating of an ONNX inference session initialized using it. diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index 5bd2f28987..5b83e441ad 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -202,24 +202,12 @@ const createMLWorker = () => { ipcRenderer.send("createMLWorker"); ipcRenderer.on("createMLWorker/port", (event) => { void windowLoaded.then(() => { - // "*"" is the origin + // "*"" is the origin to send to. window.postMessage("createMLWorker/port", "*", event.ports); }); }); }; -const computeCLIPImageEmbedding = (input: Float32Array) => - ipcRenderer.invoke("computeCLIPImageEmbedding", input); - -const computeCLIPTextEmbeddingIfAvailable = (text: string) => - ipcRenderer.invoke("computeCLIPTextEmbeddingIfAvailable", text); - -const detectFaces = (input: Float32Array) => - ipcRenderer.invoke("detectFaces", input); - -const computeFaceEmbeddings = (input: Float32Array) => - ipcRenderer.invoke("computeFaceEmbeddings", input); - // - Watch const watchGet = () => ipcRenderer.invoke("watchGet"); @@ -385,10 +373,6 @@ contextBridge.exposeInMainWorld("electron", { // - ML createMLWorker, - computeCLIPImageEmbedding, - computeCLIPTextEmbeddingIfAvailable, - detectFaces, - computeFaceEmbeddings, // - Watch From 0bc360c55caa88cb153420e6ff7a413faba6c0d0 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Wed, 31 Jul 2024 15:21:52 +0530 Subject: [PATCH 48/48] Add link --- desktop/src/preload.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/desktop/src/preload.ts b/desktop/src/preload.ts index 5b83e441ad..3058a6376f 100644 --- a/desktop/src/preload.ts +++ b/desktop/src/preload.ts @@ -314,8 +314,11 @@ const clearPendingUploads = () => ipcRenderer.invoke("clearPendingUploads"); * operation when it happens across threads. * https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Transferable_objects * - * In our case though, we're not dealing with threads but separate processes. So - * the ArrayBuffer will be copied: + * In our case though, we're not dealing with threads but separate processes. + * Electron currently only supports transferring MessagePorts: + * https://github.com/electron/electron/issues/34905 + * + * So the ArrayBuffer will be copied: * * > "parameters, errors and return values are **copied** when they're sent over * > the bridge".