From 76f51c5bd729a7fa8d023382cf631dbbf3796f08 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 31 Jul 2026 02:42:46 +0000 Subject: [PATCH 1/5] Add Parakeet TDT v3 as an optional transcription backend Offer NVIDIA Parakeet alongside Whisper Base/Small via parakeet.js, keeping VAD + pyannote diarization. WebGPU uses fp16 with WASM int8 fallback; same-origin ORT WASM is served from /vendor/ort-parakeet. Co-authored-by: Wassim Gharbi --- components/Editor.tsx | 1 + components/ImportTranscriptOption.tsx | 12 +- components/ModelSelector.tsx | 31 +- components/UploadScreen.tsx | 1 + hooks/useTranscriber.ts | 14 +- lib/models.ts | 39 +- lib/store.ts | 13 +- lib/types.ts | 6 +- next.config.ts | 2 + package-lock.json | 30 ++ package.json | 1 + scripts/copy-assets.mjs | 77 +++- tests/models-test.ts | 33 ++ workers/transcription.worker.ts | 603 ++++++++++++++++++-------- 14 files changed, 646 insertions(+), 217 deletions(-) create mode 100644 tests/models-test.ts diff --git a/components/Editor.tsx b/components/Editor.tsx index c277a0e..ba9fc57 100644 --- a/components/Editor.tsx +++ b/components/Editor.tsx @@ -232,6 +232,7 @@ export default function Editor() { + diff --git a/components/ImportTranscriptOption.tsx b/components/ImportTranscriptOption.tsx index d3ab301..da121b5 100644 --- a/components/ImportTranscriptOption.tsx +++ b/components/ImportTranscriptOption.tsx @@ -8,7 +8,7 @@ import { TRANSCRIPT_FILE_ERROR, TRANSCRIPT_ACCEPT, } from "@/lib/parseTranscript"; -import { isWhisperModel } from "@/lib/models"; +import { isAsrModel, type AsrModel } from "@/lib/models"; import { useEditorStore } from "@/lib/store"; import { ModelOption, @@ -40,7 +40,7 @@ export default function ImportTranscriptOption() { ModelOptionContextValue, "keepMenuOpen" | "closeMenu" | "select" > | null>(null); - const previousModelRef = useRef<"base" | "small">("base"); + const previousModelRef = useRef("base"); const pickGenRef = useRef(0); /** Reset import-pick state only — never touch the dropdown open state. */ @@ -53,11 +53,11 @@ export default function ImportTranscriptOption() { } }, [setModel]); - // If the user switches to Whisper while a picker/parse is in flight, invalidate - // so a late onChange/parse cannot flip model back to import. + // If the user switches to an ASR model while a picker/parse is in flight, + // invalidate so a late onChange/parse cannot flip model back to import. useEffect(() => { return useEditorStore.subscribe((state, prev) => { - if (!isWhisperModel(state.model) || state.model === prev.model) return; + if (!isAsrModel(state.model) || state.model === prev.model) return; pickGenRef.current += 1; queueMicrotask(() => { setPicking(false); @@ -171,7 +171,7 @@ export default function ImportTranscriptOption() { onSelect={(ctx) => { menuRef.current = ctx; const current = useEditorStore.getState().model; - if (isWhisperModel(current)) { + if (isAsrModel(current)) { previousModelRef.current = current; } // Do not set model to "import" until a file is chosen. Close the menu diff --git a/components/ModelSelector.tsx b/components/ModelSelector.tsx index 619642d..19986c8 100644 --- a/components/ModelSelector.tsx +++ b/components/ModelSelector.tsx @@ -25,7 +25,13 @@ import { TRANSCRIPT_LANGUAGES, type TranscriptLanguage, } from "@/lib/languages"; -import { MODELS, isWhisperModel, type ModelChoice } from "@/lib/models"; +import { + MODELS, + PARAKEET_INFO, + isParakeetModel, + isWhisperModel, + type ModelChoice, +} from "@/lib/models"; import { hydrateModelPreference, hydrateTranscriptLanguagePreference, @@ -187,9 +193,11 @@ export default function ModelSelector({ activeTrigger?.label ?? (isWhisperModel(model) ? MODELS[model].label - : model === "import" - ? "Import transcript" - : String(model)); + : isParakeetModel(model) + ? PARAKEET_INFO.label + : model === "import" + ? "Import transcript" + : String(model)); const languageInfo = TRANSCRIPT_LANGUAGES[transcriptLanguage]; const showLanguageInTrigger = isWhisperModel(model) && @@ -201,6 +209,7 @@ export default function ModelSelector({ <> + ); @@ -317,9 +326,19 @@ export function ModelOption({ const selected = selector.value === id; const resolvedLabel = - label ?? (isWhisperModel(id) ? MODELS[id].label : id); + label ?? + (isWhisperModel(id) + ? MODELS[id].label + : isParakeetModel(id) + ? PARAKEET_INFO.label + : id); const resolvedMeta = - meta ?? (isWhisperModel(id) ? MODELS[id].size : undefined); + meta ?? + (isWhisperModel(id) + ? MODELS[id].size + : isParakeetModel(id) + ? PARAKEET_INFO.size + : undefined); const optionCtx = useMemo( () => ({ diff --git a/components/UploadScreen.tsx b/components/UploadScreen.tsx index b012ebf..516c44a 100644 --- a/components/UploadScreen.tsx +++ b/components/UploadScreen.tsx @@ -277,6 +277,7 @@ export default function UploadScreen({ + diff --git a/hooks/useTranscriber.ts b/hooks/useTranscriber.ts index b413d4b..c8a696b 100644 --- a/hooks/useTranscriber.ts +++ b/hooks/useTranscriber.ts @@ -1,13 +1,13 @@ "use client"; import { useCallback, useEffect, useRef } from "react"; -import { isWhisperModel } from "@/lib/models"; +import { isAsrModel } from "@/lib/models"; import { useEditorStore } from "@/lib/store"; import type { WorkerResponse } from "@/lib/types"; let activeWorker: Worker | null = null; -/** Stop an in-flight Whisper job (e.g. after importing a transcript). */ +/** Stop an in-flight ASR job (e.g. after importing a transcript). */ export function cancelTranscription() { activeWorker?.terminate(); activeWorker = null; @@ -26,11 +26,11 @@ export function useTranscriber() { const transcribe = useCallback((audio: Float32Array, duration: number) => { const store = useEditorStore.getState(); - if (!isWhisperModel(store.model)) { - store.setError("Select Whisper Base or Small to transcribe."); + if (!isAsrModel(store.model)) { + store.setError("Select a speech model to transcribe."); return; } - const whisperModel = store.model; + const asrModel = store.model; const transcriptLanguage = store.transcriptLanguage; store.setStatus("transcribing"); store.setProgress({ message: "Loading speech model…", value: null }); @@ -44,7 +44,7 @@ export function useTranscriber() { activeWorker = workerRef.current; workerRef.current.onmessage = (event: MessageEvent) => { const s = useEditorStore.getState(); - // An imported transcript sets skipTranscription; ignore late Whisper results. + // An imported transcript sets skipTranscription; ignore late ASR results. if (s.skipTranscription) return; const msg = event.data; switch (msg.type) { @@ -73,7 +73,7 @@ export function useTranscriber() { // Transfer a copy so the original stays available for the waveform. const copy = audio.slice(); workerRef.current.postMessage( - { audio: copy, duration, model: whisperModel, language: transcriptLanguage }, + { audio: copy, duration, model: asrModel, language: transcriptLanguage }, [copy.buffer] ); }, []); diff --git a/lib/models.ts b/lib/models.ts index 5387a68..7eb288f 100644 --- a/lib/models.ts +++ b/lib/models.ts @@ -1,6 +1,9 @@ /** Transcription source choices offered on the upload screen. */ export type WhisperModel = "base" | "small"; -export type ModelChoice = WhisperModel | "import"; +/** NVIDIA Parakeet TDT 0.6B v3 via parakeet.js (ONNX / WebGPU). */ +export type ParakeetModel = "parakeet"; +export type AsrModel = WhisperModel | ParakeetModel; +export type ModelChoice = AsrModel | "import"; type DType = "fp32" | "fp16" | "q8" | "int8" | "uint8" | "q4" | "q4f16" | "bnb4"; @@ -60,32 +63,54 @@ export const MODELS: Record = { }, }; +/** + * UI metadata for Parakeet (not a transformers.js Whisper checkpoint). + * Download size depends on backend: WASM int8 ~670 MB, WebGPU fp16 ~1.2 GB. + */ +export const PARAKEET_INFO = { + /** parakeet.js model key → ysdede/parakeet-tdt-0.6b-v3-onnx */ + id: "parakeet-tdt-0.6b-v3", + label: "Parakeet TDT v3", + description: + "NVIDIA FastConformer — faster on WebGPU, strong EU-language accuracy. Auto-detects language.", + size: "~700 MB", +} as const; + export function isWhisperModel(value: unknown): value is WhisperModel { return value === "base" || value === "small"; } +export function isParakeetModel(value: unknown): value is ParakeetModel { + return value === "parakeet"; +} + +/** Models that run local ASR in the transcription worker (not import). */ +export function isAsrModel(value: unknown): value is AsrModel { + return isWhisperModel(value) || isParakeetModel(value); +} + export function isModelChoice(value: unknown): value is ModelChoice { - return isWhisperModel(value) || value === "import"; + return isAsrModel(value) || value === "import"; } const MODEL_STORAGE_KEY = "rescript.model"; -/** Read the last-selected Whisper model from localStorage (defaults to base). */ -export function loadModelPreference(): WhisperModel { +/** Read the last-selected ASR model from localStorage (defaults to base). */ +export function loadModelPreference(): AsrModel { if (typeof window === "undefined") return "base"; try { const raw = window.localStorage.getItem(MODEL_STORAGE_KEY); // Ignore a stale "import" preference — that choice is session-only until a // transcript file is picked again. - if (isWhisperModel(raw)) return raw; + if (isAsrModel(raw)) return raw; } catch { // private mode / disabled storage } return "base"; } -/** Persist the selected Whisper model for the next visit. */ -export function saveModelPreference(model: WhisperModel) { +/** Persist the selected ASR model for the next visit. */ +export function saveModelPreference(model: AsrModel) { if (typeof window === "undefined") return; try { window.localStorage.setItem(MODEL_STORAGE_KEY, model); diff --git a/lib/store.ts b/lib/store.ts index f0a3d25..844a772 100644 --- a/lib/store.ts +++ b/lib/store.ts @@ -25,6 +25,7 @@ import { trimEdgeResult, } from "./edits"; import { + isAsrModel, isModelChoice, isWhisperModel, loadModelPreference, @@ -59,19 +60,19 @@ interface EditorState { duration: number; /** Mono 16 kHz PCM of the media's audio track (used for waveform + ASR). */ audio: Float32Array | null; - /** Transcript source selected on the upload screen (Whisper or import). */ + /** Transcript source selected on the upload screen (ASR model or import). */ model: ModelChoice; - /** Language hint sent to Whisper when transcribing. */ + /** Language hint sent to Whisper when transcribing (Parakeet auto-detects). */ transcriptLanguage: TranscriptLanguage; /** * Caption file parsed on the upload screen when source is "import". - * Cleared when switching back to a Whisper model or after media loads. + * Cleared when switching back to an ASR model or after media loads. */ pendingTranscript: PendingTranscript | null; /** IndexedDB project id when this session is persisted; null for a fresh upload mid-pipeline. */ projectId: string | null; /** - * When true, Editor extracts audio for the waveform but skips Whisper + * When true, Editor extracts audio for the waveform but skips ASR * (restored projects / imported transcripts already have words). */ skipTranscription: boolean; @@ -289,7 +290,7 @@ export const useEditorStore = create((set, get) => ({ mediaKind: kind, projectId: null, skipTranscription: Boolean(imported), - model: imported ? "import" : isWhisperModel(current) ? current : "base", + model: imported ? "import" : isAsrModel(current) ? current : "base", pendingTranscript: null, status: "preparing", progress: { @@ -365,7 +366,7 @@ export const useEditorStore = create((set, get) => ({ }, setModel: (model) => { - if (isWhisperModel(model)) { + if (isAsrModel(model)) { saveModelPreference(model); set({ model, pendingTranscript: null }); } else { diff --git a/lib/types.ts b/lib/types.ts index cc4db3c..e1a9111 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -86,8 +86,8 @@ export interface WorkerRequest { audio: Float32Array; /** Total media duration in seconds (used for progress estimation). */ duration: number; - /** Which Whisper model to use (see lib/models.ts). */ - model: import("./models").WhisperModel; - /** Whisper language code for the transcript. */ + /** Which local ASR model to use (see lib/models.ts). */ + model: import("./models").AsrModel; + /** Whisper language code for the transcript (ignored by Parakeet auto-detect). */ language: import("./languages").TranscriptLanguage; } diff --git a/next.config.ts b/next.config.ts index f08d0fb..53bb86e 100644 --- a/next.config.ts +++ b/next.config.ts @@ -9,6 +9,8 @@ const basePath = process.env.NEXT_PUBLIC_BASE_PATH ?? ""; const nextConfig: NextConfig = { reactStrictMode: true, + // parakeet.js ships as raw ESM from src/; transpile for the worker bundle. + transpilePackages: ["parakeet.js"], ...(isExport ? { output: "export" as const, diff --git a/package-lock.json b/package-lock.json index ac8d638..9aaea6d 100644 --- a/package-lock.json +++ b/package-lock.json @@ -18,6 +18,7 @@ "coi-serviceworker": "^0.1.7", "lucide-react": "^1.27.0", "next": "16.2.12", + "parakeet.js": "^1.4.4", "react": "19.2.4", "react-dom": "19.2.4", "react-resizable-panels": "^4.12.2", @@ -8505,6 +8506,35 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/parakeet.js": { + "version": "1.4.4", + "resolved": "https://registry.npmjs.org/parakeet.js/-/parakeet.js-1.4.4.tgz", + "integrity": "sha512-+tYIDgp799bvsDye3y2lFyspYHgTOJ5YMZi0mWYABXyJjp/3NZmNaMTnS9jhf4qlIQJoBk0Nmj+UwsI3zk+Aqg==", + "license": "MIT", + "dependencies": { + "onnxruntime-web": "1.24.1" + } + }, + "node_modules/parakeet.js/node_modules/onnxruntime-common": { + "version": "1.24.1", + "resolved": "https://registry.npmjs.org/onnxruntime-common/-/onnxruntime-common-1.24.1.tgz", + "integrity": "sha512-UnV15u4p4XxoIV+jFP4hXPsW93s3QrwLSpi20HUDYHoTfI4z4sjzex3L4XDOxGGZJ/M/catrwAG2go958UQq0w==", + "license": "MIT" + }, + "node_modules/parakeet.js/node_modules/onnxruntime-web": { + "version": "1.24.1", + "resolved": "https://registry.npmjs.org/onnxruntime-web/-/onnxruntime-web-1.24.1.tgz", + "integrity": "sha512-i2u395dv+ZEQBdH+aORvlu19Bzvlg5AXJ7wjxnL350hknOP9z0UeP3pVfjkpMEWMPy2T6nCQxetKTmNia6wSzg==", + "license": "MIT", + "dependencies": { + "flatbuffers": "^25.1.24", + "guid-typescript": "^1.0.9", + "long": "^5.2.3", + "onnxruntime-common": "1.24.1", + "platform": "^1.3.6", + "protobufjs": "^7.2.4" + } + }, "node_modules/parent-module": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/parent-module/-/parent-module-1.0.1.tgz", diff --git a/package.json b/package.json index 85bd298..b38a80b 100644 --- a/package.json +++ b/package.json @@ -39,6 +39,7 @@ "coi-serviceworker": "^0.1.7", "lucide-react": "^1.27.0", "next": "16.2.12", + "parakeet.js": "^1.4.4", "react": "19.2.4", "react-dom": "19.2.4", "react-resizable-panels": "^4.12.2", diff --git a/scripts/copy-assets.mjs b/scripts/copy-assets.mjs index d732490..732388f 100644 --- a/scripts/copy-assets.mjs +++ b/scripts/copy-assets.mjs @@ -3,9 +3,17 @@ * served fully offline (no CDN requests at runtime): * - @ffmpeg/core-mt -> public/vendor/ffmpeg/ (audio extraction + export) * - onnxruntime-web -> public/vendor/ort/ (transformers.js inference) + * - parakeet.js ORT -> public/vendor/ort-parakeet/ (Parakeet TDT inference) * Runs automatically on `npm install` (postinstall). */ -import { cpSync, mkdirSync, readdirSync, readFileSync, writeFileSync } from "node:fs"; +import { + cpSync, + existsSync, + mkdirSync, + readFileSync, + readdirSync, + writeFileSync, +} from "node:fs"; import { dirname, join } from "node:path"; import { fileURLToPath } from "node:url"; @@ -30,15 +38,70 @@ for (const f of readdirSync(ffmpegClassSrc)) { } } +function copyOrtWasm(srcDist, dst) { + mkdirSync(dst, { recursive: true }); + for (const f of readdirSync(srcDist)) { + if (/^ort-wasm-simd-threaded.*\.(wasm|mjs)$/.test(f)) { + cpSync(join(srcDist, f), join(dst, f)); + } + } +} + const ortSrc = join(root, "node_modules/onnxruntime-web/dist"); const ortDst = join(root, "public/vendor/ort"); -mkdirSync(ortDst, { recursive: true }); -for (const f of readdirSync(ortSrc)) { - if (/^ort-wasm-simd-threaded.*\.(wasm|mjs)$/.test(f)) { - cpSync(join(ortSrc, f), join(ortDst, f)); +copyOrtWasm(ortSrc, ortDst); + +// Parakeet.js pins onnxruntime-web@1.24.1 (nested). Keep its WASM separate so +// the JS package and binaries stay version-matched. +const parakeetOrtSrc = join( + root, + "node_modules/parakeet.js/node_modules/onnxruntime-web/dist" +); +const parakeetOrtDst = join(root, "public/vendor/ort-parakeet"); +if (existsSync(parakeetOrtSrc)) { + copyOrtWasm(parakeetOrtSrc, parakeetOrtDst); +} else { + // Hoisted install: fall back to the top-level ORT package. + copyOrtWasm(ortSrc, parakeetOrtDst); +} + +/** + * parakeet.js@1.4.4 accepts `wasmPaths` in fromUrls/fromHub but initOrt never + * applies the argument — it only sets a jsDelivr CDN default. Patch that so + * Rescript can serve WASM same-origin (offline after first model download). + */ +function patchParakeetWasmPaths() { + const backendPath = join(root, "node_modules/parakeet.js/src/backend.js"); + if (!existsSync(backendPath)) return; + let src = readFileSync(backendPath, "utf8"); + if (src.includes("/* rescript-wasmPaths-patch */")) return; + + // Package may ship CRLF; normalize for matching then restore EOL style. + const eol = src.includes("\r\n") ? "\r\n" : "\n"; + const normalized = src.replace(/\r\n/g, "\n"); + const needle = + " // Set up WASM paths first (needed for all backends)\n" + + " if (!ort.env.wasm.wasmPaths) {"; + const replacement = + " // Set up WASM paths first (needed for all backends)\n" + + " /* rescript-wasmPaths-patch */\n" + + " if (wasmPaths) {\n" + + " ort.env.wasm.wasmPaths = wasmPaths;\n" + + " } else if (!ort.env.wasm.wasmPaths) {"; + if (!normalized.includes(needle)) { + console.warn( + "[copy-assets] Could not patch parakeet.js wasmPaths (needle not found)" + ); + return; } + let patched = normalized.replace(needle, replacement); + if (eol === "\r\n") patched = patched.replace(/\n/g, "\r\n"); + writeFileSync(backendPath, patched); + console.log("[copy-assets] Patched parakeet.js initOrt to honor wasmPaths"); } +patchParakeetWasmPaths(); + // coi-serviceworker provides COOP/COEP headers on static hosts (GitHub Pages) // that can't send them, keeping cross-origin isolation for SharedArrayBuffer. // A config prelude is prepended: always use COEP "credentialless" (needed for @@ -57,4 +120,6 @@ writeFileSync( coiPrelude + readFileSync(coiSrc, "utf8") ); -console.log("[copy-assets] ffmpeg core + onnxruntime wasm + coi-serviceworker copied to public/"); +console.log( + "[copy-assets] ffmpeg core + onnxruntime wasm + coi-serviceworker copied to public/" +); diff --git a/tests/models-test.ts b/tests/models-test.ts new file mode 100644 index 0000000..1cfd177 --- /dev/null +++ b/tests/models-test.ts @@ -0,0 +1,33 @@ +/** + * Model choice helpers for Whisper / Parakeet / import. + */ +import { + MODELS, + PARAKEET_INFO, + isAsrModel, + isModelChoice, + isParakeetModel, + isWhisperModel, +} from "../lib/models"; + +function assert(cond: boolean, msg: string) { + if (!cond) throw new Error(msg); +} + +assert(isWhisperModel("base"), "base is Whisper"); +assert(isWhisperModel("small"), "small is Whisper"); +assert(!isWhisperModel("parakeet"), "parakeet is not Whisper"); +assert(isParakeetModel("parakeet"), "parakeet is Parakeet"); +assert(isAsrModel("parakeet"), "parakeet is ASR"); +assert(isAsrModel("base"), "base is ASR"); +assert(!isAsrModel("import"), "import is not ASR"); +assert(isModelChoice("import"), "import is a model choice"); +assert(isModelChoice("parakeet"), "parakeet is a model choice"); +assert(!isModelChoice("tiny"), "tiny is not a model choice"); + +assert(PARAKEET_INFO.id === "parakeet-tdt-0.6b-v3", "parakeet hub id"); +assert(typeof PARAKEET_INFO.label === "string", "parakeet label"); +assert(typeof MODELS.base.id === "string", "whisper base id"); +assert(typeof MODELS.small.id === "string", "whisper small id"); + +console.log("models-test: ok"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index ff76be0..e5f3f77 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -2,13 +2,15 @@ * Transcription worker: runs entirely in the browser. * * 1. Silero VAD (energy fallback) finds speech segments; silence is skipped. - * 2. A Whisper-family model (see lib/models.ts) transcribes each segment with - * per-word timestamps, remapped onto the original timeline. + * 2. ASR (Whisper via transformers.js, or Parakeet TDT v3 via parakeet.js) + * transcribes each segment with per-word timestamps, remapped onto the + * original timeline. * 3. Pyannote segmentation 3.0 assigns a speaker to each word. * - * Models are fetched from the Hugging Face Hub on first use and cached in the - * browser Cache Storage; every run after that is fully offline. The ONNX - * runtime WASM binaries are served from /vendor/ort (same origin). + * Models are fetched from the Hugging Face Hub on first use and cached + * (Cache Storage for Whisper / IndexedDB for Parakeet); every run after that + * is fully offline. The ONNX runtime WASM binaries are served from + * /vendor/ort (same origin). */ import { pipeline, @@ -21,7 +23,14 @@ import { type AutomaticSpeechRecognitionPipeline, } from "@huggingface/transformers"; import type { Word, WorkerRequest, WorkerResponse } from "@/lib/types"; -import { MODELS, type WhisperModel } from "@/lib/models"; +import { + MODELS, + PARAKEET_INFO, + isParakeetModel, + isWhisperModel, + type AsrModel, + type WhisperModel, +} from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; import { alignWordsToSpeech } from "@/lib/align"; import { @@ -34,9 +43,12 @@ import { import { isWebGpuDeviceLostError } from "@/lib/webgpu"; env.allowLocalModels = false; +const ORT_WASM_PATHS = `${process.env.NEXT_PUBLIC_BASE_PATH ?? ""}/vendor/ort/`; +/** Parakeet.js pins onnxruntime-web@1.24.1 — keep its WASM on a separate path. */ +const PARAKEET_ORT_WASM_PATHS = `${process.env.NEXT_PUBLIC_BASE_PATH ?? ""}/vendor/ort-parakeet/`; // Serve onnxruntime-web WASM from our own origin (offline friendly). if (env.backends?.onnx?.wasm) { - env.backends.onnx.wasm.wasmPaths = `${process.env.NEXT_PUBLIC_BASE_PATH ?? ""}/vendor/ort/`; + env.backends.onnx.wasm.wasmPaths = ORT_WASM_PATHS; } const DIARIZATION_MODEL = "onnx-community/pyannote-segmentation-3.0"; @@ -466,192 +478,431 @@ function assignSpeakers(words: Word[], segments: DiarizationSegment[]) { } } -self.onmessage = async (event: MessageEvent) => { - const { audio, duration, model, language } = event.data; - try { - const choice: WhisperModel = model ?? "base"; - const transcriptLanguage = language ?? "en"; - - // Overlap Whisper + Silero downloads; diarizer warms in the background. - getDiarizer().catch(() => {}); - const [asr, vad] = await Promise.all([getAsr(choice), getVad()]); - let transcriber = asr; - - post({ type: "progress", message: "Detecting speech…", value: 0 }); - const { segments: speechSegments, frames: speechFrames } = - await detectSpeechSegments(audio, vad); - - const { verbatimPrompt } = MODELS[choice]; - const promptedIds = verbatimPrompt - ? buildPromptedDecoderIds(transcriber, verbatimPrompt, transcriptLanguage) - : null; - if (verbatimPrompt && !promptedIds) { - console.warn("Could not build verbatim prompt tokens; using default decoding."); +type ParakeetInstance = { + transcribe: ( + audio: Float32Array, + sampleRate?: number, + opts?: { + returnTimestamps?: boolean; + timeOffset?: number; } + ) => Promise<{ + utterance_text: string; + words: Array<{ text: string; start_time: number; end_time: number }>; + }>; +}; - const speechSamples = speechSegments.reduce( - (n, s) => n + (s.endSample - s.startSample), - 0 - ); +let parakeetPromise: Promise | null = null; +let parakeetBackend: "webgpu" | "wasm" = "wasm"; - post({ type: "progress", message: "Transcribing…", value: 0 }); - - let partial = ""; - // Use 29s instead of 30: transformers.js has a known word-timestamp bug - // at exactly chunk_length_s=30 (#1357 / #1358); 29 is the common workaround. - const chunkLength = 29; - const stride = 5; - const timePrecision = - // @ts-expect-error feature_extractor config is untyped - (transcriber.processor.feature_extractor.config.chunk_length ?? 30) / - // @ts-expect-error model config is untyped - (transcriber.model.config.max_source_positions ?? 1500); - - let speechDone = 0; - let transcribed = 0; - let chunkFloor = 0; - let chunkTokens = 0; - let avgChunkDelta = - speechSamples > 0 - ? Math.min(0.15, ((chunkLength - stride) * VAD_SAMPLE_RATE) / speechSamples) - : 0.05; - - const reportProgress = (segmentLocalT: number, segmentSamples: number) => { - const local = Math.min( - segmentSamples, - Math.max(0, segmentLocalT * VAD_SAMPLE_RATE) - ); - const next = Math.max( - transcribed, - Math.min(1, speechSamples > 0 ? (speechDone + local) / speechSamples : 1) - ); - const realDelta = next - chunkFloor; - if (realDelta > 0) avgChunkDelta = avgChunkDelta * 0.5 + realDelta * 0.5; - chunkFloor = next; - chunkTokens = 0; - transcribed = next; - post({ type: "progress", message: "Transcribing…", value: transcribed }); - }; +function makeParakeetHubProgress(label: string) { + const files = new Map(); + let best = 0; + return (p: { loaded: number; total: number; file: string }) => { + if (!p.file || !(p.total > 0)) return; + files.set(p.file, { loaded: p.loaded, total: p.total }); + let loaded = 0; + let total = 0; + for (const f of files.values()) { + loaded += f.loaded; + total += f.total; + } + if (total <= 0) return; + best = Math.max(best, Math.min(1, loaded / total)); + post({ type: "progress", message: label, value: best }); + }; +} + +/** + * parakeet.js accepts `wasmPaths` on fromHub/fromUrls. A postinstall patch in + * scripts/copy-assets.mjs makes initOrt honor that argument (upstream 1.4.4 + * ignored it and defaulted to a jsDelivr CDN). + */ +async function getParakeet(): Promise { + if (!parakeetPromise) { + parakeetPromise = (async () => { + const { fromHub } = await import("parakeet.js"); + const device = await pickDevice(); + const progress = makeParakeetHubProgress("Downloading speech model…"); + const common = { + preprocessorBackend: "js" as const, + progress, + wasmPaths: PARAKEET_ORT_WASM_PATHS, + }; - /** Nudge the bar forward between chunk boundaries as tokens stream in. */ - const interpolateProgress = () => { - chunkTokens++; - // n/(n+8): 0.11 at token 1, 0.5 at token 8, 0.9 at token 72 — strictly - // increasing, so it can never get stuck as long as tokens keep coming. - const frac = chunkTokens / (chunkTokens + 8); - const interpolated = Math.min(0.999, chunkFloor + frac * avgChunkDelta); - if (interpolated > transcribed) { - transcribed = interpolated; - post({ type: "progress", message: "Transcribing…", value: transcribed }); + // WebGPU cannot run the int8 encoder; fp16 (~1.2 GB) is the practical + // WebGPU path. WASM int8 (~670 MB) is the compatibility / size fallback. + if (device === "webgpu") { + try { + post({ + type: "progress", + message: "Downloading speech model…", + value: null, + }); + const model = await fromHub(PARAKEET_INFO.id, { + ...common, + backend: "webgpu", + encoderQuant: "fp16", + decoderQuant: "int8", + }); + parakeetBackend = "webgpu"; + asrDevice = "webgpu"; + return model as ParakeetInstance; + } catch (err) { + console.warn( + "Parakeet WebGPU/fp16 load failed; falling back to WASM int8.", + err + ); + } } - }; - const asrOptions = { - chunk_length_s: chunkLength, - stride_length_s: stride, - return_timestamps: "word" as const, - // Anti-repetition: Whisper-base on multi-minute audio often falls into - // loops like "little bit of a little bit of a…" near chunk boundaries - // or silence. Keep penalty mild — 1.15 truncates multi-speaker clips - // mid-utterance (second speaker dropped on continuous speech). - no_repeat_ngram_size: 4, - repetition_penalty: 1.05, - ...(promptedIds - ? { decoder_input_ids: promptedIds } - : { language: transcriptLanguage }), + post({ + type: "progress", + message: "Downloading speech model…", + value: null, + }); + const model = await fromHub(PARAKEET_INFO.id, { + ...common, + backend: "wasm", + encoderQuant: "int8", + decoderQuant: "int8", + }); + parakeetBackend = "wasm"; + asrDevice = "wasm"; + return model as ParakeetInstance; + })(); + parakeetPromise.catch(() => { + parakeetPromise = null; + }); + } + return parakeetPromise; +} + +/** Map Parakeet word timestamps onto the original media timeline. */ +function wordsFromParakeet( + words: Array<{ text: string; start_time: number; end_time: number }>, + offsetS: number, + segmentDuration: number, + mediaDuration: number +): Word[] { + const clampLocal = (t: number) => Math.min(Math.max(t, 0), segmentDuration); + const usable = words + .map((w) => ({ + text: w.text.trim(), + start: w.start_time, + end: w.end_time, + })) + .filter((w) => w.text.length > 0); + + return usable.map((w, i) => { + const localStart = clampLocal(w.start); + const next = usable[i + 1]; + const nextStart = next ? clampLocal(next.start) : segmentDuration; + const localEnd = + Number.isFinite(w.end) && w.end <= segmentDuration + 0.05 + ? clampLocal(w.end) + : Math.min(localStart + FALLBACK_WORD_S, Math.max(localStart, nextStart)); + + let start = offsetS + localStart; + let end = offsetS + Math.max(localEnd, localStart); + if (mediaDuration > 0) { + start = Math.min(start, mediaDuration); + end = Math.min(end, mediaDuration); + } + start = Math.max(0, start); + return { + id: i, + text: w.text, + start, + end: Math.max(end, start + 0.02), + speaker: 0, + deleted: false, }; + }); +} - const rawWords: Word[] = []; - const leadPadSamples = Math.floor(WHISPER_LEAD_PAD_S * VAD_SAMPLE_RATE); - for (const seg of speechSegments) { - const segmentSamples = seg.endSample - seg.startSample; - // Copy into a fresh buffer with leading silence. Views with a non-zero - // byteOffset have caused incomplete ASR with onnxruntime-web; starting - // mid-speech with no lead-in also drops later speakers on mixed clips. - const slice = new Float32Array(leadPadSamples + segmentSamples); - slice.set(audio.subarray(seg.startSample, seg.endSample), leadPadSamples); - const sliceDuration = slice.length / VAD_SAMPLE_RATE; - const offsetS = seg.startSample / VAD_SAMPLE_RATE - WHISPER_LEAD_PAD_S; - - // Snapshot progress so a failed WebGPU attempt can be rolled back before - // the WASM retry of this same segment. - const partialBefore = partial; - const progressBefore = { transcribed, chunkFloor, chunkTokens }; - - const runSlice = async () => { - // Each generate() window consumes `chunkLength - 2 * stride` seconds of - // new audio, and the streamer's timestamps rewind to ~0 when the next - // window starts. A timestamp lower than the last one seen marks that - // boundary; accumulate the offset to recover segment-local time. - const windowJumpS = chunkLength - 2 * stride; - let windowOffsetS = 0; - let lastChunkStartT = 0; - const tokenizer = transcriber.tokenizer as ConstructorParameters< - typeof WhisperTextStreamer - >[0]; - const streamer = new WhisperTextStreamer(tokenizer, { - skip_prompt: true, - time_precision: timePrecision, - on_chunk_start: (t: number) => { - if (t < lastChunkStartT) windowOffsetS += windowJumpS; - lastChunkStartT = t; - reportProgress( - Math.max(0, windowOffsetS + t - WHISPER_LEAD_PAD_S), - segmentSamples - ); - }, - callback_function: (text: string) => { - partial += text; - post({ type: "partial", text: partial }); - interpolateProgress(); - }, - }); - const output = await transcriber(slice, { ...asrOptions, streamer }); - const result = Array.isArray(output) ? output[0] : output; - return (result.chunks ?? []) as AsrChunk[]; - }; +async function finishWithDiarization( + words: Word[], + audio: Float32Array +): Promise { + try { + post({ type: "progress", message: "Identifying speakers…", value: null }); + const segments = await diarize(audio); + assignSpeakers(words, segments); + } catch (err) { + console.warn("Speaker diarization failed; using a single speaker.", err); + } + return words; +} + +async function runParakeet( + audio: Float32Array, + duration: number +): Promise { + getDiarizer().catch(() => {}); + const [model, vad] = await Promise.all([getParakeet(), getVad()]); + + post({ type: "progress", message: "Detecting speech…", value: 0 }); + const { segments: speechSegments } = await detectSpeechSegments(audio, vad); + + post({ type: "progress", message: "Transcribing…", value: 0 }); + const speechSamples = speechSegments.reduce( + (n, s) => n + (s.endSample - s.startSample), + 0 + ); + + const rawWords: Word[] = []; + let partial = ""; + let speechDone = 0; + + for (const seg of speechSegments) { + const segmentSamples = seg.endSample - seg.startSample; + // Fresh buffer: non-zero byteOffset views have caused incomplete ASR with + // onnxruntime-web in the Whisper path; keep the same hygiene here. + const slice = audio.slice(seg.startSample, seg.endSample); + const sliceDuration = slice.length / VAD_SAMPLE_RATE; + const offsetS = seg.startSample / VAD_SAMPLE_RATE; + + const runSlice = () => + model.transcribe(slice, VAD_SAMPLE_RATE, { + returnTimestamps: true, + timeOffset: 0, + }); - let chunks: AsrChunk[]; - try { - chunks = await runSlice(); - } catch (err) { - // Windows screen lock tears down WebGPU mid-OrtRun. Fall back to WASM - // and retry this segment once so the job can finish. - if (asrDevice !== "webgpu" || !isWebGpuDeviceLostError(err)) throw err; - console.warn( - "WebGPU lost during transcription (often after screen lock); falling back to WASM.", - err - ); - partial = partialBefore; - transcribed = progressBefore.transcribed; - chunkFloor = progressBefore.chunkFloor; - chunkTokens = progressBefore.chunkTokens; - post({ type: "partial", text: partial }); - transcriber = await fallbackAsrToWasm(choice); - chunks = await runSlice(); + let result: Awaited>; + try { + result = await runSlice(); + } catch (err) { + if (parakeetBackend !== "webgpu" || !isWebGpuDeviceLostError(err)) { + throw err; } + console.warn( + "WebGPU lost during Parakeet transcription; reloading on WASM.", + err + ); + forceWasm = true; + parakeetPromise = null; + const wasmModel = await getParakeet(); + result = await wasmModel.transcribe(slice, VAD_SAMPLE_RATE, { + returnTimestamps: true, + timeOffset: 0, + }); + } + + rawWords.push( + ...wordsFromParakeet(result.words ?? [], offsetS, sliceDuration, duration) + ); + const piece = (result.utterance_text ?? "").trim(); + if (piece) { + partial = partial ? `${partial} ${piece}` : piece; + post({ type: "partial", text: partial }); + } + + speechDone += segmentSamples; + const value = + speechSamples > 0 ? Math.min(1, speechDone / speechSamples) : 1; + post({ type: "progress", message: "Transcribing…", value }); + } - rawWords.push(...wordsFromChunks(chunks, offsetS, sliceDuration, duration)); - speechDone += segmentSamples; - reportProgress(0, 0); + const cleaned = cleanTranscript(rawWords); + return finishWithDiarization(cleaned, audio); +} + +async function runWhisper( + audio: Float32Array, + duration: number, + choice: WhisperModel, + transcriptLanguage: WorkerRequest["language"] +): Promise { + // Overlap Whisper + Silero downloads; diarizer warms in the background. + getDiarizer().catch(() => {}); + const [asr, vad] = await Promise.all([getAsr(choice), getVad()]); + let transcriber = asr; + + post({ type: "progress", message: "Detecting speech…", value: 0 }); + const { segments: speechSegments, frames: speechFrames } = + await detectSpeechSegments(audio, vad); + + const { verbatimPrompt } = MODELS[choice]; + const promptedIds = verbatimPrompt + ? buildPromptedDecoderIds(transcriber, verbatimPrompt, transcriptLanguage) + : null; + if (verbatimPrompt && !promptedIds) { + console.warn("Could not build verbatim prompt tokens; using default decoding."); + } + + const speechSamples = speechSegments.reduce( + (n, s) => n + (s.endSample - s.startSample), + 0 + ); + + post({ type: "progress", message: "Transcribing…", value: 0 }); + + let partial = ""; + // Use 29s instead of 30: transformers.js has a known word-timestamp bug + // at exactly chunk_length_s=30 (#1357 / #1358); 29 is the common workaround. + const chunkLength = 29; + const stride = 5; + const timePrecision = + // @ts-expect-error feature_extractor config is untyped + (transcriber.processor.feature_extractor.config.chunk_length ?? 30) / + // @ts-expect-error model config is untyped + (transcriber.model.config.max_source_positions ?? 1500); + + let speechDone = 0; + let transcribed = 0; + let chunkFloor = 0; + let chunkTokens = 0; + let avgChunkDelta = + speechSamples > 0 + ? Math.min(0.15, ((chunkLength - stride) * VAD_SAMPLE_RATE) / speechSamples) + : 0.05; + + const reportProgress = (segmentLocalT: number, segmentSamples: number) => { + const local = Math.min( + segmentSamples, + Math.max(0, segmentLocalT * VAD_SAMPLE_RATE) + ); + const next = Math.max( + transcribed, + Math.min(1, speechSamples > 0 ? (speechDone + local) / speechSamples : 1) + ); + const realDelta = next - chunkFloor; + if (realDelta > 0) avgChunkDelta = avgChunkDelta * 0.5 + realDelta * 0.5; + chunkFloor = next; + chunkTokens = 0; + transcribed = next; + post({ type: "progress", message: "Transcribing…", value: transcribed }); + }; + + /** Nudge the bar forward between chunk boundaries as tokens stream in. */ + const interpolateProgress = () => { + chunkTokens++; + // n/(n+8): 0.11 at token 1, 0.5 at token 8, 0.9 at token 72 — strictly + // increasing, so it can never get stuck as long as tokens keep coming. + const frac = chunkTokens / (chunkTokens + 8); + const interpolated = Math.min(0.999, chunkFloor + frac * avgChunkDelta); + if (interpolated > transcribed) { + transcribed = interpolated; + post({ type: "progress", message: "Transcribing…", value: transcribed }); } + }; - // Post-process: collapse leftover n-gram loops and drop known hallucination - // phrases ("I'm sorry", "thanks for watching", …) that slip past decoding. - const cleaned = cleanTranscript(rawWords); + const asrOptions = { + chunk_length_s: chunkLength, + stride_length_s: stride, + return_timestamps: "word" as const, + // Anti-repetition: Whisper-base on multi-minute audio often falls into + // loops like "little bit of a little bit of a…" near chunk boundaries + // or silence. Keep penalty mild — 1.15 truncates multi-speaker clips + // mid-utterance (second speaker dropped on continuous speech). + no_repeat_ngram_size: 4, + repetition_penalty: 1.05, + ...(promptedIds + ? { decoder_input_ids: promptedIds } + : { language: transcriptLanguage }), + }; - // Whisper's DTW word timestamps run consistently late (~0.2 s on the test - // clips). Realign them against the VAD flags before diarization, so speakers - // are assigned from corrected times too. - const words = alignWordsToSpeech(cleaned, speechFrames, { duration }); + const rawWords: Word[] = []; + const leadPadSamples = Math.floor(WHISPER_LEAD_PAD_S * VAD_SAMPLE_RATE); + for (const seg of speechSegments) { + const segmentSamples = seg.endSample - seg.startSample; + // Copy into a fresh buffer with leading silence. Views with a non-zero + // byteOffset have caused incomplete ASR with onnxruntime-web; starting + // mid-speech with no lead-in also drops later speakers on mixed clips. + const slice = new Float32Array(leadPadSamples + segmentSamples); + slice.set(audio.subarray(seg.startSample, seg.endSample), leadPadSamples); + const sliceDuration = slice.length / VAD_SAMPLE_RATE; + const offsetS = seg.startSample / VAD_SAMPLE_RATE - WHISPER_LEAD_PAD_S; + + // Snapshot progress so a failed WebGPU attempt can be rolled back before + // the WASM retry of this same segment. + const partialBefore = partial; + const progressBefore = { transcribed, chunkFloor, chunkTokens }; + + const runSlice = async () => { + // Each generate() window consumes `chunkLength - 2 * stride` seconds of + // new audio, and the streamer's timestamps rewind to ~0 when the next + // window starts. A timestamp lower than the last one seen marks that + // boundary; accumulate the offset to recover segment-local time. + const windowJumpS = chunkLength - 2 * stride; + let windowOffsetS = 0; + let lastChunkStartT = 0; + const tokenizer = transcriber.tokenizer as ConstructorParameters< + typeof WhisperTextStreamer + >[0]; + const streamer = new WhisperTextStreamer(tokenizer, { + skip_prompt: true, + time_precision: timePrecision, + on_chunk_start: (t: number) => { + if (t < lastChunkStartT) windowOffsetS += windowJumpS; + lastChunkStartT = t; + reportProgress( + Math.max(0, windowOffsetS + t - WHISPER_LEAD_PAD_S), + segmentSamples + ); + }, + callback_function: (text: string) => { + partial += text; + post({ type: "partial", text: partial }); + interpolateProgress(); + }, + }); + const output = await transcriber(slice, { ...asrOptions, streamer }); + const result = Array.isArray(output) ? output[0] : output; + return (result.chunks ?? []) as AsrChunk[]; + }; - // Best-effort speaker diarization; a failure should not lose the transcript. + let chunks: AsrChunk[]; try { - post({ type: "progress", message: "Identifying speakers…", value: null }); - const segments = await diarize(audio); - assignSpeakers(words, segments); + chunks = await runSlice(); } catch (err) { - console.warn("Speaker diarization failed; using a single speaker.", err); + // Windows screen lock tears down WebGPU mid-OrtRun. Fall back to WASM + // and retry this segment once so the job can finish. + if (asrDevice !== "webgpu" || !isWebGpuDeviceLostError(err)) throw err; + console.warn( + "WebGPU lost during transcription (often after screen lock); falling back to WASM.", + err + ); + partial = partialBefore; + transcribed = progressBefore.transcribed; + chunkFloor = progressBefore.chunkFloor; + chunkTokens = progressBefore.chunkTokens; + post({ type: "partial", text: partial }); + transcriber = await fallbackAsrToWasm(choice); + chunks = await runSlice(); + } + + rawWords.push(...wordsFromChunks(chunks, offsetS, sliceDuration, duration)); + speechDone += segmentSamples; + reportProgress(0, 0); + } + + // Post-process: collapse leftover n-gram loops and drop known hallucination + // phrases ("I'm sorry", "thanks for watching", …) that slip past decoding. + const cleaned = cleanTranscript(rawWords); + + // Whisper's DTW word timestamps run consistently late (~0.2 s on the test + // clips). Realign them against the VAD flags before diarization, so speakers + // are assigned from corrected times too. + const words = alignWordsToSpeech(cleaned, speechFrames, { duration }); + + return finishWithDiarization(words, audio); +} + +self.onmessage = async (event: MessageEvent) => { + const { audio, duration, model, language } = event.data; + try { + const choice: AsrModel = model ?? "base"; + const transcriptLanguage = language ?? "en"; + + let words: Word[]; + if (isParakeetModel(choice)) { + words = await runParakeet(audio, duration); + } else if (isWhisperModel(choice)) { + words = await runWhisper(audio, duration, choice, transcriptLanguage); + } else { + throw new Error(`Unknown speech model: ${String(choice)}`); } post({ type: "complete", words }); From 49a64b5d02ae1314e1ba8515ed2ba947931f6b21 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 2 Aug 2026 09:48:03 +0000 Subject: [PATCH 2/5] Load Parakeet through weightlift ModelManager MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Register Parakeet alongside Whisper in the shared ASR registry so download progress, IndexedDB cache labeling, and WebGPU→WASM fallback use the same weightlift path after merging main. Co-authored-by: Wassim Gharbi --- lib/models.ts | 4 +- tests/models-test.ts | 4 + workers/transcription.worker.ts | 311 ++++++++++++++++++-------------- 3 files changed, 178 insertions(+), 141 deletions(-) diff --git a/lib/models.ts b/lib/models.ts index 7eb288f..b41ca41 100644 --- a/lib/models.ts +++ b/lib/models.ts @@ -68,8 +68,10 @@ export const MODELS: Record = { * Download size depends on backend: WASM int8 ~670 MB, WebGPU fp16 ~1.2 GB. */ export const PARAKEET_INFO = { - /** parakeet.js model key → ysdede/parakeet-tdt-0.6b-v3-onnx */ + /** parakeet.js model key (weightlift registry id). */ id: "parakeet-tdt-0.6b-v3", + /** Hugging Face repo used by parakeet.js hub downloads / IndexedDB cache keys. */ + repoId: "ysdede/parakeet-tdt-0.6b-v3-onnx", label: "Parakeet TDT v3", description: "NVIDIA FastConformer — faster on WebGPU, strong EU-language accuracy. Auto-detects language.", diff --git a/tests/models-test.ts b/tests/models-test.ts index 1cfd177..183690e 100644 --- a/tests/models-test.ts +++ b/tests/models-test.ts @@ -26,6 +26,10 @@ assert(isModelChoice("parakeet"), "parakeet is a model choice"); assert(!isModelChoice("tiny"), "tiny is not a model choice"); assert(PARAKEET_INFO.id === "parakeet-tdt-0.6b-v3", "parakeet hub id"); +assert( + PARAKEET_INFO.repoId === "ysdede/parakeet-tdt-0.6b-v3-onnx", + "parakeet HF repo id" +); assert(typeof PARAKEET_INFO.label === "string", "parakeet label"); assert(typeof MODELS.base.id === "string", "whisper base id"); assert(typeof MODELS.small.id === "string", "whisper small id"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 4671ece..3649ef3 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -7,10 +7,10 @@ * original timeline. * 3. Pyannote segmentation 3.0 assigns a speaker to each word. * - * Models are fetched from the Hugging Face Hub on first use and cached - * (Cache Storage for Whisper / IndexedDB for Parakeet); every run after that - * is fully offline. The ONNX runtime WASM binaries are served from - * /vendor/ort (same origin). + * Both ASR backends are registered in a weightlift ModelManager so download + * progress, cache labeling, and WebGPU→WASM fallback share one path. Weights + * land in Cache Storage (Whisper) or IndexedDB (Parakeet); later runs are + * offline. ORT WASM is served same-origin from /vendor/ort* . */ import { pipeline, @@ -22,7 +22,7 @@ import { env, type AutomaticSpeechRecognitionPipeline, } from "@huggingface/transformers"; -import { ModelManager } from "weightlift"; +import { ModelManager, type ModelDefinition } from "weightlift"; import { fallbackDevicePolicy, transformersModel, @@ -77,30 +77,160 @@ const post = (msg: WorkerResponse, transfer: Transferable[] = []) => /** Device the current ASR pipeline is running on. */ let asrDevice: "webgpu" | "wasm" = "wasm"; +type ParakeetInstance = { + transcribe: ( + audio: Float32Array, + sampleRate?: number, + opts?: { + returnTimestamps?: boolean; + timeOffset?: number; + } + ) => Promise<{ + utterance_text: string; + words: Array<{ text: string; start_time: number; end_time: number }>; + }>; +}; + +const PARAKEET_CACHE_DB = "parakeet-cache-db"; +const PARAKEET_CACHE_STORE = "file-store"; + +/** Whether Parakeet ONNX weights already sit in parakeet.js IndexedDB. */ +async function isParakeetCached(): Promise { + if (typeof indexedDB === "undefined") return false; + // Avoid opening (and thereby creating) the DB when nothing has been cached. + try { + if (typeof indexedDB.databases === "function") { + const dbs = await indexedDB.databases(); + if (!dbs.some((d) => d.name === PARAKEET_CACHE_DB)) return false; + } + } catch { + // databases() can throw in private mode; fall through to open(). + } + + const repoId = PARAKEET_INFO.repoId; + // Hub keys: `hf-${repoId}-main--${filename}` (empty subfolder). + const candidates = [ + `hf-${repoId}-main--encoder-model.int8.onnx`, + `hf-${repoId}-main--encoder-model.fp16.onnx`, + ]; + try { + const db = await new Promise((resolve, reject) => { + const req = indexedDB.open(PARAKEET_CACHE_DB); + req.onerror = () => reject(req.error ?? new Error("IndexedDB open failed")); + req.onsuccess = () => resolve(req.result); + }); + if (!db.objectStoreNames.contains(PARAKEET_CACHE_STORE)) { + db.close(); + return false; + } + const hit = await new Promise((resolve, reject) => { + const tx = db.transaction([PARAKEET_CACHE_STORE], "readonly"); + const store = tx.objectStore(PARAKEET_CACHE_STORE); + let pending = candidates.length; + let found = false; + for (const key of candidates) { + const req = store.get(key); + req.onsuccess = () => { + const blob = req.result as Blob | undefined; + if (blob && blob.size > 1_000_000) found = true; + pending -= 1; + if (pending === 0) resolve(found); + }; + req.onerror = () => reject(req.error ?? new Error("IndexedDB get failed")); + } + }); + db.close(); + return hit; + } catch { + return false; + } +} + +/** + * Parakeet via parakeet.js — custom weightlift definition (not transformers.js). + * WebGPU uses fp16 encoder; WASM int8 is the size / compatibility fallback. + */ +function parakeetModel(): ModelDefinition { + return { + isCached: isParakeetCached, + load: async ({ progress }) => { + const { fromHub } = await import("parakeet.js"); + const onProgress = (p: { loaded: number; total: number; file: string }) => { + if (!p.file) return; + progress.dispatch({ + type: "progress", + file: p.file, + loaded: p.loaded, + ...(p.total > 0 ? { total: p.total } : {}), + }); + }; + const common = { + preprocessorBackend: "js" as const, + progress: onProgress, + wasmPaths: PARAKEET_ORT_WASM_PATHS, + }; + + // WebGPU cannot run the int8 encoder; fp16 (~1.2 GB) is the practical + // WebGPU path. WASM int8 (~670 MB) is the compatibility / size fallback. + const device = await fallbackDevicePolicy.pickDevice(); + if (device === "webgpu") { + try { + const model = await fromHub(PARAKEET_INFO.id, { + ...common, + backend: "webgpu", + encoderQuant: "fp16", + decoderQuant: "int8", + }); + asrDevice = "webgpu"; + return model as ParakeetInstance; + } catch (err) { + console.warn( + "Parakeet WebGPU/fp16 load failed; falling back to WASM int8.", + err + ); + fallbackDevicePolicy.preferWasm(); + } + } + + const model = await fromHub(PARAKEET_INFO.id, { + ...common, + backend: "wasm", + encoderQuant: "int8", + decoderQuant: "int8", + }); + asrDevice = "wasm"; + return model as ParakeetInstance; + }, + }; +} + /** - * ASR registry keyed by Hugging Face model id. Definitions are registered up - * front; getAsr() only loads by id. unloadAll() after a WebGPU loss forces a - * clean reload on WASM. + * ASR registry: Whisper ids are Hugging Face model ids; Parakeet uses + * PARAKEET_INFO.id. Definitions are registered up front; loaders only take an + * id. unloadAll() after a WebGPU loss forces a clean reload on WASM. */ const asrModels = new ModelManager({ - models: Object.fromEntries( - (Object.keys(MODELS) as WhisperModel[]).map((choice) => { - const { id, dtype } = MODELS[choice]; - return [ - id, - transformersModel({ - pipeline, - task: "automatic-speech-recognition", - modelId: id, - dtype, - cacheKey: env.cacheKey ?? "transformers-cache", - onDevice: (device) => { - asrDevice = device; - }, - }), - ]; - }) - ), + models: { + ...Object.fromEntries( + (Object.keys(MODELS) as WhisperModel[]).map((choice) => { + const { id, dtype } = MODELS[choice]; + return [ + id, + transformersModel({ + pipeline, + task: "automatic-speech-recognition", + modelId: id, + dtype, + cacheKey: env.cacheKey ?? "transformers-cache", + onDevice: (device) => { + asrDevice = device; + }, + }), + ]; + }) + ), + [PARAKEET_INFO.id]: parakeetModel(), + }, }); asrModels.subscribe((snap) => { const id = snap.loading[0]; @@ -121,12 +251,16 @@ async function getAsr(choice: WhisperModel) { return asrModels.load(MODELS[choice].id); } +async function getParakeet() { + return asrModels.load(PARAKEET_INFO.id); +} + /** - * Drop dead WebGPU pipelines and reload the requested model on WASM. + * Drop dead WebGPU pipelines and reload on WASM. * A lost GPU device invalidates every WebGPU session, so clear the whole - * ASR cache (same as the old asrPromises.clear()) — not just `choice`. + * ASR cache — not just the model that was running. */ -async function fallbackAsrToWasm(choice: WhisperModel) { +async function fallbackAsrToWasm() { fallbackDevicePolicy.preferWasm(); asrDevice = "wasm"; await asrModels.unloadAll(); @@ -135,7 +269,6 @@ async function fallbackAsrToWasm(choice: WhisperModel) { message: "GPU interrupted — continuing on CPU…", value: null, }); - return getAsr(choice); } /** @@ -409,106 +542,6 @@ function assignSpeakers(words: Word[], segments: DiarizationSegment[]) { } } -type ParakeetInstance = { - transcribe: ( - audio: Float32Array, - sampleRate?: number, - opts?: { - returnTimestamps?: boolean; - timeOffset?: number; - } - ) => Promise<{ - utterance_text: string; - words: Array<{ text: string; start_time: number; end_time: number }>; - }>; -}; - -let parakeetPromise: Promise | null = null; -let parakeetBackend: "webgpu" | "wasm" = "wasm"; - -function makeParakeetHubProgress(label: string) { - const files = new Map(); - let best = 0; - return (p: { loaded: number; total: number; file: string }) => { - if (!p.file || !(p.total > 0)) return; - files.set(p.file, { loaded: p.loaded, total: p.total }); - let loaded = 0; - let total = 0; - for (const f of files.values()) { - loaded += f.loaded; - total += f.total; - } - if (total <= 0) return; - best = Math.max(best, Math.min(1, loaded / total)); - post({ type: "progress", message: label, value: best }); - }; -} - -/** - * parakeet.js accepts `wasmPaths` on fromHub/fromUrls. A postinstall patch in - * scripts/copy-assets.mjs makes initOrt honor that argument (upstream 1.4.4 - * ignored it and defaulted to a jsDelivr CDN). - */ -async function getParakeet(): Promise { - if (!parakeetPromise) { - parakeetPromise = (async () => { - const { fromHub } = await import("parakeet.js"); - const device = await pickDevice(); - const progress = makeParakeetHubProgress("Downloading speech model…"); - const common = { - preprocessorBackend: "js" as const, - progress, - wasmPaths: PARAKEET_ORT_WASM_PATHS, - }; - - // WebGPU cannot run the int8 encoder; fp16 (~1.2 GB) is the practical - // WebGPU path. WASM int8 (~670 MB) is the compatibility / size fallback. - if (device === "webgpu") { - try { - post({ - type: "progress", - message: "Downloading speech model…", - value: null, - }); - const model = await fromHub(PARAKEET_INFO.id, { - ...common, - backend: "webgpu", - encoderQuant: "fp16", - decoderQuant: "int8", - }); - parakeetBackend = "webgpu"; - asrDevice = "webgpu"; - return model as ParakeetInstance; - } catch (err) { - console.warn( - "Parakeet WebGPU/fp16 load failed; falling back to WASM int8.", - err - ); - } - } - - post({ - type: "progress", - message: "Downloading speech model…", - value: null, - }); - const model = await fromHub(PARAKEET_INFO.id, { - ...common, - backend: "wasm", - encoderQuant: "int8", - decoderQuant: "int8", - }); - parakeetBackend = "wasm"; - asrDevice = "wasm"; - return model as ParakeetInstance; - })(); - parakeetPromise.catch(() => { - parakeetPromise = null; - }); - } - return parakeetPromise; -} - /** Map Parakeet word timestamps onto the original media timeline. */ function wordsFromParakeet( words: Array<{ text: string; start_time: number; end_time: number }>, @@ -571,7 +604,8 @@ async function runParakeet( duration: number ): Promise { getDiarizer().catch(() => {}); - const [model, vad] = await Promise.all([getParakeet(), getVad()]); + const [loaded, vad] = await Promise.all([getParakeet(), getVad()]); + let model = loaded; post({ type: "progress", message: "Detecting speech…", value: 0 }); const { segments: speechSegments } = await detectSpeechSegments(audio, vad); @@ -604,20 +638,16 @@ async function runParakeet( try { result = await runSlice(); } catch (err) { - if (parakeetBackend !== "webgpu" || !isWebGpuDeviceLostError(err)) { + if (asrDevice !== "webgpu" || !isWebGpuDeviceLostError(err)) { throw err; } console.warn( "WebGPU lost during Parakeet transcription; reloading on WASM.", err ); - forceWasm = true; - parakeetPromise = null; - const wasmModel = await getParakeet(); - result = await wasmModel.transcribe(slice, VAD_SAMPLE_RATE, { - returnTimestamps: true, - timeOffset: 0, - }); + await fallbackAsrToWasm(); + model = await getParakeet(); + result = await runSlice(); } rawWords.push( @@ -800,7 +830,8 @@ async function runWhisper( chunkFloor = progressBefore.chunkFloor; chunkTokens = progressBefore.chunkTokens; post({ type: "partial", text: partial }); - transcriber = await fallbackAsrToWasm(choice); + await fallbackAsrToWasm(); + transcriber = await getAsr(choice); chunks = await runSlice(); } From cf41d196cc17cbc877f45956684dc02241a87050 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 2 Aug 2026 11:07:00 +0000 Subject: [PATCH 3/5] Unify ASR models into a single MODELS registry Fold PARAKEET_INFO into MODELS with a backend discriminant so UI and the worker share one source for labels, sizes, and loader ids. Menus now render from ASR_ORDER. Co-authored-by: Wassim Gharbi --- components/Editor.tsx | 7 ++-- components/ModelSelector.tsx | 38 ++++++------------ components/UploadScreen.tsx | 7 ++-- lib/models.ts | 70 +++++++++++++++++++++------------ lib/store.ts | 1 - tests/models-test.ts | 21 ++++++++-- tests/vad-regression-test.ts | 6 ++- workers/transcription.worker.ts | 57 +++++++++++++-------------- 8 files changed, 114 insertions(+), 93 deletions(-) diff --git a/components/Editor.tsx b/components/Editor.tsx index ba9fc57..e371da9 100644 --- a/components/Editor.tsx +++ b/components/Editor.tsx @@ -23,6 +23,7 @@ import ModelSelector, { ModelOptionSeparator, } from "./ModelSelector"; import ImportTranscriptOption from "./ImportTranscriptOption"; +import { ASR_ORDER } from "@/lib/models"; /** How long the desktop mode-change overlay stays up. Matches the macOS * `setBounds(..., animate)` duration plus a small buffer so the layout @@ -230,9 +231,9 @@ export default function Editor() { <> {isElectron && - - - + {ASR_ORDER.map((id) => ( + + ))} diff --git a/components/ModelSelector.tsx b/components/ModelSelector.tsx index 19986c8..1155ef2 100644 --- a/components/ModelSelector.tsx +++ b/components/ModelSelector.tsx @@ -26,9 +26,9 @@ import { type TranscriptLanguage, } from "@/lib/languages"; import { + ASR_ORDER, MODELS, - PARAKEET_INFO, - isParakeetModel, + isAsrModel, isWhisperModel, type ModelChoice, } from "@/lib/models"; @@ -191,13 +191,11 @@ export default function ModelSelector({ activeTrigger?.icon ?? (model === "import" ? FileText : AudioLines); const baseTriggerLabel = activeTrigger?.label ?? - (isWhisperModel(model) + (isAsrModel(model) ? MODELS[model].label - : isParakeetModel(model) - ? PARAKEET_INFO.label - : model === "import" - ? "Import transcript" - : String(model)); + : model === "import" + ? "Import transcript" + : String(model)); const languageInfo = TRANSCRIPT_LANGUAGES[transcriptLanguage]; const showLanguageInTrigger = isWhisperModel(model) && @@ -207,9 +205,9 @@ export default function ModelSelector({ // Always mount options (hidden when closed) so custom triggers stay registered. const options = children ?? ( <> - - - + {ASR_ORDER.map((id) => ( + + ))} ); @@ -303,7 +301,7 @@ export default function ModelSelector({ ); } -/** Default option row: icon + label + optional meta. Whisper ids fill in from MODELS. */ +/** Default option row: icon + label + optional meta. ASR ids fill in from MODELS. */ export function ModelOption({ id, label, @@ -325,20 +323,8 @@ export function ModelOption({ const selector = useSelectorCtx(); const selected = selector.value === id; - const resolvedLabel = - label ?? - (isWhisperModel(id) - ? MODELS[id].label - : isParakeetModel(id) - ? PARAKEET_INFO.label - : id); - const resolvedMeta = - meta ?? - (isWhisperModel(id) - ? MODELS[id].size - : isParakeetModel(id) - ? PARAKEET_INFO.size - : undefined); + const resolvedLabel = label ?? (isAsrModel(id) ? MODELS[id].label : id); + const resolvedMeta = meta ?? (isAsrModel(id) ? MODELS[id].size : undefined); const optionCtx = useMemo( () => ({ diff --git a/components/UploadScreen.tsx b/components/UploadScreen.tsx index 516c44a..8655a85 100644 --- a/components/UploadScreen.tsx +++ b/components/UploadScreen.tsx @@ -23,6 +23,7 @@ import ModelSelector, { ModelOptionSeparator, } from "./ModelSelector"; import ImportTranscriptOption from "./ImportTranscriptOption"; +import { ASR_ORDER } from "@/lib/models"; import { useCrossOriginIsolated } from "@/hooks/useCrossOriginIsolated"; import { detectMediaKind, MEDIA_ACCEPT } from "@/lib/media"; import { formatTime } from "@/lib/edits"; @@ -275,9 +276,9 @@ export default function UploadScreen({
- - - + {ASR_ORDER.map((id) => ( + + ))} diff --git a/lib/models.ts b/lib/models.ts index b41ca41..56762d4 100644 --- a/lib/models.ts +++ b/lib/models.ts @@ -7,13 +7,18 @@ export type ModelChoice = AsrModel | "import"; type DType = "fp32" | "fp16" | "q8" | "int8" | "uint8" | "q4" | "q4f16" | "bnb4"; -export interface ModelInfo { - /** Hugging Face model id (ONNX export compatible with transformers.js). */ - id: string; +/** Shared UI fields for every local ASR backend. */ +type ModelDisplay = { label: string; description: string; /** Approximate download size shown in the UI. */ size: string; +}; + +export type WhisperModelInfo = ModelDisplay & { + backend: "whisper"; + /** Hugging Face model id (ONNX export compatible with transformers.js). */ + id: string; /** dtype configuration per device. */ dtype: { webgpu: Record; @@ -30,21 +35,40 @@ export interface ModelInfo { * required for word-level timestamps, which this editor depends on.) */ verbatimPrompt?: string; -} +}; + +export type ParakeetModelInfo = ModelDisplay & { + backend: "parakeet"; + /** parakeet.js model key (also the weightlift registry id). */ + id: string; + /** Hugging Face repo used by parakeet.js hub downloads / IndexedDB cache keys. */ + repoId: string; +}; + +export type ModelInfo = WhisperModelInfo | ParakeetModelInfo; -/** Display order for Whisper rows in the homepage source dropdown. */ -export const WHISPER_ORDER: WhisperModel[] = ["base", "small"]; +/** Display order for ASR rows in the source dropdown. */ +export const ASR_ORDER: AsrModel[] = ["base", "small", "parakeet"]; const WHISPER_DTYPE = { // q4 decoder: q8 fails session creation on onnxruntime-web 1.26 // (Missing required scale … MatMulNBits). webgpu: { encoder_model: "fp32", decoder_model_merged: "q4" }, wasm: { encoder_model: "fp32", decoder_model_merged: "q4" }, -} satisfies ModelInfo["dtype"]; +} satisfies WhisperModelInfo["dtype"]; -/** Whisper models that can run in the transcription worker. */ -export const MODELS: Record = { +/** + * Local ASR models that can run in the transcription worker. + * Shared display fields live on every entry; backend-specific knobs + * (`dtype` / `verbatimPrompt` vs `repoId`) are gated by `backend`. + */ +export const MODELS: { + base: WhisperModelInfo; + small: WhisperModelInfo; + parakeet: ParakeetModelInfo; +} = { base: { + backend: "whisper", id: "onnx-community/whisper-base_timestamped", label: "Whisper Base", description: "Faster download and transcription. Good for most clips.", @@ -55,29 +79,25 @@ export const MODELS: Record = { // speaker on mixed clips). Prefer post-process / filler tools instead. }, small: { + backend: "whisper", id: "onnx-community/whisper-small_timestamped", label: "Whisper Small", description: "More accurate on longer or noisier audio. Larger download.", size: "~600 MB", dtype: WHISPER_DTYPE, }, + parakeet: { + backend: "parakeet", + id: "parakeet-tdt-0.6b-v3", + repoId: "ysdede/parakeet-tdt-0.6b-v3-onnx", + label: "Parakeet TDT v3", + description: + "NVIDIA FastConformer — faster on WebGPU, strong EU-language accuracy. Auto-detects language.", + // WASM int8 ~670 MB; WebGPU fp16 ~1.2 GB. + size: "~700 MB", + }, }; -/** - * UI metadata for Parakeet (not a transformers.js Whisper checkpoint). - * Download size depends on backend: WASM int8 ~670 MB, WebGPU fp16 ~1.2 GB. - */ -export const PARAKEET_INFO = { - /** parakeet.js model key (weightlift registry id). */ - id: "parakeet-tdt-0.6b-v3", - /** Hugging Face repo used by parakeet.js hub downloads / IndexedDB cache keys. */ - repoId: "ysdede/parakeet-tdt-0.6b-v3-onnx", - label: "Parakeet TDT v3", - description: - "NVIDIA FastConformer — faster on WebGPU, strong EU-language accuracy. Auto-detects language.", - size: "~700 MB", -} as const; - export function isWhisperModel(value: unknown): value is WhisperModel { return value === "base" || value === "small"; } @@ -88,7 +108,7 @@ export function isParakeetModel(value: unknown): value is ParakeetModel { /** Models that run local ASR in the transcription worker (not import). */ export function isAsrModel(value: unknown): value is AsrModel { - return isWhisperModel(value) || isParakeetModel(value); + return typeof value === "string" && Object.prototype.hasOwnProperty.call(MODELS, value); } export function isModelChoice(value: unknown): value is ModelChoice { diff --git a/lib/store.ts b/lib/store.ts index 844a772..9e0a9be 100644 --- a/lib/store.ts +++ b/lib/store.ts @@ -27,7 +27,6 @@ import { import { isAsrModel, isModelChoice, - isWhisperModel, loadModelPreference, saveModelPreference, } from "./models"; diff --git a/tests/models-test.ts b/tests/models-test.ts index 183690e..72f30b5 100644 --- a/tests/models-test.ts +++ b/tests/models-test.ts @@ -2,8 +2,8 @@ * Model choice helpers for Whisper / Parakeet / import. */ import { + ASR_ORDER, MODELS, - PARAKEET_INFO, isAsrModel, isModelChoice, isParakeetModel, @@ -25,13 +25,26 @@ assert(isModelChoice("import"), "import is a model choice"); assert(isModelChoice("parakeet"), "parakeet is a model choice"); assert(!isModelChoice("tiny"), "tiny is not a model choice"); -assert(PARAKEET_INFO.id === "parakeet-tdt-0.6b-v3", "parakeet hub id"); +assert(MODELS.parakeet.backend === "parakeet", "parakeet backend"); +assert(MODELS.parakeet.id === "parakeet-tdt-0.6b-v3", "parakeet hub id"); assert( - PARAKEET_INFO.repoId === "ysdede/parakeet-tdt-0.6b-v3-onnx", + MODELS.parakeet.repoId === "ysdede/parakeet-tdt-0.6b-v3-onnx", "parakeet HF repo id" ); -assert(typeof PARAKEET_INFO.label === "string", "parakeet label"); +assert(typeof MODELS.parakeet.label === "string", "parakeet label"); +assert(MODELS.base.backend === "whisper", "base backend"); assert(typeof MODELS.base.id === "string", "whisper base id"); assert(typeof MODELS.small.id === "string", "whisper small id"); +assert(MODELS.base.dtype.webgpu.encoder_model === "fp32", "whisper dtype"); + +assert( + ASR_ORDER.includes("parakeet") && ASR_ORDER.includes("base"), + "ASR_ORDER lists whisper + parakeet" +); +for (const id of ASR_ORDER) { + assert(isAsrModel(id), `${id} in ASR_ORDER is an ASR model`); + assert(typeof MODELS[id].label === "string", `${id} has label`); + assert(typeof MODELS[id].size === "string", `${id} has size`); +} console.log("models-test: ok"); diff --git a/tests/vad-regression-test.ts b/tests/vad-regression-test.ts index a64a901..e554a45 100644 --- a/tests/vad-regression-test.ts +++ b/tests/vad-regression-test.ts @@ -30,10 +30,12 @@ function ffmpegAvailable(): boolean { return probe.status === 0; } -// Long decoder prompts truncate long-form ASR — models must not set them. +// Long decoder prompts truncate long-form ASR — Whisper models must not set them. for (const id of Object.keys(MODELS) as (keyof typeof MODELS)[]) { + const info = MODELS[id]; + if (info.backend !== "whisper") continue; assert( - !MODELS[id].verbatimPrompt, + !info.verbatimPrompt, `${id} must not set verbatimPrompt (truncates multi-speaker clips)` ); } diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 3649ef3..22ca4d3 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -30,7 +30,6 @@ import { import type { Word, WorkerRequest, WorkerResponse } from "@/lib/types"; import { MODELS, - PARAKEET_INFO, isParakeetModel, isWhisperModel, type AsrModel, @@ -107,7 +106,7 @@ async function isParakeetCached(): Promise { // databases() can throw in private mode; fall through to open(). } - const repoId = PARAKEET_INFO.repoId; + const repoId = MODELS.parakeet.repoId; // Hub keys: `hf-${repoId}-main--${filename}` (empty subfolder). const candidates = [ `hf-${repoId}-main--encoder-model.int8.onnx`, @@ -175,7 +174,7 @@ function parakeetModel(): ModelDefinition { const device = await fallbackDevicePolicy.pickDevice(); if (device === "webgpu") { try { - const model = await fromHub(PARAKEET_INFO.id, { + const model = await fromHub(MODELS.parakeet.id, { ...common, backend: "webgpu", encoderQuant: "fp16", @@ -192,7 +191,7 @@ function parakeetModel(): ModelDefinition { } } - const model = await fromHub(PARAKEET_INFO.id, { + const model = await fromHub(MODELS.parakeet.id, { ...common, backend: "wasm", encoderQuant: "int8", @@ -205,32 +204,32 @@ function parakeetModel(): ModelDefinition { } /** - * ASR registry: Whisper ids are Hugging Face model ids; Parakeet uses - * PARAKEET_INFO.id. Definitions are registered up front; loaders only take an - * id. unloadAll() after a WebGPU loss forces a clean reload on WASM. + * ASR registry keyed by each model's `id` from MODELS. Definitions are + * registered up front; loaders only take an id. unloadAll() after a WebGPU + * loss forces a clean reload on WASM. */ const asrModels = new ModelManager({ - models: { - ...Object.fromEntries( - (Object.keys(MODELS) as WhisperModel[]).map((choice) => { - const { id, dtype } = MODELS[choice]; - return [ - id, - transformersModel({ - pipeline, - task: "automatic-speech-recognition", - modelId: id, - dtype, - cacheKey: env.cacheKey ?? "transformers-cache", - onDevice: (device) => { - asrDevice = device; - }, - }), - ]; - }) - ), - [PARAKEET_INFO.id]: parakeetModel(), - }, + models: Object.fromEntries( + (Object.keys(MODELS) as AsrModel[]).map((choice) => { + const info = MODELS[choice]; + if (info.backend === "parakeet") { + return [info.id, parakeetModel()]; + } + return [ + info.id, + transformersModel({ + pipeline, + task: "automatic-speech-recognition", + modelId: info.id, + dtype: info.dtype, + cacheKey: env.cacheKey ?? "transformers-cache", + onDevice: (device) => { + asrDevice = device; + }, + }), + ]; + }) + ), }); asrModels.subscribe((snap) => { const id = snap.loading[0]; @@ -252,7 +251,7 @@ async function getAsr(choice: WhisperModel) { } async function getParakeet() { - return asrModels.load(PARAKEET_INFO.id); + return asrModels.load(MODELS.parakeet.id); } /** From 7a04394dd5f9526de9eef3ed742b9f804fd56df6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 2 Aug 2026 11:11:50 +0000 Subject: [PATCH 4/5] Rename asrModels to models and ASR_ORDER to MODEL_ORDER Co-authored-by: Wassim Gharbi --- components/Editor.tsx | 4 ++-- components/ModelSelector.tsx | 4 ++-- components/UploadScreen.tsx | 4 ++-- lib/models.ts | 4 ++-- tests/models-test.ts | 10 +++++----- workers/transcription.worker.ts | 10 +++++----- 6 files changed, 18 insertions(+), 18 deletions(-) diff --git a/components/Editor.tsx b/components/Editor.tsx index e371da9..ea57f58 100644 --- a/components/Editor.tsx +++ b/components/Editor.tsx @@ -23,7 +23,7 @@ import ModelSelector, { ModelOptionSeparator, } from "./ModelSelector"; import ImportTranscriptOption from "./ImportTranscriptOption"; -import { ASR_ORDER } from "@/lib/models"; +import { MODEL_ORDER } from "@/lib/models"; /** How long the desktop mode-change overlay stays up. Matches the macOS * `setBounds(..., animate)` duration plus a small buffer so the layout @@ -231,7 +231,7 @@ export default function Editor() { <> {isElectron && - {ASR_ORDER.map((id) => ( + {MODEL_ORDER.map((id) => ( ))} diff --git a/components/ModelSelector.tsx b/components/ModelSelector.tsx index 1155ef2..b0a108f 100644 --- a/components/ModelSelector.tsx +++ b/components/ModelSelector.tsx @@ -26,7 +26,7 @@ import { type TranscriptLanguage, } from "@/lib/languages"; import { - ASR_ORDER, + MODEL_ORDER, MODELS, isAsrModel, isWhisperModel, @@ -205,7 +205,7 @@ export default function ModelSelector({ // Always mount options (hidden when closed) so custom triggers stay registered. const options = children ?? ( <> - {ASR_ORDER.map((id) => ( + {MODEL_ORDER.map((id) => ( ))} diff --git a/components/UploadScreen.tsx b/components/UploadScreen.tsx index 8655a85..adc91b1 100644 --- a/components/UploadScreen.tsx +++ b/components/UploadScreen.tsx @@ -23,7 +23,7 @@ import ModelSelector, { ModelOptionSeparator, } from "./ModelSelector"; import ImportTranscriptOption from "./ImportTranscriptOption"; -import { ASR_ORDER } from "@/lib/models"; +import { MODEL_ORDER } from "@/lib/models"; import { useCrossOriginIsolated } from "@/hooks/useCrossOriginIsolated"; import { detectMediaKind, MEDIA_ACCEPT } from "@/lib/media"; import { formatTime } from "@/lib/edits"; @@ -276,7 +276,7 @@ export default function UploadScreen({
- {ASR_ORDER.map((id) => ( + {MODEL_ORDER.map((id) => ( ))} diff --git a/lib/models.ts b/lib/models.ts index 56762d4..363c380 100644 --- a/lib/models.ts +++ b/lib/models.ts @@ -47,8 +47,8 @@ export type ParakeetModelInfo = ModelDisplay & { export type ModelInfo = WhisperModelInfo | ParakeetModelInfo; -/** Display order for ASR rows in the source dropdown. */ -export const ASR_ORDER: AsrModel[] = ["base", "small", "parakeet"]; +/** Display order for model rows in the source dropdown. */ +export const MODEL_ORDER: AsrModel[] = ["base", "small", "parakeet"]; const WHISPER_DTYPE = { // q4 decoder: q8 fails session creation on onnxruntime-web 1.26 diff --git a/tests/models-test.ts b/tests/models-test.ts index 72f30b5..71b4c7e 100644 --- a/tests/models-test.ts +++ b/tests/models-test.ts @@ -2,7 +2,7 @@ * Model choice helpers for Whisper / Parakeet / import. */ import { - ASR_ORDER, + MODEL_ORDER, MODELS, isAsrModel, isModelChoice, @@ -38,11 +38,11 @@ assert(typeof MODELS.small.id === "string", "whisper small id"); assert(MODELS.base.dtype.webgpu.encoder_model === "fp32", "whisper dtype"); assert( - ASR_ORDER.includes("parakeet") && ASR_ORDER.includes("base"), - "ASR_ORDER lists whisper + parakeet" + MODEL_ORDER.includes("parakeet") && MODEL_ORDER.includes("base"), + "MODEL_ORDER lists whisper + parakeet" ); -for (const id of ASR_ORDER) { - assert(isAsrModel(id), `${id} in ASR_ORDER is an ASR model`); +for (const id of MODEL_ORDER) { + assert(isAsrModel(id), `${id} in MODEL_ORDER is an ASR model`); assert(typeof MODELS[id].label === "string", `${id} has label`); assert(typeof MODELS[id].size === "string", `${id} has size`); } diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 22ca4d3..4e0d9a7 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -208,7 +208,7 @@ function parakeetModel(): ModelDefinition { * registered up front; loaders only take an id. unloadAll() after a WebGPU * loss forces a clean reload on WASM. */ -const asrModels = new ModelManager({ +const models = new ModelManager({ models: Object.fromEntries( (Object.keys(MODELS) as AsrModel[]).map((choice) => { const info = MODELS[choice]; @@ -231,7 +231,7 @@ const asrModels = new ModelManager({ }) ), }); -asrModels.subscribe((snap) => { +models.subscribe((snap) => { const id = snap.loading[0]; if (!id) return; const rec = snap.models[id]; @@ -247,11 +247,11 @@ asrModels.subscribe((snap) => { }); async function getAsr(choice: WhisperModel) { - return asrModels.load(MODELS[choice].id); + return models.load(MODELS[choice].id); } async function getParakeet() { - return asrModels.load(MODELS.parakeet.id); + return models.load(MODELS.parakeet.id); } /** @@ -262,7 +262,7 @@ async function getParakeet() { async function fallbackAsrToWasm() { fallbackDevicePolicy.preferWasm(); asrDevice = "wasm"; - await asrModels.unloadAll(); + await models.unloadAll(); post({ type: "progress", message: "GPU interrupted — continuing on CPU…", From d54ba6984280ff57646a98a5c3dcfa847d7d8c15 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 2 Aug 2026 11:17:51 +0000 Subject: [PATCH 5/5] Split transcript source from speech models Keep MODELS as real speech backends only. Move "import" onto TranscriptSource (store.source) so model helpers no longer have to special-case caption import. Co-authored-by: Wassim Gharbi --- components/ImportTranscriptOption.tsx | 83 ++++++++++++++------------- components/ModelSelector.tsx | 42 +++++++------- components/UploadScreen.tsx | 7 +-- hooks/useTranscriber.ts | 8 +-- lib/autosave.ts | 2 +- lib/models.ts | 31 ++++------ lib/projects.ts | 25 +++++--- lib/source.ts | 11 ++++ lib/store.ts | 48 +++++++--------- lib/types.ts | 4 +- tests/models-test.ts | 20 +++---- workers/transcription.worker.ts | 6 +- 12 files changed, 149 insertions(+), 138 deletions(-) create mode 100644 lib/source.ts diff --git a/components/ImportTranscriptOption.tsx b/components/ImportTranscriptOption.tsx index da121b5..b403360 100644 --- a/components/ImportTranscriptOption.tsx +++ b/components/ImportTranscriptOption.tsx @@ -8,7 +8,7 @@ import { TRANSCRIPT_FILE_ERROR, TRANSCRIPT_ACCEPT, } from "@/lib/parseTranscript"; -import { isAsrModel, type AsrModel } from "@/lib/models"; +import { isModelId, type ModelId } from "@/lib/models"; import { useEditorStore } from "@/lib/store"; import { ModelOption, @@ -30,8 +30,8 @@ import { export default function ImportTranscriptOption() { const pendingTranscript = useEditorStore((s) => s.pendingTranscript); const setPendingTranscript = useEditorStore((s) => s.setPendingTranscript); - const setModel = useEditorStore((s) => s.setModel); - const selected = useEditorStore((s) => s.model === "import"); + const setSource = useEditorStore((s) => s.setSource); + const selected = useEditorStore((s) => s.source === "import"); const [reading, setReading] = useState(false); const [error, setError] = useState(null); const [picking, setPicking] = useState(false); @@ -40,7 +40,7 @@ export default function ImportTranscriptOption() { ModelOptionContextValue, "keepMenuOpen" | "closeMenu" | "select" > | null>(null); - const previousModelRef = useRef("base"); + const previousModelRef = useRef("base"); const pickGenRef = useRef(0); /** Reset import-pick state only — never touch the dropdown open state. */ @@ -49,15 +49,15 @@ export default function ImportTranscriptOption() { setReading(false); setError(null); if (!useEditorStore.getState().pendingTranscript) { - setModel(previousModelRef.current); + setSource(previousModelRef.current); } - }, [setModel]); + }, [setSource]); - // If the user switches to an ASR model while a picker/parse is in flight, - // invalidate so a late onChange/parse cannot flip model back to import. + // If the user switches to a speech model while a picker/parse is in flight, + // invalidate so a late onChange/parse cannot flip source back to import. useEffect(() => { return useEditorStore.subscribe((state, prev) => { - if (!isAsrModel(state.model) || state.model === prev.model) return; + if (!isModelId(state.source) || state.source === prev.source) return; pickGenRef.current += 1; queueMicrotask(() => { setPicking(false); @@ -131,19 +131,19 @@ export default function ImportTranscriptOption() { setPicking(false); setError(TRANSCRIPT_FILE_ERROR); setPendingTranscript(null); - setModel("import"); + setSource("import"); menu?.keepMenuOpen(); return; } setReading(true); setPicking(false); setError(null); - setModel("import"); // so the closed trigger can show progress + setSource("import"); // so the closed trigger can show progress try { const words = await parseTranscriptFile(file); if (pickGenRef.current !== gen) return; setPendingTranscript({ name: file.name, words }); - setModel("import"); + setSource("import"); menu?.closeMenu(); } catch (err) { if (pickGenRef.current !== gen) return; @@ -154,7 +154,7 @@ export default function ImportTranscriptOption() { ? err.message : "Could not read that transcript." ); - setModel("import"); + setSource("import"); menu?.keepMenuOpen(); } finally { if (pickGenRef.current === gen) setReading(false); @@ -170,11 +170,11 @@ export default function ImportTranscriptOption() { autoTrigger={false} onSelect={(ctx) => { menuRef.current = ctx; - const current = useEditorStore.getState().model; - if (isAsrModel(current)) { + const current = useEditorStore.getState().source; + if (isModelId(current)) { previousModelRef.current = current; } - // Do not set model to "import" until a file is chosen. Close the menu + // Do not set source to "import" until a file is chosen. Close the menu // before the OS dialog so cancel cannot leave it pinned open. pickGenRef.current += 1; setPicking(true); @@ -210,8 +210,12 @@ function ImportTrigger({ "import", { label, - icon: FileText, - iconClassName: error ? "text-red-500" : "text-zinc-500 dark:text-zinc-400", + icon: busy ? Loader2 : FileText, + iconClassName: busy + ? "animate-spin text-zinc-500" + : error + ? "text-red-500" + : "text-zinc-500", busy, }, enabled @@ -229,27 +233,24 @@ function ImportStatus({ picking: boolean; }) { const { selected } = useModelOption(); - const pendingTranscript = useEditorStore((s) => s.pendingTranscript); - if (!picking && !selected && !reading && !error) return null; - if (!reading && !pendingTranscript && !error && !picking) return null; - return ( - - {reading ? ( - - - Reading file… - - ) : error ? ( - error - ) : pendingTranscript ? ( - `${pendingTranscript.name} · ${pendingTranscript.words.length} words` - ) : picking ? ( - "Choose an SRT, VTT, or JSON file…" - ) : null} - - ); + if (reading) { + return ( +

+ Reading file… +

+ ); + } + if (error) { + return ( +

{error}

+ ); + } + if (picking && !selected) { + return ( +

+ Choose a file… +

+ ); + } + return null; } diff --git a/components/ModelSelector.tsx b/components/ModelSelector.tsx index b0a108f..11c6409 100644 --- a/components/ModelSelector.tsx +++ b/components/ModelSelector.tsx @@ -28,10 +28,10 @@ import { import { MODEL_ORDER, MODELS, - isAsrModel, + isModelId, isWhisperModel, - type ModelChoice, } from "@/lib/models"; +import type { TranscriptSource } from "@/lib/source"; import { hydrateModelPreference, hydrateTranscriptLanguagePreference, @@ -41,7 +41,7 @@ import Popover, { PopoverContent, PopoverTrigger } from "./Popover"; export type ModelOptionContextValue = { /** Currently selected source id. */ - value: ModelChoice; + value: TranscriptSource; selected: boolean; select: () => void; /** Close the dropdown after a normal selection. */ @@ -59,8 +59,8 @@ export type OptionTrigger = { }; type SelectorContextValue = { - value: ModelChoice; - setValue: (id: ModelChoice) => void; + value: TranscriptSource; + setValue: (id: TranscriptSource) => void; closeMenu: () => void; keepMenuOpen: () => void; registerTrigger: (id: string, trigger: OptionTrigger) => void; @@ -88,7 +88,7 @@ export function useModelOption(): ModelOptionContextValue { /** Let a custom option drive the closed trigger while it is selected. */ export function useOptionTrigger( - id: ModelChoice, + id: TranscriptSource, trigger: OptionTrigger, enabled = true ) { @@ -126,8 +126,8 @@ export default function ModelSelector({ /** Called when an option needs the parent panel to stay open (embedded). */ onKeepOpen?: () => void; }) { - const model = useEditorStore((s) => s.model); - const setModel = useEditorStore((s) => s.setModel); + const source = useEditorStore((s) => s.source); + const setSource = useEditorStore((s) => s.setSource); const transcriptLanguage = useEditorStore((s) => s.transcriptLanguage); const [open, setOpen] = useState(false); const [triggers, setTriggers] = useState>({}); @@ -174,31 +174,31 @@ export default function ModelSelector({ const ctx = useMemo( () => ({ - value: model, - setValue: setModel, + value: source, + setValue: setSource, closeMenu, keepMenuOpen, registerTrigger, unregisterTrigger, }), - [model, setModel, closeMenu, keepMenuOpen, registerTrigger, unregisterTrigger] + [source, setSource, closeMenu, keepMenuOpen, registerTrigger, unregisterTrigger] ); - const activeTrigger = triggers[model]; + const activeTrigger = triggers[source]; // Prefer the option's registered trigger. Fall back carefully so an unmounted // custom option (e.g. import) never shows the raw id + default wave icon. const TriggerIcon = - activeTrigger?.icon ?? (model === "import" ? FileText : AudioLines); + activeTrigger?.icon ?? (source === "import" ? FileText : AudioLines); const baseTriggerLabel = activeTrigger?.label ?? - (isAsrModel(model) - ? MODELS[model].label - : model === "import" + (isModelId(source) + ? MODELS[source].label + : source === "import" ? "Import transcript" - : String(model)); + : String(source)); const languageInfo = TRANSCRIPT_LANGUAGES[transcriptLanguage]; const showLanguageInTrigger = - isWhisperModel(model) && + isWhisperModel(source) && !activeTrigger?.busy && transcriptLanguage !== "en"; @@ -312,7 +312,7 @@ export function ModelOption({ /** When false, a child owns the closed trigger via `useOptionTrigger`. */ autoTrigger = true, }: { - id: ModelChoice; + id: TranscriptSource; label?: string; meta?: string; icon?: LucideIcon; @@ -323,8 +323,8 @@ export function ModelOption({ const selector = useSelectorCtx(); const selected = selector.value === id; - const resolvedLabel = label ?? (isAsrModel(id) ? MODELS[id].label : id); - const resolvedMeta = meta ?? (isAsrModel(id) ? MODELS[id].size : undefined); + const resolvedLabel = label ?? (isModelId(id) ? MODELS[id].label : id); + const resolvedMeta = meta ?? (isModelId(id) ? MODELS[id].size : undefined); const optionCtx = useMemo( () => ({ diff --git a/components/UploadScreen.tsx b/components/UploadScreen.tsx index adc91b1..b2d8888 100644 --- a/components/UploadScreen.tsx +++ b/components/UploadScreen.tsx @@ -167,7 +167,7 @@ export default function UploadScreen({ // would fail immediately and lose the file to that reload. const isolation = useCrossOriginIsolated(); const ready = isolation === "ready"; - const model = useEditorStore((s) => s.model); + const source = useEditorStore((s) => s.source); const pendingTranscript = useEditorStore((s) => s.pendingTranscript); const openProject = useEditorStore((s) => s.openProject); const removeProject = useEditorStore((s) => s.removeProject); @@ -206,8 +206,7 @@ export default function UploadScreen({ alert("Please choose a video or audio file."); return; } - const { model: source, pendingTranscript: pending } = - useEditorStore.getState(); + const { source, pendingTranscript: pending } = useEditorStore.getState(); if (source === "import") { if (!pending) { alert("Choose a transcript file from the source menu first."); @@ -343,7 +342,7 @@ export default function UploadScreen({ browse

- {model === "import" + {source === "import" ? pendingTranscript ? `Will use ${pendingTranscript.name} · MP4, WebM, MOV, MP3, WAV, …` : "Pick a transcript in the menu above, then drop your media" diff --git a/hooks/useTranscriber.ts b/hooks/useTranscriber.ts index c8a696b..7c1ae43 100644 --- a/hooks/useTranscriber.ts +++ b/hooks/useTranscriber.ts @@ -1,7 +1,7 @@ "use client"; import { useCallback, useEffect, useRef } from "react"; -import { isAsrModel } from "@/lib/models"; +import { isModelId } from "@/lib/models"; import { useEditorStore } from "@/lib/store"; import type { WorkerResponse } from "@/lib/types"; @@ -26,11 +26,11 @@ export function useTranscriber() { const transcribe = useCallback((audio: Float32Array, duration: number) => { const store = useEditorStore.getState(); - if (!isAsrModel(store.model)) { + if (!isModelId(store.source)) { store.setError("Select a speech model to transcribe."); return; } - const asrModel = store.model; + const model = store.source; const transcriptLanguage = store.transcriptLanguage; store.setStatus("transcribing"); store.setProgress({ message: "Loading speech model…", value: null }); @@ -73,7 +73,7 @@ export function useTranscriber() { // Transfer a copy so the original stays available for the waveform. const copy = audio.slice(); workerRef.current.postMessage( - { audio: copy, duration, model: asrModel, language: transcriptLanguage }, + { audio: copy, duration, model, language: transcriptLanguage }, [copy.buffer] ); }, []); diff --git a/lib/autosave.ts b/lib/autosave.ts index 40f73e6..c3fb400 100644 --- a/lib/autosave.ts +++ b/lib/autosave.ts @@ -65,7 +65,7 @@ async function writeSnapshot() { name: s.videoFile.name, mediaKind: s.mediaKind, duration: s.duration, - model: s.model, + source: s.source, transcriptLanguage: s.transcriptLanguage, words: s.words, showDeleted: s.showDeleted, diff --git a/lib/models.ts b/lib/models.ts index 363c380..53b8eaa 100644 --- a/lib/models.ts +++ b/lib/models.ts @@ -1,13 +1,12 @@ -/** Transcription source choices offered on the upload screen. */ +/** Local speech models offered on the upload screen. */ export type WhisperModel = "base" | "small"; /** NVIDIA Parakeet TDT 0.6B v3 via parakeet.js (ONNX / WebGPU). */ export type ParakeetModel = "parakeet"; -export type AsrModel = WhisperModel | ParakeetModel; -export type ModelChoice = AsrModel | "import"; +export type ModelId = WhisperModel | ParakeetModel; type DType = "fp32" | "fp16" | "q8" | "int8" | "uint8" | "q4" | "q4f16" | "bnb4"; -/** Shared UI fields for every local ASR backend. */ +/** Shared UI fields for every local speech backend. */ type ModelDisplay = { label: string; description: string; @@ -48,7 +47,7 @@ export type ParakeetModelInfo = ModelDisplay & { export type ModelInfo = WhisperModelInfo | ParakeetModelInfo; /** Display order for model rows in the source dropdown. */ -export const MODEL_ORDER: AsrModel[] = ["base", "small", "parakeet"]; +export const MODEL_ORDER: ModelId[] = ["base", "small", "parakeet"]; const WHISPER_DTYPE = { // q4 decoder: q8 fails session creation on onnxruntime-web 1.26 @@ -58,7 +57,7 @@ const WHISPER_DTYPE = { } satisfies WhisperModelInfo["dtype"]; /** - * Local ASR models that can run in the transcription worker. + * Local speech models that can run in the transcription worker. * Shared display fields live on every entry; backend-specific knobs * (`dtype` / `verbatimPrompt` vs `repoId`) are gated by `backend`. */ @@ -106,33 +105,27 @@ export function isParakeetModel(value: unknown): value is ParakeetModel { return value === "parakeet"; } -/** Models that run local ASR in the transcription worker (not import). */ -export function isAsrModel(value: unknown): value is AsrModel { +/** Whether `value` is a key of {@link MODELS}. */ +export function isModelId(value: unknown): value is ModelId { return typeof value === "string" && Object.prototype.hasOwnProperty.call(MODELS, value); } -export function isModelChoice(value: unknown): value is ModelChoice { - return isAsrModel(value) || value === "import"; -} - const MODEL_STORAGE_KEY = "rescript.model"; -/** Read the last-selected ASR model from localStorage (defaults to base). */ -export function loadModelPreference(): AsrModel { +/** Read the last-selected speech model from localStorage (defaults to base). */ +export function loadModelPreference(): ModelId { if (typeof window === "undefined") return "base"; try { const raw = window.localStorage.getItem(MODEL_STORAGE_KEY); - // Ignore a stale "import" preference — that choice is session-only until a - // transcript file is picked again. - if (isAsrModel(raw)) return raw; + if (isModelId(raw)) return raw; } catch { // private mode / disabled storage } return "base"; } -/** Persist the selected ASR model for the next visit. */ -export function saveModelPreference(model: AsrModel) { +/** Persist the selected speech model for the next visit. */ +export function saveModelPreference(model: ModelId) { if (typeof window === "undefined") return; try { window.localStorage.setItem(MODEL_STORAGE_KEY, model); diff --git a/lib/projects.ts b/lib/projects.ts index 9630c46..1e0c986 100644 --- a/lib/projects.ts +++ b/lib/projects.ts @@ -6,8 +6,7 @@ * (oldest by updatedAt are pruned). */ -import type { ModelChoice } from "./models"; -import { isModelChoice } from "./models"; +import { isTranscriptSource, type TranscriptSource } from "./source"; import type { TranscriptLanguage } from "./languages"; import { DEFAULT_TRANSCRIPT_LANGUAGE, @@ -26,12 +25,21 @@ export interface ProjectMeta { name: string; mediaKind: MediaKind; duration: number; - model: ModelChoice; + source: TranscriptSource; transcriptLanguage: TranscriptLanguage; updatedAt: number; createdAt: number; } +/** Read source from a stored row; older saves used `model`. */ +function projectSource(row: { + source?: unknown; + model?: unknown; +}): TranscriptSource { + const raw = row.source ?? row.model; + return isTranscriptSource(raw) ? raw : "base"; +} + export interface ProjectRecord extends ProjectMeta { words: Word[]; showDeleted: boolean; @@ -126,7 +134,7 @@ export async function listProjects(): Promise { name: r.name, mediaKind: r.mediaKind, duration: r.duration, - model: r.model, + source: projectSource(r), transcriptLanguage: isTranscriptLanguage(r.transcriptLanguage) ? r.transcriptLanguage : DEFAULT_TRANSCRIPT_LANGUAGE, @@ -140,10 +148,13 @@ export async function getProject(id: string): Promise { const db = await openDb(); const tx = db.transaction(STORE, "readonly"); const row = await idbReq( - tx.objectStore(STORE).get(id) as IDBRequest + tx.objectStore(STORE).get(id) as IDBRequest< + (ProjectRecord & { model?: unknown }) | undefined + > ); await txDone(tx); - return row ?? null; + if (!row) return null; + return { ...row, source: projectSource(row) }; } /** Insert or replace a project, then prune to MAX_PROJECTS. Returns the id. */ @@ -167,7 +178,7 @@ export async function putProject(input: ProjectWrite): Promise { name: input.name, mediaKind: input.mediaKind, duration: input.duration, - model: isModelChoice(input.model) ? input.model : "base", + source: isTranscriptSource(input.source) ? input.source : "base", transcriptLanguage: isTranscriptLanguage(input.transcriptLanguage) ? input.transcriptLanguage : DEFAULT_TRANSCRIPT_LANGUAGE, diff --git a/lib/source.ts b/lib/source.ts new file mode 100644 index 0000000..ab5c671 --- /dev/null +++ b/lib/source.ts @@ -0,0 +1,11 @@ +import { isModelId, type ModelId } from "./models"; + +/** + * How the transcript is obtained on the upload screen: + * a local speech model id, or an imported caption file. + */ +export type TranscriptSource = ModelId | "import"; + +export function isTranscriptSource(value: unknown): value is TranscriptSource { + return isModelId(value) || value === "import"; +} diff --git a/lib/store.ts b/lib/store.ts index 9e0a9be..9f2520d 100644 --- a/lib/store.ts +++ b/lib/store.ts @@ -24,13 +24,8 @@ import { shrinkManualCuts, trimEdgeResult, } from "./edits"; -import { - isAsrModel, - isModelChoice, - loadModelPreference, - saveModelPreference, -} from "./models"; -import type { ModelChoice } from "./models"; +import { isModelId, loadModelPreference, saveModelPreference } from "./models"; +import { isTranscriptSource, type TranscriptSource } from "./source"; import { DEFAULT_TRANSCRIPT_LANGUAGE, isTranscriptLanguage, @@ -59,13 +54,13 @@ interface EditorState { duration: number; /** Mono 16 kHz PCM of the media's audio track (used for waveform + ASR). */ audio: Float32Array | null; - /** Transcript source selected on the upload screen (ASR model or import). */ - model: ModelChoice; + /** Transcript source selected on the upload screen (speech model or import). */ + source: TranscriptSource; /** Language hint sent to Whisper when transcribing (Parakeet auto-detects). */ transcriptLanguage: TranscriptLanguage; /** * Caption file parsed on the upload screen when source is "import". - * Cleared when switching back to an ASR model or after media loads. + * Cleared when switching back to a speech model or after media loads. */ pendingTranscript: PendingTranscript | null; /** IndexedDB project id when this session is persisted; null for a fresh upload mid-pipeline. */ @@ -121,7 +116,7 @@ interface EditorState { openProject: (id: string) => Promise; /** Delete a saved project; if it is the active one, resets to the home screen. */ removeProject: (id: string) => Promise; - setModel: (m: ModelChoice) => void; + setSource: (s: TranscriptSource) => void; setTranscriptLanguage: (language: TranscriptLanguage) => void; setPendingTranscript: (t: PendingTranscript | null) => void; setDuration: (d: number) => void; @@ -245,7 +240,7 @@ export const useEditorStore = create((set, get) => ({ mediaKind: null, duration: 0, audio: null, - model: "base", + source: "base", transcriptLanguage: DEFAULT_TRANSCRIPT_LANGUAGE, pendingTranscript: null, projectId: null, @@ -282,14 +277,14 @@ export const useEditorStore = create((set, get) => ({ if (imported && imported.length === 0) return; const prev = get().mediaUrl; if (prev) URL.revokeObjectURL(prev); - const current = get().model; + const current = get().source; set({ videoFile: file, mediaUrl: URL.createObjectURL(file), mediaKind: kind, projectId: null, skipTranscription: Boolean(imported), - model: imported ? "import" : isAsrModel(current) ? current : "base", + source: imported ? "import" : isModelId(current) ? current : "base", pendingTranscript: null, status: "preparing", progress: { @@ -328,7 +323,7 @@ export const useEditorStore = create((set, get) => ({ mediaUrl: URL.createObjectURL(file), mediaKind: record.mediaKind, duration: record.duration, - model: isModelChoice(record.model) ? record.model : "base", + source: isTranscriptSource(record.source) ? record.source : "base", transcriptLanguage: isTranscriptLanguage(record.transcriptLanguage) ? record.transcriptLanguage : DEFAULT_TRANSCRIPT_LANGUAGE, @@ -364,12 +359,12 @@ export const useEditorStore = create((set, get) => ({ } }, - setModel: (model) => { - if (isAsrModel(model)) { - saveModelPreference(model); - set({ model, pendingTranscript: null }); + setSource: (source) => { + if (isModelId(source)) { + saveModelPreference(source); + set({ source, pendingTranscript: null }); } else { - set({ model }); + set({ source }); } }, setTranscriptLanguage: (transcriptLanguage) => { @@ -426,7 +421,7 @@ export const useEditorStore = create((set, get) => ({ status: "ready", progress: { message: "", value: null }, skipTranscription: true, - model: "import", + source: "import", }); bumpAutosave(); }, @@ -708,7 +703,7 @@ export const useEditorStore = create((set, get) => ({ mediaKind: null, duration: 0, audio: null, - model: loadModelPreference(), + source: loadModelPreference(), transcriptLanguage: loadTranscriptLanguagePreference(), pendingTranscript: null, projectId: null, @@ -735,12 +730,13 @@ export const useEditorStore = create((set, get) => ({ }, })); -/** Apply the stored model choice after mount (avoids SSR/localStorage mismatch). */ +/** Apply the stored model preference after mount (avoids SSR/localStorage mismatch). */ export function hydrateModelPreference() { const stored = loadModelPreference(); - if (stored !== useEditorStore.getState().model) { - useEditorStore.setState({ model: stored }); - } + const current = useEditorStore.getState().source; + // Don't clobber an in-progress import selection. + if (current === "import" || stored === current) return; + useEditorStore.setState({ source: stored }); } /** Apply the stored transcript language after mount (avoids SSR/localStorage mismatch). */ diff --git a/lib/types.ts b/lib/types.ts index e1a9111..3750ea5 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -86,8 +86,8 @@ export interface WorkerRequest { audio: Float32Array; /** Total media duration in seconds (used for progress estimation). */ duration: number; - /** Which local ASR model to use (see lib/models.ts). */ - model: import("./models").AsrModel; + /** Which local speech model to use (see lib/models.ts). */ + model: import("./models").ModelId; /** Whisper language code for the transcript (ignored by Parakeet auto-detect). */ language: import("./languages").TranscriptLanguage; } diff --git a/tests/models-test.ts b/tests/models-test.ts index 71b4c7e..e668ded 100644 --- a/tests/models-test.ts +++ b/tests/models-test.ts @@ -1,14 +1,14 @@ /** - * Model choice helpers for Whisper / Parakeet / import. + * Model + transcript-source helpers. */ import { MODEL_ORDER, MODELS, - isAsrModel, - isModelChoice, + isModelId, isParakeetModel, isWhisperModel, } from "../lib/models"; +import { isTranscriptSource } from "../lib/source"; function assert(cond: boolean, msg: string) { if (!cond) throw new Error(msg); @@ -18,12 +18,12 @@ assert(isWhisperModel("base"), "base is Whisper"); assert(isWhisperModel("small"), "small is Whisper"); assert(!isWhisperModel("parakeet"), "parakeet is not Whisper"); assert(isParakeetModel("parakeet"), "parakeet is Parakeet"); -assert(isAsrModel("parakeet"), "parakeet is ASR"); -assert(isAsrModel("base"), "base is ASR"); -assert(!isAsrModel("import"), "import is not ASR"); -assert(isModelChoice("import"), "import is a model choice"); -assert(isModelChoice("parakeet"), "parakeet is a model choice"); -assert(!isModelChoice("tiny"), "tiny is not a model choice"); +assert(isModelId("parakeet"), "parakeet is a model id"); +assert(isModelId("base"), "base is a model id"); +assert(!isModelId("import"), "import is not a model id"); +assert(isTranscriptSource("import"), "import is a transcript source"); +assert(isTranscriptSource("parakeet"), "parakeet is a transcript source"); +assert(!isTranscriptSource("tiny"), "tiny is not a transcript source"); assert(MODELS.parakeet.backend === "parakeet", "parakeet backend"); assert(MODELS.parakeet.id === "parakeet-tdt-0.6b-v3", "parakeet hub id"); @@ -42,7 +42,7 @@ assert( "MODEL_ORDER lists whisper + parakeet" ); for (const id of MODEL_ORDER) { - assert(isAsrModel(id), `${id} in MODEL_ORDER is an ASR model`); + assert(isModelId(id), `${id} in MODEL_ORDER is a model id`); assert(typeof MODELS[id].label === "string", `${id} has label`); assert(typeof MODELS[id].size === "string", `${id} has size`); } diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 4e0d9a7..7730d4c 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -32,7 +32,7 @@ import { MODELS, isParakeetModel, isWhisperModel, - type AsrModel, + type ModelId, type WhisperModel, } from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; @@ -210,7 +210,7 @@ function parakeetModel(): ModelDefinition { */ const models = new ModelManager({ models: Object.fromEntries( - (Object.keys(MODELS) as AsrModel[]).map((choice) => { + (Object.keys(MODELS) as ModelId[]).map((choice) => { const info = MODELS[choice]; if (info.backend === "parakeet") { return [info.id, parakeetModel()]; @@ -854,7 +854,7 @@ async function runWhisper( self.onmessage = async (event: MessageEvent) => { const { audio, duration, model, language } = event.data; try { - const choice: AsrModel = model ?? "base"; + const choice: ModelId = model ?? "base"; const transcriptLanguage = language ?? "en"; let words: Word[];