VoiceMessageTranscriber: Keep Modal Loaded Until Stop/Reload/Etc

src/equicordplugins/voiceMessageTranscriber.desktop/components/TranscriptionModal.tsx

@@ -14,7 +14,7 @@ import { Modal, useEffect, useRef, useState } from "@webpack/common";
 import { settings } from "../settings";
 import { decodeAudio, fetchVoiceMessage } from "../utils/audio";
 import { cl, formatTranscript } from "../utils/misc";
-import { TranscriptionResult, TranscriptionWorker } from "../utils/worker";
+import { TranscriptionJob, TranscriptionResult } from "../utils/worker";
 
 export function TranscriptionModal(props: { modalProps: RenderModalProps, src: string, options: { language: string, task: string; }; }) {
     const { modalProps, src, options } = props;
@@ -25,7 +25,7 @@ export function TranscriptionModal(props: { modalProps: RenderModalProps, src: s
     const [copied, setCopied] = useState(false);
     const [retryCount, setRetryCount] = useState(0);
 
-    const workerRef = useRef<TranscriptionWorker | null>(null);
+    const jobRef = useRef<TranscriptionJob | null>(null);
 
     useEffect(() => {
         let active = true;
@@ -41,30 +41,30 @@ export function TranscriptionModal(props: { modalProps: RenderModalProps, src: s
                 const audioData = await decodeAudio(blob);
 
                 if (!active) return;
-                workerRef.current?.terminate();
+                jobRef.current?.cancel();
-                workerRef.current = new TranscriptionWorker(
+                jobRef.current = new TranscriptionJob({
-                    s => {
+                    onStatus: s => {
                         if (active) setStatus(s);
                     },
-                    out => {
+                    onComplete: out => {
                         if (active) {
                             setResult(out);
                             setStatus("complete");
                         }
                     },
-                    err => {
+                    onError: err => {
                         if (active) {
                             setError(String(err));
                             setStatus("error");
                         }
                     },
-                    partial => {
+                    onPartial: partial => {
                         if (active) setResult(partial);
                     }
-                );
+                });
 
                 const { quantized, selectedModel, useGpu } = settings.store;
-                workerRef.current.run(audioData, {
+                jobRef.current.run(audioData, {
                     model: selectedModel,
                     quantized,
                     useGpu,
@@ -81,7 +81,7 @@ export function TranscriptionModal(props: { modalProps: RenderModalProps, src: s
 
         return () => {
             active = false;
-            workerRef.current?.terminate();
+            jobRef.current?.cancel();
         };
     }, [retryCount]);
 

src/equicordplugins/voiceMessageTranscriber.desktop/components/VoiceMessageButton.tsx

@@ -16,7 +16,7 @@ import { openModal, ScrollerAuto, Tooltip, useEffect, useRef, useState } from "@
 import { settings } from "../settings";
 import { decodeAudio, fetchVoiceMessage } from "../utils/audio";
 import { cl, formatTranscript } from "../utils/misc";
-import { TranscriptionResult, TranscriptionWorker } from "../utils/worker";
+import { TranscriptionJob, TranscriptionResult } from "../utils/worker";
 import { LanguageSelectionModal } from "./LanguageSelectionModal";
 
 const ChannelListIcon = findComponentByCodeLazy("1-1-1ZM2 8a1");
@@ -30,7 +30,7 @@ export function VoiceMessageButton({ src }: { src: string; }) {
     const [error, setError] = useState<string | null>(null);
     const [copied, setCopied] = useState(false);
 
-    const workerRef = useRef<TranscriptionWorker | null>(null);
+    const jobRef = useRef<TranscriptionJob | null>(null);
     const activeRunId = useRef(0);
     const buttonRef = useRef<HTMLButtonElement>(null);
 
@@ -53,15 +53,15 @@ export function VoiceMessageButton({ src }: { src: string; }) {
     useEffect(() => {
         return () => {
             activeRunId.current++;
-            workerRef.current?.terminate();
+            jobRef.current?.cancel();
-            workerRef.current = null;
+            jobRef.current = null;
         };
     }, []);
 
     const startTranscription = async () => {
         const runId = ++activeRunId.current;
-        workerRef.current?.terminate();
+        jobRef.current?.cancel();
-        workerRef.current = null;
+        jobRef.current = null;
 
         setIsOpen(true);
         setError(null);
@@ -77,28 +77,28 @@ export function VoiceMessageButton({ src }: { src: string; }) {
 
             if (runId !== activeRunId.current) return;
 
-            workerRef.current = new TranscriptionWorker(
+            jobRef.current = new TranscriptionJob({
-                s => {
+                onStatus: s => {
                     if (runId === activeRunId.current) setStatus(s);
                 },
-                out => {
+                onComplete: out => {
                     if (runId === activeRunId.current) {
                         setResult(out);
                         setStatus("complete");
                     }
                 },
-                err => {
+                onError: err => {
                     if (runId === activeRunId.current) {
                         setError(String(err));
                         setStatus("error");
                     }
                 },
-                partial => {
+                onPartial: partial => {
                     if (runId === activeRunId.current) setResult(partial);
                 }
-            );
+            });
 
-            workerRef.current.run(audioData, {
+            jobRef.current.run(audioData, {
                 model: selectedModel,
                 quantized,
                 useGpu,

src/equicordplugins/voiceMessageTranscriber.desktop/index.tsx

@@ -14,6 +14,7 @@ import { setManaBaseRadioGroup } from "./components/LanguageSelectionModal";
 import { VoiceMessageButton } from "./components/VoiceMessageButton";
 import { migrateLegacyModel, settings } from "./settings";
 import { deleteLegacyCache } from "./utils/cache";
+import { terminateWorker } from "./utils/worker";
 
 const VoiceMessageButtonWrapped = ErrorBoundary.wrap(VoiceMessageButton, { noop: true });
 
@@ -48,6 +49,10 @@ export default definePlugin({
         deleteLegacyCache();
     },
 
+    stop() {
+        terminateWorker();
+    },
+
     button(src: string) {
         return <VoiceMessageButtonWrapped src={src} />;
     },

src/equicordplugins/voiceMessageTranscriber.desktop/utils/worker.ts

@@ -17,6 +17,8 @@ env.useBrowserCache = false;
 env.useWasmCache = false;
 
 const pendingRequests = new Map();
+const cancelledRuns = new Set();
+let queue = Promise.resolve();
 
 self.addEventListener('message', (event) => {
     const { type, id, response, status, error, headers } = event.data;
@@ -35,7 +37,9 @@ self.addEventListener('message', (event) => {
             }
         }
     } else if (type === 'run') {
-        runTranscription(event.data);
+        queue = queue.then(() => runTranscription(event.data));
+    } else if (type === 'cancel') {
+        cancelledRuns.add(event.data.runId);
     }
 });
 
@@ -82,6 +86,9 @@ async function compressionRatio(text) {
 }
 
 let transcriber = null;
+let loadedKey = null;
+
+class Cancelled extends Error {}
 
 // transformers.js has no language detection and falls back to english, which translates everything else
 async function detectLanguage(audio) {
@@ -102,16 +109,32 @@ async function detectLanguage(audio) {
     return best;
 }
 
-async function runTranscription({ audio, model, quantized, useGpu, language, task }) {
+async function runTranscription({ runId, audio, model, quantized, useGpu, language, task }) {
+    const post = msg => self.postMessage({ ...msg, runId });
+    const checkCancelled = () => {
+        if (cancelledRuns.has(runId)) throw new Cancelled();
+    };
+
     try {
-        if (!transcriber) {
+        checkCancelled();
-            self.postMessage({ type: 'status', status: 'loading' });
+
+        // keep the model loaded between runs, only reload when the settings change
+        const key = [model, quantized, useGpu].join();
+        if (loadedKey !== key) {
+            post({ type: 'status', status: 'loading' });
+            const previous = transcriber;
+            transcriber = null;
+            loadedKey = null;
+            await previous?.dispose();
             transcriber = await loadTranscriber(model, quantized, useGpu);
+            loadedKey = key;
         }
 
-        self.postMessage({ type: 'status', status: 'transcribing' });
+        checkCancelled();
+        post({ type: 'status', status: 'transcribing' });
 
         language ??= await detectLanguage(audio);
+        checkCancelled();
 
         const time_precision =
             transcriber.processor.feature_extractor.config.chunk_length /
@@ -130,9 +153,10 @@ async function runTranscription({ audio, model, quantized, useGpu, language, tas
                     chunks.push({ timestamp: [windowIndex * windowStep + time, null], text: '' });
                 },
                 callback_function: piece => {
+                    checkCancelled();
                     text += piece;
                     if (chunks.length) chunks[chunks.length - 1].text += piece;
-                    self.postMessage({ type: 'partial', output: { text, chunks } });
+                    post({ type: 'partial', output: { text, chunks } });
                 },
                 on_chunk_end: time => {
                     if (chunks.length) chunks[chunks.length - 1].timestamp[1] = windowIndex * windowStep + time;
@@ -160,10 +184,12 @@ async function runTranscription({ audio, model, quantized, useGpu, language, tas
             output = await transcribe({ no_repeat_ngram_size: 4 });
         }
 
-        self.postMessage({ type: 'complete', output });
+        post({ type: 'complete', output });
 
     } catch (e) {
-        self.postMessage({ type: 'error', error: e.toString() });
+        if (!(e instanceof Cancelled)) post({ type: 'error', error: e.toString() });
+    } finally {
+        cancelledRuns.delete(runId);
     }
 }
 `;
@@ -178,93 +204,133 @@ export interface TranscriptionResult {
     chunks: TranscriptionChunk[];
 }
 
-export class TranscriptionWorker {
+export interface TranscriptionOptions {
-    private worker: Worker;
+    model: string;
-    private onStatus: (status: string) => void;
+    quantized: boolean;
-    private onComplete: (output: TranscriptionResult) => void;
+    useGpu: boolean;
-    private onError: (error: unknown) => void;
+    language?: string;
-    private onPartial: (output: TranscriptionResult) => void;
+    task?: string;
+}
 
-    constructor(
+interface TranscriptionCallbacks {
-        onStatus: (status: string) => void,
+    onStatus: (status: string) => void;
-        onComplete: (output: TranscriptionResult) => void,
+    onComplete: (output: TranscriptionResult) => void;
-        onError: (error: unknown) => void,
+    onError: (error: unknown) => void;
-        onPartial: (output: TranscriptionResult) => void
+    onPartial: (output: TranscriptionResult) => void;
-    ) {
+}
-        this.onStatus = onStatus;
-        this.onComplete = onComplete;
-        this.onError = onError;
-        this.onPartial = onPartial;
 
-        const blob = new Blob([workerCode], { type: "text/javascript" });
+let worker: Worker | null = null;
-        const objectUrl = URL.createObjectURL(blob);
+let nextRunId = 0;
-        this.worker = new Worker(objectUrl, { type: "module" });
+const runs = new Map<number, TranscriptionCallbacks>();
-        URL.revokeObjectURL(objectUrl);
-        this.worker.onmessage = this.handleMessage.bind(this);
-    }
 
-    private getMimeType(url: string): string {
+function getMimeType(url: string): string {
-        if (url.endsWith(".wasm")) return "application/wasm";
+    if (url.endsWith(".wasm")) return "application/wasm";
-        if (url.endsWith(".json")) return "application/json";
+    if (url.endsWith(".json")) return "application/json";
-        if (url.endsWith(".onnx")) return "application/octet-stream";
+    return "application/octet-stream";
-        return "application/octet-stream";
+}
-    }
 
-    private async handleMessage(event: MessageEvent) {
+async function handleFetchRequest(target: Worker, id: string, url: string) {
-        const { type, id, url, status, output, error } = event.data;
+    try {
+        let data = await getCachedFile(url);
+        if (!(data instanceof ArrayBuffer)) {
+            const res = await fetch(url);
+            if (!res.ok) {
+                target.postMessage({ type: "fetch_response", id, status: res.status });
+                return;
+            }
 
-        switch (type) {
+            data = await res.arrayBuffer();
-            case "fetch_request":
+            await cacheFile(url, data);
-                try {
-                    let data = await getCachedFile(url);
-                    if (!(data instanceof ArrayBuffer)) {
-                        const res = await fetch(url);
-                        if (!res.ok) {
-                            this.worker.postMessage({ type: "fetch_response", id, status: res.status });
-                            break;
-                        }
-
-                        data = await res.arrayBuffer();
-                        await cacheFile(url, data);
-                    }
-
-                    this.worker.postMessage({
-                        type: "fetch_response",
-                        id,
-                        response: data,
-                        headers: {
-                            "Content-Length": data.byteLength.toString(),
-                            "Content-Type": this.getMimeType(url)
-                        }
-                    }, [data]);
-                } catch (err) {
-                    this.worker.postMessage({
-                        type: "fetch_response",
-                        id,
-                        error: String(err)
-                    });
-                }
-                break;
-            case "status":
-                this.onStatus(status);
-                break;
-            case "complete":
-                this.onComplete(output);
-                break;
-            case "partial":
-                this.onPartial(output);
-                break;
-            case "error":
-                this.onError(error);
-                break;
         }
-    }
 
-    public run(audio: Float32Array, options: { model: string; quantized: boolean; useGpu: boolean; language?: string; task?: string; }) {
+        target.postMessage({
-        this.worker.postMessage({ type: "run", audio, ...options });
+            type: "fetch_response",
-    }
+            id,
-
+            response: data,
-    public terminate() {
+            headers: {
-        this.worker.terminate();
+                "Content-Length": data.byteLength.toString(),
+                "Content-Type": getMimeType(url)
+            }
+        }, [data]);
+    } catch (err) {
+        target.postMessage({
+            type: "fetch_response",
+            id,
+            error: String(err)
+        });
+    }
+}
+
+function handleMessage(target: Worker, event: MessageEvent) {
+    const { type, id, url, runId, status, output, error } = event.data;
+
+    if (type === "fetch_request") {
+        handleFetchRequest(target, id, url);
+        return;
+    }
+
+    const run = runs.get(runId);
+    if (!run) return;
+
+    switch (type) {
+        case "status":
+            run.onStatus(status);
+            break;
+        case "partial":
+            run.onPartial(output);
+            break;
+        case "complete":
+            runs.delete(runId);
+            run.onComplete(output);
+            break;
+        case "error":
+            runs.delete(runId);
+            run.onError(error);
+            break;
+    }
+}
+
+function getWorker() {
+    if (worker) return worker;
+
+    const blob = new Blob([workerCode], { type: "text/javascript" });
+    const objectUrl = URL.createObjectURL(blob);
+    const created = new Worker(objectUrl, { type: "module" });
+    URL.revokeObjectURL(objectUrl);
+
+    created.onmessage = event => handleMessage(created, event);
+    created.onerror = event => {
+        event.preventDefault();
+        terminateWorker(event.message || "Transcription worker crashed");
+    };
+
+    return worker = created;
+}
+
+export function terminateWorker(reason = "Transcription worker was stopped") {
+    worker?.terminate();
+    worker = null;
+
+    const pending = [...runs.values()];
+    runs.clear();
+    for (const run of pending) run.onError(reason);
+}
+
+export class TranscriptionJob {
+    private runId: number | null = null;
+
+    constructor(private callbacks: TranscriptionCallbacks) { }
+
+    public run(audio: Float32Array, options: TranscriptionOptions) {
+        this.cancel();
+        const runId = this.runId = nextRunId++;
+        runs.set(runId, this.callbacks);
+        getWorker().postMessage({ type: "run", runId, audio, ...options });
+    }
+
+    public cancel() {
+        if (this.runId === null) return;
+        if (runs.delete(this.runId)) worker?.postMessage({ type: "cancel", runId: this.runId });
+        this.runId = null;
     }
 }