diff --git a/components/Editor.tsx b/components/Editor.tsx index c277a0e..ea57f58 100644 --- a/components/Editor.tsx +++ b/components/Editor.tsx @@ -23,6 +23,7 @@ import ModelSelector, { ModelOptionSeparator, } from "./ModelSelector"; import ImportTranscriptOption from "./ImportTranscriptOption"; +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 @@ -230,8 +231,9 @@ export default function Editor() { <> {isElectron && - - + {MODEL_ORDER.map((id) => ( + + ))} diff --git a/components/ImportTranscriptOption.tsx b/components/ImportTranscriptOption.tsx index d3ab301..b403360 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 { 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" | "small">("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 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 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 (!isWhisperModel(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 (isWhisperModel(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 619642d..11c6409 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 { + MODEL_ORDER, + MODELS, + isModelId, + isWhisperModel, +} from "@/lib/models"; +import type { TranscriptSource } from "@/lib/source"; import { hydrateModelPreference, hydrateTranscriptLanguagePreference, @@ -35,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. */ @@ -53,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; @@ -82,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 ) { @@ -120,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>({}); @@ -168,39 +174,40 @@ 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 ?? - (isWhisperModel(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"; // Always mount options (hidden when closed) so custom triggers stay registered. const options = children ?? ( <> - - + {MODEL_ORDER.map((id) => ( + + ))} ); @@ -294,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, @@ -305,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; @@ -316,10 +323,8 @@ export function ModelOption({ const selector = useSelectorCtx(); const selected = selector.value === id; - const resolvedLabel = - label ?? (isWhisperModel(id) ? MODELS[id].label : id); - const resolvedMeta = - meta ?? (isWhisperModel(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 b012ebf..b2d8888 100644 --- a/components/UploadScreen.tsx +++ b/components/UploadScreen.tsx @@ -23,6 +23,7 @@ import ModelSelector, { ModelOptionSeparator, } from "./ModelSelector"; import ImportTranscriptOption from "./ImportTranscriptOption"; +import { MODEL_ORDER } from "@/lib/models"; import { useCrossOriginIsolated } from "@/hooks/useCrossOriginIsolated"; import { detectMediaKind, MEDIA_ACCEPT } from "@/lib/media"; import { formatTime } from "@/lib/edits"; @@ -166,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); @@ -205,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."); @@ -275,8 +275,9 @@ export default function UploadScreen({
- - + {MODEL_ORDER.map((id) => ( + + ))} @@ -341,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 b413d4b..7c1ae43 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 { isModelId } 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 (!isModelId(store.source)) { + store.setError("Select a speech model to transcribe."); return; } - const whisperModel = store.model; + const model = store.source; 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, 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 5387a68..53b8eaa 100644 --- a/lib/models.ts +++ b/lib/models.ts @@ -1,16 +1,23 @@ -/** Transcription source choices offered on the upload screen. */ +/** Local speech models 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 ModelId = WhisperModel | ParakeetModel; 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 speech 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; @@ -27,21 +34,40 @@ export interface ModelInfo { * required for word-level timestamps, which this editor depends on.) */ verbatimPrompt?: string; -} +}; -/** Display order for Whisper rows in the homepage source dropdown. */ -export const WHISPER_ORDER: WhisperModel[] = ["base", "small"]; +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 model rows in the source dropdown. */ +export const MODEL_ORDER: ModelId[] = ["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 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`. + */ +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.", @@ -52,40 +78,54 @@ 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", + }, }; export function isWhisperModel(value: unknown): value is WhisperModel { return value === "base" || value === "small"; } -export function isModelChoice(value: unknown): value is ModelChoice { - return isWhisperModel(value) || value === "import"; +export function isParakeetModel(value: unknown): value is ParakeetModel { + return value === "parakeet"; +} + +/** 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); } 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 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 (isWhisperModel(raw)) return raw; + if (isModelId(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 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 f0a3d25..9f2520d 100644 --- a/lib/store.ts +++ b/lib/store.ts @@ -24,13 +24,8 @@ import { shrinkManualCuts, trimEdgeResult, } from "./edits"; -import { - isModelChoice, - isWhisperModel, - 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,19 +54,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). */ - model: ModelChoice; - /** Language hint sent to Whisper when transcribing. */ + /** 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 a Whisper 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. */ 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; @@ -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" : isWhisperModel(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 (isWhisperModel(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 cc4db3c..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 Whisper model to use (see lib/models.ts). */ - model: import("./models").WhisperModel; - /** Whisper language code for the transcript. */ + /** 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/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 a06a41e..b3642fe 100644 --- a/package-lock.json +++ b/package-lock.json @@ -20,6 +20,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", @@ -8779,6 +8780,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 48f686b..2b72b95 100644 --- a/package.json +++ b/package.json @@ -42,6 +42,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..e668ded --- /dev/null +++ b/tests/models-test.ts @@ -0,0 +1,50 @@ +/** + * Model + transcript-source helpers. + */ +import { + MODEL_ORDER, + MODELS, + isModelId, + isParakeetModel, + isWhisperModel, +} from "../lib/models"; +import { isTranscriptSource } from "../lib/source"; + +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(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"); +assert( + MODELS.parakeet.repoId === "ysdede/parakeet-tdt-0.6b-v3-onnx", + "parakeet HF repo id" +); +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( + MODEL_ORDER.includes("parakeet") && MODEL_ORDER.includes("base"), + "MODEL_ORDER lists whisper + parakeet" +); +for (const id of MODEL_ORDER) { + 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`); +} + +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 bd9c258..7730d4c 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). + * 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, @@ -20,13 +22,19 @@ import { env, type AutomaticSpeechRecognitionPipeline, } from "@huggingface/transformers"; -import { ModelManager } from "weightlift"; +import { ModelManager, type ModelDefinition } from "weightlift"; import { fallbackDevicePolicy, transformersModel, } from "weightlift/transformers"; import type { Word, WorkerRequest, WorkerResponse } from "@/lib/types"; -import { MODELS, type WhisperModel } from "@/lib/models"; +import { + MODELS, + isParakeetModel, + isWhisperModel, + type ModelId, + type WhisperModel, +} from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; import { alignWordsToSpeech } from "@/lib/align"; import { @@ -39,9 +47,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"; @@ -65,22 +76,152 @@ 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 = MODELS.parakeet.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(MODELS.parakeet.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(MODELS.parakeet.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 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({ +const models = new ModelManager({ models: Object.fromEntries( - (Object.keys(MODELS) as WhisperModel[]).map((choice) => { - const { id, dtype } = MODELS[choice]; + (Object.keys(MODELS) as ModelId[]).map((choice) => { + const info = MODELS[choice]; + if (info.backend === "parakeet") { + return [info.id, parakeetModel()]; + } return [ - id, + info.id, transformersModel({ pipeline, task: "automatic-speech-recognition", - modelId: id, - dtype, + modelId: info.id, + dtype: info.dtype, cacheKey: env.cacheKey ?? "transformers-cache", onDevice: (device) => { asrDevice = device; @@ -90,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]; @@ -106,24 +247,27 @@ asrModels.subscribe((snap) => { }); async function getAsr(choice: WhisperModel) { - return asrModels.load(MODELS[choice].id); + return models.load(MODELS[choice].id); +} + +async function getParakeet() { + return models.load(MODELS.parakeet.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(); + await models.unloadAll(); post({ type: "progress", message: "GPU interrupted — continuing on CPU…", value: null, }); - return getAsr(choice); } /** @@ -397,192 +541,329 @@ 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"; +/** 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)); - // 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."); + 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 speechSamples = speechSegments.reduce( - (n, s) => n + (s.endSample - s.startSample), - 0 - ); +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; +} - 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 }); - }; +async function runParakeet( + audio: Float32Array, + duration: number +): Promise { + getDiarizer().catch(() => {}); + 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); + + 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, + }); - /** 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 }); + let result: Awaited>; + try { + result = await runSlice(); + } catch (err) { + if (asrDevice !== "webgpu" || !isWebGpuDeviceLostError(err)) { + throw err; } - }; + console.warn( + "WebGPU lost during Parakeet transcription; reloading on WASM.", + err + ); + await fallbackAsrToWasm(); + model = await getParakeet(); + result = await runSlice(); + } - 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 }), - }; + 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 }); + } - 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[]; - }; + speechDone += segmentSamples; + const value = + speechSamples > 0 ? Math.min(1, speechDone / speechSamples) : 1; + post({ type: "progress", message: "Transcribing…", value }); + } - 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(); - } + 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."); + } - rawWords.push(...wordsFromChunks(chunks, offsetS, sliceDuration, duration)); - speechDone += segmentSamples; - reportProgress(0, 0); + 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 }); + await fallbackAsrToWasm(); + transcriber = await getAsr(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: ModelId = 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 });