change how worker works

electron-builder.json5

@@ -11,64 +11,6 @@
     "dist",
     "dist-electron"
   ],
-  "extraResources": [
-    {
-      "from": "node_modules/@huggingface",
-      "to": "app.asar.unpacked/node_modules/@huggingface"
-    },
-    {
-      "from": "node_modules/@xenova/transformers",
-      "to": "app.asar.unpacked/node_modules/@xenova/transformers"
-    },
-    {
-      "from": "node_modules/color",
-      "to": "app.asar.unpacked/node_modules/color"
-    },
-    {
-      "from": "node_modules/color-convert",
-      "to": "app.asar.unpacked/node_modules/color-convert"
-    },
-    {
-      "from": "node_modules/color-name",
-      "to": "app.asar.unpacked/node_modules/color-name"
-    },
-    {
-      "from": "node_modules/color-string",
-      "to": "app.asar.unpacked/node_modules/color-string"
-    },
-    {
-      "from": "node_modules/detect-libc",
-      "to": "app.asar.unpacked/node_modules/detect-libc"
-    },
-    {
-      "from": "node_modules/is-arrayish",
-      "to": "app.asar.unpacked/node_modules/is-arrayish"
-    },
-    {
-      "from": "node_modules/onnxruntime-common",
-      "to": "app.asar.unpacked/node_modules/onnxruntime-common"
-    },
-    {
-      "from": "node_modules/onnxruntime-node",
-      "to": "app.asar.unpacked/node_modules/onnxruntime-node"
-    },
-    {
-      "from": "node_modules/onnxruntime-web",
-      "to": "app.asar.unpacked/node_modules/onnxruntime-web"
-    },
-    {
-      "from": "node_modules/sharp",
-      "to": "app.asar.unpacked/node_modules/sharp"
-    },
-    {
-      "from": "node_modules/simple-swizzle",
-      "to": "app.asar.unpacked/node_modules/simple-swizzle"
-    },
-    {
-      "from": "src/worker.mjs",
-      "to": "app.asar.unpacked/src/worker.mjs"
-    }
-  ],
   "mac": {
     "target": [
       "dmg"

electron/main/index.ts

@@ -2,7 +2,6 @@ import path from 'node:path'
 import { createRequire } from 'node:module'
 import { fileURLToPath } from 'node:url'
 import os from 'node:os'
-import { Worker } from 'node:worker_threads'
 import { BrowserWindow, app, ipcMain, shell } from 'electron'
 
 import { WebSocketServer } from 'ws'
@@ -11,6 +10,8 @@ import { emit_osc, empty_queue } from './modules/osc'
 import { initialize_ws } from './modules/ws'
 import { check_update } from './modules/check_update'
 
+import { TranslationPipeline } from './modules/transformers'
+
 const store = new Store()
 
 const require = createRequire(import.meta.url)
@@ -109,11 +110,11 @@ async function createWindow() {
   })
 }
 
-const transformersWorkerPath = `file://${path.join(process.env.APP_ROOT, 'src', 'worker.mjs').replace('app.asar', 'app.asar.unpacked')}`
-const transformersWorker = new Worker(new URL(transformersWorkerPath, import.meta.url))
+// const transformersWorkerPath = `file://${path.join(process.env.APP_ROOT, 'src', 'worker.mjs').replace('app.asar', 'app.asar.unpacked')}`
+// const transformersWorker = new Worker(new URL(transformersWorkerPath, import.meta.url))
 
-const transformersPath = `file://${path.join(process.env.APP_ROOT, 'node_modules', '@xenova', 'transformers', 'src', 'transformers.js').replace('app.asar', 'app.asar.unpacked')}`
-transformersWorker.postMessage({ type: 'transformers-init', data: transformersPath })
+// const transformersPath = `file://${path.join(process.env.APP_ROOT, 'node_modules', '@xenova', 'transformers', 'src', 'transformers.js').replace('app.asar', 'app.asar.unpacked')}`
+// transformersWorker.postMessage({ type: 'transformers-init', data: transformersPath })
 
 app.whenReady().then(createWindow)
 
@@ -225,11 +226,7 @@ ipcMain.on('update-check', async () => {
 // → Footer ('transformers-translate-render')
 
 ipcMain.on('transformers-translate', async (event, args) => {
-  transformersWorker.postMessage({ type: 'transformers-translate', data: args })
-})
-
-transformersWorker.on('message', async (message) => {
-  if (message.type === 'transformers-translate-output') {
-    win.webContents.send('transformers-translate-render', message.data)
-  }
-})
+  // transformersWorker.postMessage({ type: 'transformers-translate', data: args })
+  const translator = new TranslationPipeline()
+  translator.translate(win, args)
+})

electron/main/modules/transformers.ts

@@ -0,0 +1,73 @@
+// see: https://github.com/xenova/transformers.js
+import { PipelineType, pipeline } from '@xenova/transformers'
+
+export class TranslationPipeline {
+  static task: PipelineType = 'translation'
+  static model = 'Xenova/nllb-200-distilled-600M'
+  static instance: any = null
+
+  static async getInstance(progress_callback: any = null) {
+    if (this.instance === null)
+      this.instance = pipeline(this.task, this.model, { progress_callback })
+
+    return this.instance
+  }
+
+  async translate(win: any, data: any) {
+    // call translator. downloads and caches model if first load
+    const translator = await TranslationPipeline.getInstance((x: any) => {
+        // self.postMessage(x)
+        win.webContents.send('transformers-translate-render', x)
+    })
+
+    // console.log(data)
+    const output = await translator(data.text, {
+      tgt_lang: data.tgt_lang,
+      src_lang: data.src_lang,
+
+      // partial outputs
+      callback_function: (x: any) => {
+        win.webContents.send('transformers-translate-render', {
+          status: 'update',
+          output: translator.tokenizer.decode(x[0].output_token_ids, { skip_special_tokens: true }),
+          index: data.index,
+        })
+      },
+    })
+
+    // send back to main thread
+    win.webContents.send('transformers-translate-render', {
+      status: 'complete',
+      output,
+      index: data.index,
+    })
+  }
+}
+
+// self.addEventListener('message', async (event) => {
+//   // call translator. downloads and caches model if first load
+//   const translator = await TranslationPipeline.getInstance((x: any) => {
+//     self.postMessage(x)
+//   })
+
+//   const output = await translator(event.data.text, {
+//     tgt_lang: event.data.tgt_lang,
+//     src_lang: event.data.src_lang,
+
+//     // partial outputs
+//     callback_function: (x: any) => {
+//       self.postMessage({
+//         status: 'update',
+//         output: translator.tokenizer.decode(x[0].output_token_ids, { skip_special_tokens: true }),
+//         index: event.data.index,
+//       })
+//     },
+//   })
+
+//   // send back to main thread
+//   self.postMessage({
+//     status: 'complete',
+//     output,
+//     index: event.data.index,
+//   })
+// })