Merge pull request #243 from ente-io/clip-desktop

Clip desktop

.gitignore

@@ -9,3 +9,4 @@ buildingSteps.md
 build/.DS_Store
 .env
 .electron-symbols/
+models/

package.json

@@ -1,7 +1,7 @@
 {
     "name": "ente",
     "productName": "ente",
-    "version": "1.6.49",
+    "version": "1.6.50-alpha.18",
     "private": true,
     "description": "Desktop client for ente.io",
     "main": "app/main.js",
@@ -30,7 +30,8 @@
                 ]
             },
             "category": "public.app-category.photography",
-            "hardenedRuntime": true
+            "hardenedRuntime": true,
+            "x64ArchFiles": "Contents/Resources/ggmlclip-mac"
         },
         "afterSign": "electron-builder-notarize",
         "extraFiles": [
@@ -87,7 +88,7 @@
         "concurrently": "^7.0.0",
         "cross-env": "^7.0.3",
         "electron": "^25.8.4",
-        "electron-builder": "^23.0.3",
+        "electron-builder": "^24.6.4",
         "electron-builder-notarize": "^1.2.0",
         "electron-download": "^4.1.1",
         "eslint": "^7.23.0",

src/api/clip.ts

@@ -0,0 +1,29 @@
+import { ipcRenderer } from 'electron';
+import { writeStream } from '../services/fs';
+
+export async function computeImageEmbedding(
+    imageData: Uint8Array
+): Promise<Float32Array> {
+    let tempInputFilePath = null;
+    try {
+        tempInputFilePath = await ipcRenderer.invoke('get-temp-file-path', '');
+        const imageStream = new Response(imageData.buffer).body;
+        await writeStream(tempInputFilePath, imageStream);
+        const embedding = await ipcRenderer.invoke(
+            'compute-image-embedding',
+            tempInputFilePath
+        );
+        return embedding;
+    } finally {
+        if (tempInputFilePath) {
+            await ipcRenderer.invoke('remove-temp-file', tempInputFilePath);
+        }
+    }
+}
+
+export async function computeTextEmbedding(
+    text: string
+): Promise<Float32Array> {
+    const embedding = await ipcRenderer.invoke('compute-text-embedding', text);
+    return embedding;
+}

src/api/common.ts

@@ -27,6 +27,15 @@ export const openDirectory = async (dirPath: string): Promise<void> => {
     }
 };
 
+export const getPlatform = async (): Promise<'mac' | 'windows' | 'linux'> => {
+    try {
+        return await ipcRenderer.invoke('get-platform');
+    } catch (e) {
+        logError(e, 'failed to get platform');
+        throw e;
+    }
+};
+
 export {
     logToDisk,
     openLogDirectory,

src/preload.ts

@@ -41,6 +41,7 @@ import {
     getAppVersion,
     openDirectory,
     updateOptOutOfCrashReports,
+    getPlatform,
 } from './api/common';
 import { fixHotReloadNext12 } from './utils/preload';
 import {
@@ -59,6 +60,7 @@ import {
     logRendererProcessMemoryUsage,
 } from './utils/processStats';
 import { runFFmpegCmd } from './api/ffmpeg';
+import { computeImageEmbedding, computeTextEmbedding } from './api/clip';
 
 fixHotReloadNext12();
 setupLogging();
@@ -114,4 +116,7 @@ windowObject['ElectronAPIs'] = {
     rename,
     deleteFile,
     updateOptOutOfCrashReports,
+    computeImageEmbedding,
+    computeTextEmbedding,
+    getPlatform,
 };

src/services/clipService.ts

@@ -0,0 +1,204 @@
+import * as log from 'electron-log';
+import util from 'util';
+import { logErrorSentry } from './sentry';
+import { isDev } from '../utils/common';
+import { app } from 'electron';
+import path from 'path';
+import { existsSync } from 'fs';
+import fs from 'fs/promises';
+const shellescape = require('any-shell-escape');
+const execAsync = util.promisify(require('child_process').exec);
+import fetch from 'node-fetch';
+import { writeNodeStream } from './fs';
+import { getPlatform } from '../utils/common/platform';
+
+const CLIP_MODEL_PATH_PLACEHOLDER = 'CLIP_MODEL';
+const GGMLCLIP_PATH_PLACEHOLDER = 'GGML_PATH';
+const INPUT_PATH_PLACEHOLDER = 'INPUT';
+
+const IMAGE_EMBEDDING_EXTRACT_CMD: string[] = [
+    GGMLCLIP_PATH_PLACEHOLDER,
+    '-mv',
+    CLIP_MODEL_PATH_PLACEHOLDER,
+    '--image',
+    INPUT_PATH_PLACEHOLDER,
+];
+
+const TEXT_EMBEDDING_EXTRACT_CMD: string[] = [
+    GGMLCLIP_PATH_PLACEHOLDER,
+    '-mt',
+    CLIP_MODEL_PATH_PLACEHOLDER,
+    '--text',
+    INPUT_PATH_PLACEHOLDER,
+];
+
+const TEXT_MODEL_DOWNLOAD_URL =
+    'https://models.ente.io/clip-vit-base-patch32_ggml-text-model-f16.gguf';
+const IMAGE_MODEL_DOWNLOAD_URL =
+    'https://models.ente.io/clip-vit-base-patch32_ggml-vision-model-f16.gguf';
+
+const TEXT_MODEL_NAME = 'clip-vit-base-patch32_ggml-text-model-f16.gguf';
+const IMAGE_MODEL_NAME = 'clip-vit-base-patch32_ggml-vision-model-f16.gguf';
+
+const IMAGE_MODEL_SIZE_IN_BYTES = 175957504; // 167.8 MB
+const TEXT_MODEL_SIZE_IN_BYTES = 127853440; // 121.9 MB
+const MODEL_SAVE_FOLDER = 'models';
+
+function getModelSavePath(modelName: string) {
+    let userDataDir: string;
+    if (isDev) {
+        userDataDir = '.';
+    } else {
+        userDataDir = app.getPath('userData');
+    }
+    return path.join(userDataDir, MODEL_SAVE_FOLDER, modelName);
+}
+
+async function downloadModel(saveLocation: string, url: string) {
+    // confirm that the save location exists
+    const saveDir = path.dirname(saveLocation);
+    if (!existsSync(saveDir)) {
+        log.info('creating model save dir');
+        await fs.mkdir(saveDir, { recursive: true });
+    }
+    log.info('downloading clip model');
+    const resp = await fetch(url);
+    await writeNodeStream(saveLocation, resp.body, true);
+    log.info('clip model downloaded');
+}
+
+let imageModelDownloadInProgress: Promise<void> = null;
+
+export async function getClipImageModelPath() {
+    const modelSavePath = getModelSavePath(IMAGE_MODEL_NAME);
+    if (imageModelDownloadInProgress) {
+        log.info('waiting for image model download to finish');
+        await imageModelDownloadInProgress;
+    } else {
+        if (!existsSync(modelSavePath)) {
+            log.info('clip image model not found, downloading');
+            imageModelDownloadInProgress = downloadModel(
+                modelSavePath,
+                IMAGE_MODEL_DOWNLOAD_URL
+            );
+            await imageModelDownloadInProgress;
+        } else {
+            const localFileSize = (await fs.stat(modelSavePath)).size;
+            if (localFileSize !== IMAGE_MODEL_SIZE_IN_BYTES) {
+                log.info('clip model size mismatch, downloading again');
+                imageModelDownloadInProgress = downloadModel(
+                    modelSavePath,
+                    IMAGE_MODEL_DOWNLOAD_URL
+                );
+                await imageModelDownloadInProgress;
+            }
+        }
+    }
+    return modelSavePath;
+}
+
+let textModelDownloadInProgress: Promise<void> = null;
+
+export async function getClipTextModelPath() {
+    const modelSavePath = getModelSavePath(TEXT_MODEL_NAME);
+    if (textModelDownloadInProgress) {
+        log.info('waiting for text model download to finish');
+        await textModelDownloadInProgress;
+    } else {
+        if (!existsSync(modelSavePath)) {
+            log.info('clip text model not found, downloading');
+            textModelDownloadInProgress = downloadModel(
+                modelSavePath,
+                TEXT_MODEL_DOWNLOAD_URL
+            );
+            await textModelDownloadInProgress;
+        } else {
+            const localFileSize = (await fs.stat(modelSavePath)).size;
+            if (localFileSize !== TEXT_MODEL_SIZE_IN_BYTES) {
+                log.info('clip model size mismatch, downloading again');
+                textModelDownloadInProgress = downloadModel(
+                    modelSavePath,
+                    TEXT_MODEL_DOWNLOAD_URL
+                );
+                await textModelDownloadInProgress;
+            }
+        }
+    }
+    return modelSavePath;
+}
+
+function getGGMLClipPath() {
+    return isDev
+        ? path.join('./build', `ggmlclip-${getPlatform()}`)
+        : path.join(process.resourcesPath, `ggmlclip-${getPlatform()}`);
+}
+
+export async function computeImageEmbedding(
+    inputFilePath: string
+): Promise<Float32Array> {
+    try {
+        const clipModelPath = await getClipImageModelPath();
+        const ggmlclipPath = getGGMLClipPath();
+        const cmd = IMAGE_EMBEDDING_EXTRACT_CMD.map((cmdPart) => {
+            if (cmdPart === GGMLCLIP_PATH_PLACEHOLDER) {
+                return ggmlclipPath;
+            } else if (cmdPart === CLIP_MODEL_PATH_PLACEHOLDER) {
+                return clipModelPath;
+            } else if (cmdPart === INPUT_PATH_PLACEHOLDER) {
+                return inputFilePath;
+            } else {
+                return cmdPart;
+            }
+        });
+
+        const escapedCmd = shellescape(cmd);
+        log.info('running clip command', escapedCmd);
+        const startTime = Date.now();
+        const { stdout } = await execAsync(escapedCmd);
+        log.info('clip command execution time ', Date.now() - startTime);
+        // parse stdout and return embedding
+        // get the last line of stdout
+        const lines = stdout.split('\n');
+        const lastLine = lines[lines.length - 1];
+        const embedding = JSON.parse(lastLine);
+        const embeddingArray = new Float32Array(embedding);
+        return embeddingArray;
+    } catch (err) {
+        logErrorSentry(err, 'Error in computeImageEmbedding');
+    }
+}
+
+export async function computeTextEmbedding(
+    text: string
+): Promise<Float32Array> {
+    try {
+        const clipModelPath = await getClipTextModelPath();
+        const ggmlclipPath = getGGMLClipPath();
+        const cmd = TEXT_EMBEDDING_EXTRACT_CMD.map((cmdPart) => {
+            if (cmdPart === GGMLCLIP_PATH_PLACEHOLDER) {
+                return ggmlclipPath;
+            } else if (cmdPart === CLIP_MODEL_PATH_PLACEHOLDER) {
+                return clipModelPath;
+            } else if (cmdPart === INPUT_PATH_PLACEHOLDER) {
+                return text;
+            } else {
+                return cmdPart;
+            }
+        });
+
+        const escapedCmd = shellescape(cmd);
+        log.info('running clip command', escapedCmd);
+        const startTime = Date.now();
+        const { stdout } = await execAsync(escapedCmd);
+        log.info('clip command execution time ', Date.now() - startTime);
+        // parse stdout and return embedding
+        // get the last line of stdout
+        const lines = stdout.split('\n');
+        const lastLine = lines[lines.length - 1];
+        const embedding = JSON.parse(lastLine);
+        const embeddingArray = new Float32Array(embedding);
+        return embeddingArray;
+    } catch (err) {
+        logErrorSentry(err, 'Error in computeTextEmbedding');
+    }
+}

src/services/fs.ts

@@ -6,6 +6,8 @@ import StreamZip from 'node-stream-zip';
 import { Readable } from 'stream';
 import { logError } from './logging';
 import { existsSync } from 'fs';
+import { log } from 'electron-log';
+import { convertBytesToHumanReadable } from '../utils/logging';
 
 // https://stackoverflow.com/a/63111390
 export const getDirFilePaths = async (dirPath: string) => {
@@ -226,18 +228,26 @@ export const convertBrowserStreamToNode = (
     return rs;
 };
 
-export async function writeStream(
+export async function writeNodeStream(
     filePath: string,
-    fileStream: ReadableStream<Uint8Array>
+    fileStream: NodeJS.ReadableStream,
+    enableLogging = false
 ) {
     const writeable = fs.createWriteStream(filePath);
-    const readable = convertBrowserStreamToNode(fileStream);
 
-    readable.on('error', (error) => {
+    fileStream.on('error', (error) => {
         writeable.destroy(error); // Close the writable stream with an error
     });
 
-    readable.pipe(writeable);
+    fileStream.pipe(writeable);
+
+    let downloaded = 0;
+    if (enableLogging) {
+        fileStream.on('data', (chunk) => {
+            downloaded += chunk.length;
+            log(`Received ${convertBytesToHumanReadable(downloaded)} of data.`);
+        });
+    }
 
     await new Promise((resolve, reject) => {
         writeable.on('finish', resolve);
@@ -250,6 +260,14 @@ export async function writeStream(
     });
 }
 
+export async function writeStream(
+    filePath: string,
+    fileStream: ReadableStream<Uint8Array>
+) {
+    const readable = convertBrowserStreamToNode(fileStream);
+    await writeNodeStream(filePath, readable);
+}
+
 export async function readTextFile(filePath: string) {
     if (!existsSync(filePath)) {
         throw new Error('File does not exist');

src/utils/common/platform.ts

@@ -1,11 +1,19 @@
 export function isPlatform(platform: 'mac' | 'windows' | 'linux') {
-    if (process.platform === 'darwin') {
-        return platform === 'mac';
-    } else if (process.platform === 'win32') {
-        return platform === 'windows';
-    } else if (process.platform === 'linux') {
-        return platform === 'linux';
-    } else {
-        return false;
+    return getPlatform() === platform;
+}
+
+export function getPlatform(): 'mac' | 'windows' | 'linux' {
+    switch (process.platform) {
+        case 'aix':
+        case 'freebsd':
+        case 'linux':
+        case 'openbsd':
+        case 'android':
+            return 'linux';
+        case 'darwin':
+        case 'sunos':
+            return 'mac';
+        case 'win32':
+            return 'windows';
     }
 }

src/utils/ipcComms.ts

@@ -27,6 +27,11 @@ import { deleteTempFile, runFFmpegCmd } from '../services/ffmpeg';
 import { generateTempFilePath } from './temp';
 import { setOptOutOfCrashReports } from '../services/userPreference';
 import { updateOptOutOfCrashReports } from '../main';
+import {
+    computeImageEmbedding,
+    computeTextEmbedding,
+} from '../services/clipService';
+import { getPlatform } from './common/platform';
 
 export default function setupIpcComs(
     tray: Tray,
@@ -166,4 +171,13 @@ export default function setupIpcComs(
         setOptOutOfCrashReports(optOut);
         updateOptOutOfCrashReports(optOut);
     });
+    ipcMain.handle('compute-image-embedding', (_, inputFilePath) => {
+        return computeImageEmbedding(inputFilePath);
+    });
+    ipcMain.handle('compute-text-embedding', (_, text) => {
+        return computeTextEmbedding(text);
+    });
+    ipcMain.handle('get-platform', () => {
+        return getPlatform();
+    });
 }

src/utils/logging.ts

@@ -23,3 +23,16 @@ export function makeID(length: number) {
     }
     return result;
 }
+
+export function convertBytesToHumanReadable(
+    bytes: number,
+    precision = 2
+): string {
+    if (bytes === 0 || isNaN(bytes)) {
+        return '0 MB';
+    }
+
+    const i = Math.floor(Math.log(bytes) / Math.log(1024));
+    const sizes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'];
+    return (bytes / Math.pow(1024, i)).toFixed(precision) + ' ' + sizes[i];
+}

src/utils/processStats.ts

@@ -1,5 +1,6 @@
 import ElectronLog from 'electron-log';
 import { webFrame } from 'electron/renderer';
+import { convertBytesToHumanReadable } from './logging';
 
 const LOGGING_INTERVAL_IN_MICROSECONDS = 30 * 1000; // 30 seconds
 
@@ -292,13 +293,3 @@ const getNormalizedWebFrameResourceUsage = () => {
         },
     };
 };
-
-function convertBytesToHumanReadable(bytes: number, precision = 2): string {
-    if (bytes === 0 || isNaN(bytes)) {
-        return '0 MB';
-    }
-
-    const i = Math.floor(Math.log(bytes) / Math.log(1024));
-    const sizes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'];
-    return (bytes / Math.pow(1024, i)).toFixed(precision) + ' ' + sizes[i];
-}

ui

@@ -1 +1 @@
-Subproject commit 45f83e212ae851b25cfad941d233ccc91a320881
+Subproject commit 1991d2825b85eb3ce9b31fc427b978458b55ee38