diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/components/DeleteCacheButton.tsx b/src/equicordplugins/voiceMessageTranscriber.desktop/components/DeleteCacheButton.tsx index f1ed5167d..7766878f9 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/components/DeleteCacheButton.tsx +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/components/DeleteCacheButton.tsx @@ -4,28 +4,18 @@ * SPDX-License-Identifier: GPL-3.0-or-later */ -import { DataStore } from "@api/index"; import { Button } from "@components/Button"; import { useEffect, useState } from "@webpack/common"; +import { clearCache, getCacheSize } from "../utils/cache"; + export function DeleteCacheButton() { const [size, setSize] = useState(0); - const [deleteKeys, setDeleteKeys] = useState([]); useEffect(() => { let unmounted = false; - DataStore.entries().then(entries => { - if (unmounted) return; - let totalSize = 0; - const keys: string[] = []; - for (const [key, value] of entries) { - if (typeof key === "string" && (key.startsWith("whisper-") || key.startsWith("onnx-"))) { - totalSize += (value as string).length; - keys.push(key); - } - } - setSize(totalSize); - setDeleteKeys(keys); + getCacheSize().then(total => { + if (!unmounted) setSize(total); }); return () => { unmounted = true; }; }, []); @@ -33,9 +23,7 @@ export function DeleteCacheButton() { return ; diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/components/TranscriptionModal.tsx b/src/equicordplugins/voiceMessageTranscriber.desktop/components/TranscriptionModal.tsx index 5e07ccc83..a2eecebb9 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/components/TranscriptionModal.tsx +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/components/TranscriptionModal.tsx @@ -63,14 +63,14 @@ export function TranscriptionModal(props: { modalProps: RenderModalProps, src: s } ); - const { quantized, selectedModel } = settings.store; - workerRef.current.run( - audioData, - selectedModel, + const { quantized, selectedModel, useGpu } = settings.store; + workerRef.current.run(audioData, { + model: selectedModel, quantized, - options.language === "auto" ? undefined : options.language, - options.task - ); + useGpu, + language: options.language === "auto" ? undefined : options.language, + task: options.task + }); } catch (err) { if (active) { setError(String(err)); diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/components/VoiceMessageButton.tsx b/src/equicordplugins/voiceMessageTranscriber.desktop/components/VoiceMessageButton.tsx index 8d161f6ad..959841746 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/components/VoiceMessageButton.tsx +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/components/VoiceMessageButton.tsx @@ -22,7 +22,7 @@ import { LanguageSelectionModal } from "./LanguageSelectionModal"; const ChannelListIcon = findComponentByCodeLazy("1-1-1ZM2 8a1"); export function VoiceMessageButton({ src }: { src: string; }) { - const { embed, maintainHorizontal, quantized, selectedModel } = settings.use(["embed", "maintainHorizontal", "quantized", "selectedModel"]); + const { embed, maintainHorizontal, quantized, selectedModel, useGpu } = settings.use(["embed", "maintainHorizontal", "quantized", "selectedModel", "useGpu"]); const [isOpen, setIsOpen] = useState(false); const [status, setStatus] = useState("idle"); const [result, setResult] = useState(null); @@ -98,13 +98,12 @@ export function VoiceMessageButton({ src }: { src: string; }) { } ); - workerRef.current.run( - audioData, - selectedModel, + workerRef.current.run(audioData, { + model: selectedModel, quantized, - undefined, - "transcribe" - ); + useGpu, + task: "transcribe" + }); } catch (err) { if (runId === activeRunId.current) { setError(String(err)); diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/index.tsx b/src/equicordplugins/voiceMessageTranscriber.desktop/index.tsx index 17e1486e8..abe141cb3 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/index.tsx +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/index.tsx @@ -12,7 +12,8 @@ import definePlugin from "@utils/types"; import { setManaBaseRadioGroup } from "./components/LanguageSelectionModal"; import { VoiceMessageButton } from "./components/VoiceMessageButton"; -import { settings } from "./settings"; +import { migrateLegacyModel, settings } from "./settings"; +import { deleteLegacyCache } from "./utils/cache"; const VoiceMessageButtonWrapped = ErrorBoundary.wrap(VoiceMessageButton, { noop: true }); @@ -42,6 +43,11 @@ export default definePlugin({ }, settings, + start() { + migrateLegacyModel(); + deleteLegacyCache(); + }, + button(src: string) { return ; }, diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/settings.tsx b/src/equicordplugins/voiceMessageTranscriber.desktop/settings.tsx index 7c77411d2..0dd902126 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/settings.tsx +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/settings.tsx @@ -10,21 +10,33 @@ import { OptionType } from "@utils/types"; import { DeleteCacheButton } from "./components/DeleteCacheButton"; import { cl } from "./utils/misc"; -const MODEL_SIZES: Record = { - "Xenova/whisper-tiny": { quantized: "~40 MB", full: "~150 MB" }, - "Xenova/whisper-base": { quantized: "~77 MB", full: "~290 MB" }, - "Xenova/whisper-small": { quantized: "~250 MB", full: "~960 MB" }, - "Xenova/whisper-medium": { quantized: "~765 MB", full: "~3.1 GB" }, +const MODEL_SIZES: Record = { + "onnx-community/whisper-tiny": { quantized: "~40 MB", full: "~150 MB", gpu: "~75 MB" }, + "onnx-community/whisper-base": { quantized: "~77 MB", full: "~290 MB", gpu: "~145 MB" }, + "onnx-community/whisper-small": { quantized: "~250 MB", full: "~970 MB", gpu: "~490 MB" }, + "onnx-community/whisper-large-v3-turbo": { quantized: "~1.1 GB", full: "~3.2 GB", gpu: "~1.6 GB" }, }; +const LEGACY_MODELS: Record = { + "Xenova/whisper-tiny": "onnx-community/whisper-tiny", + "Xenova/whisper-base": "onnx-community/whisper-base", + "Xenova/whisper-small": "onnx-community/whisper-small", + "Xenova/whisper-medium": "onnx-community/whisper-large-v3-turbo", +}; + +export function migrateLegacyModel() { + const migrated = LEGACY_MODELS[settings.store.selectedModel]; + if (migrated) settings.store.selectedModel = migrated; +} + function renderModelOption(option?: { label: string; value: string; }) { if (!option) return null; return ; } function ModelOption({ option }: { option: { label: string; value: string; }; }) { - const isQuantized = settings.use(["quantized"]).quantized ?? true; - const size = MODEL_SIZES[option.value]?.[isQuantized ? "quantized" : "full"]; + const { quantized, useGpu } = settings.use(["quantized", "useGpu"]); + const size = MODEL_SIZES[option.value]?.[useGpu ? "gpu" : quantized ? "quantized" : "full"]; return (
@@ -57,20 +69,20 @@ export const settings = definePluginSettings({ options: [ { label: "Tiny (Fastest, lowest accuracy)", - value: "Xenova/whisper-tiny", + value: "onnx-community/whisper-tiny", }, { label: "Base (Recommended)", - value: "Xenova/whisper-base", + value: "onnx-community/whisper-base", default: true }, { label: "Small", - value: "Xenova/whisper-small" + value: "onnx-community/whisper-small" }, { - label: "Medium (Slowest, best accuracy)", - value: "Xenova/whisper-medium" + label: "Large v3 Turbo (Best accuracy, GPU recommended)", + value: "onnx-community/whisper-large-v3-turbo" } ], componentProps: { @@ -79,9 +91,15 @@ export const settings = definePluginSettings({ }, restartNeeded: false }, + useGpu: { + type: OptionType.BOOLEAN, + description: "Run the model on your GPU (WebGPU) when supported. Much faster, and makes the larger models usable. Falls back to CPU if unavailable.", + default: true, + restartNeeded: false + }, quantized: { type: OptionType.BOOLEAN, - description: "Use quantized models (smaller size, slight quality loss).", + description: "Use quantized models when running on CPU (smaller size, slight quality loss).", default: true, restartNeeded: false }, diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/style.css b/src/equicordplugins/voiceMessageTranscriber.desktop/style.css index feb8ce510..bb10775bf 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/style.css +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/style.css @@ -153,6 +153,11 @@ padding: 8px 12px; max-height: 180px; overflow-y: auto; + scrollbar-width: none; +} + +.vc-transcription-embed-body::-webkit-scrollbar { + display: none; } .vc-transcription-embed-text { diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/utils/cache.ts b/src/equicordplugins/voiceMessageTranscriber.desktop/utils/cache.ts new file mode 100644 index 000000000..f96986b84 --- /dev/null +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/utils/cache.ts @@ -0,0 +1,37 @@ +/* + * Vencord, a Discord client mod + * Copyright (c) 2026 Vendicated and contributors + * SPDX-License-Identifier: GPL-3.0-or-later + */ + +import { DataStore } from "@api/index"; + +const SIZES_KEY = "__sizes"; +const LEGACY_PREFIX = "VoiceMessageTranscriber_"; + +let store: ReturnType | undefined; +const getStore = () => store ??= DataStore.createStore("VoiceMessageTranscriberData", "VoiceMessageTranscriberStore"); + +export function getCachedFile(url: string) { + return DataStore.get(url, getStore()); +} + +export async function cacheFile(url: string, data: ArrayBuffer) { + await DataStore.set(url, data, getStore()); + await DataStore.update>(SIZES_KEY, sizes => ({ ...sizes, [url]: data.byteLength }), getStore()); +} + +export async function getCacheSize() { + const sizes = await DataStore.get>(SIZES_KEY, getStore()) ?? {}; + return Object.values(sizes).reduce((total, size) => total + size, 0); +} + +export function clearCache() { + return DataStore.clear(getStore()); +} + +export async function deleteLegacyCache() { + const keys = await DataStore.keys(); + const legacy = keys.filter(key => typeof key === "string" && key.startsWith(LEGACY_PREFIX)); + if (legacy.length) await DataStore.delMany(legacy); +} diff --git a/src/equicordplugins/voiceMessageTranscriber.desktop/utils/worker.ts b/src/equicordplugins/voiceMessageTranscriber.desktop/utils/worker.ts index 968dd97d7..25c73d2ab 100644 --- a/src/equicordplugins/voiceMessageTranscriber.desktop/utils/worker.ts +++ b/src/equicordplugins/voiceMessageTranscriber.desktop/utils/worker.ts @@ -4,19 +4,22 @@ * SPDX-License-Identifier: GPL-3.0-or-later */ -import { DataStore } from "@api/index"; -import { lodash } from "@webpack/common"; +import { cacheFile, getCachedFile } from "./cache"; const workerCode = ` -import { pipeline, env } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.17.2'; +import { pipeline, env, Tensor, WhisperTextStreamer } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.3.0/dist/transformers.min.js'; + +const CHUNK_LENGTH_S = 30; +const STRIDE_LENGTH_S = 5; env.allowLocalModels = false; env.useBrowserCache = false; +env.useWasmCache = false; const pendingRequests = new Map(); self.addEventListener('message', (event) => { - const { type, id, response, error, headers } = event.data; + const { type, id, response, status, error, headers } = event.data; if (type === 'fetch_response') { const resolver = pendingRequests.get(id); @@ -25,10 +28,10 @@ self.addEventListener('message', (event) => { if (error) { resolver.reject(new Error(error)); } else { - const res = new Response(response, { + resolver.resolve(new Response(response ?? null, { + status: status ?? 200, headers: headers || { 'Content-Type': 'application/octet-stream' } - }); - resolver.resolve(res); + })); } } } else if (type === 'run') { @@ -36,20 +39,40 @@ self.addEventListener('message', (event) => { } }); -const originalFetch = globalThis.fetch; -globalThis.fetch = async (input, init) => { - const url = input.toString(); - if (url.includes('huggingface.co') || url.includes('cdn.jsdelivr.net')) { - const id = Math.random().toString(36).substring(7); - return new Promise((resolve, reject) => { - pendingRequests.set(id, { resolve, reject }); - self.postMessage({ type: 'fetch_request', url, id }); - }); +const originalFetch = globalThis.fetch.bind(globalThis); +function proxyFetch(input, init) { + const url = input instanceof Request ? input.url : input.toString(); + if (!url.startsWith('https://huggingface.co/') && !url.startsWith('https://cdn.jsdelivr.net/')) { + return originalFetch(input, init); } - return originalFetch(input, init); -}; + const id = Math.random().toString(36).substring(7); + return new Promise((resolve, reject) => { + pendingRequests.set(id, { resolve, reject }); + self.postMessage({ type: 'fetch_request', url, id }); + }); +} +env.fetch = proxyFetch; +globalThis.fetch = proxyFetch; -let transcriber = null; +async function supportsWebGpu() { + try { + const adapter = await navigator.gpu?.requestAdapter(); + return !!adapter?.features.has('shader-f16'); + } catch { + return false; + } +} + +async function loadTranscriber(model, quantized, useGpu) { + if (useGpu && await supportsWebGpu()) { + try { + return await pipeline('automatic-speech-recognition', model, { device: 'webgpu', dtype: 'fp16' }); + } catch (e) { + console.warn('[VoiceMessageTranscriber] WebGPU failed, falling back to CPU', e); + } + } + return pipeline('automatic-speech-recognition', model, { device: 'wasm', dtype: quantized ? 'q8' : 'fp32' }); +} async function compressionRatio(text) { const bytes = new TextEncoder().encode(text); @@ -58,70 +81,73 @@ async function compressionRatio(text) { return bytes.length / compressed.byteLength; } -async function runTranscription({ audio, model, quantized, language, task }) { +let transcriber = null; + +// transformers.js has no language detection and falls back to english, which translates everything else +async function detectLanguage(audio) { + const { input_features } = await transcriber.processor(audio.subarray(0, 16000 * CHUNK_LENGTH_S)); + const { decoder_start_token_id, lang_to_id } = transcriber.model.generation_config; + const decoder_input_ids = new Tensor('int64', [BigInt(decoder_start_token_id)], [1, 1]); + let { logits } = await transcriber.model({ input_features, decoder_input_ids }); + if (logits.type !== 'float32') logits = logits.to('float32'); + + let best = null; + let bestScore = -Infinity; + for (const [token, id] of Object.entries(lang_to_id)) { + if (logits.data[id] > bestScore) { + bestScore = logits.data[id]; + best = token.slice(2, -2); + } + } + return best; +} + +async function runTranscription({ audio, model, quantized, useGpu, language, task }) { try { if (!transcriber) { self.postMessage({ type: 'status', status: 'loading' }); - transcriber = await pipeline('automatic-speech-recognition', model, { - quantized: quantized, - progress_callback: (data) => { - self.postMessage({ type: 'progress', data }); - } - }); + transcriber = await loadTranscriber(model, quantized, useGpu); } self.postMessage({ type: 'status', status: 'transcribing' }); + language ??= await detectLanguage(audio); + const time_precision = transcriber.processor.feature_extractor.config.chunk_length / transcriber.model.config.max_source_positions; - - let chunks_to_process; - - function chunk_callback(chunk) { - let last = chunks_to_process[chunks_to_process.length - 1]; - - Object.assign(last, chunk); - last.finalised = true; - - if (!chunk.is_last) { - chunks_to_process.push({ - tokens: [], - finalised: false, - }); - } - } - - function callback_function(item) { - let last = chunks_to_process[chunks_to_process.length - 1]; - - last.tokens = [...item[0].output_token_ids]; - - let data = transcriber.tokenizer._decode_asr(chunks_to_process, { - time_precision: time_precision, - return_timestamps: true, - force_full_sequences: false, - }); - - self.postMessage({ - type: 'partial', - output: { - text: data[0], - chunks: data[1].chunks - } - }); - } + const windowStep = CHUNK_LENGTH_S - 2 * STRIDE_LENGTH_S; async function transcribe(extraOptions) { - chunks_to_process = [{ tokens: [], finalised: false }]; + let text = ''; + const chunks = []; + let windowIndex = 0; + + const streamer = new WhisperTextStreamer(transcriber.tokenizer, { + time_precision, + skip_prompt: true, + on_chunk_start: time => { + chunks.push({ timestamp: [windowIndex * windowStep + time, null], text: '' }); + }, + callback_function: piece => { + text += piece; + if (chunks.length) chunks[chunks.length - 1].text += piece; + self.postMessage({ type: 'partial', output: { text, chunks } }); + }, + on_chunk_end: time => { + if (chunks.length) chunks[chunks.length - 1].timestamp[1] = windowIndex * windowStep + time; + }, + on_finalize: () => { + windowIndex++; + } + }); + return transcriber(audio, { - top_k: 0, do_sample: false, - chunk_length_s: 30, - stride_length_s: 5, + chunk_length_s: CHUNK_LENGTH_S, + stride_length_s: STRIDE_LENGTH_S, return_timestamps: true, - callback_function, - chunk_callback, + streamer, language, task: task === "translate" ? "translate" : undefined, ...extraOptions @@ -190,35 +216,27 @@ export class TranscriptionWorker { switch (type) { case "fetch_request": try { - const cachedData = await DataStore.get(`VoiceMessageTranscriber_${url}`); - - if (cachedData && lodash.isArrayBuffer(cachedData)) { - this.worker.postMessage({ - type: "fetch_response", - id, - response: cachedData, - headers: { - "Content-Length": cachedData.byteLength.toString(), - "Content-Type": this.getMimeType(url) - } - }); - } else { + let data = await getCachedFile(url); + if (!(data instanceof ArrayBuffer)) { const res = await fetch(url); - if (!res.ok) throw new Error("Failed to fetch " + url); + if (!res.ok) { + this.worker.postMessage({ type: "fetch_response", id, status: res.status }); + break; + } - const buffer = await res.arrayBuffer(); - await DataStore.set(`VoiceMessageTranscriber_${url}`, buffer); - - this.worker.postMessage({ - type: "fetch_response", - id, - response: buffer, - headers: { - "Content-Length": res.headers.get("Content-Length") || buffer.byteLength.toString(), - "Content-Type": this.getMimeType(url) - } - }); + 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", @@ -242,15 +260,8 @@ export class TranscriptionWorker { } } - public run(audio: Float32Array, model: string, quantized: boolean = true, language?: string, task?: string) { - this.worker.postMessage({ - type: "run", - audio, - model, - quantized, - language, - task - }); + public run(audio: Float32Array, options: { model: string; quantized: boolean; useGpu: boolean; language?: string; task?: string; }) { + this.worker.postMessage({ type: "run", audio, ...options }); } public terminate() {