diff --git a/lib/webgpu.ts b/lib/webgpu.ts new file mode 100644 index 0000000..23467e0 --- /dev/null +++ b/lib/webgpu.ts @@ -0,0 +1,19 @@ +/** + * True for ONNX Runtime / WebGPU failures caused by a lost GPU device + * (Windows screen lock, driver reset, Chromium GPU process teardown, …). + */ +export function isWebGpuDeviceLostError(err: unknown): boolean { + const msg = (err instanceof Error ? err.message : String(err ?? "")).toLowerCase(); + if (!msg) return false; + return ( + msg.includes("failed to download data from buffer") || + msg.includes("external instance reference no longer exists") || + msg.includes("buffer was destroyed before") || + msg.includes("buffer was unmapped before mapping was resolved") || + msg.includes("device was lost") || + msg.includes("webgpu device lost") || + msg.includes("gpudevice lost") || + (msg.includes("buffer_manager.cc") && msg.includes("download")) || + (msg.includes("mapasync") && msg.includes("gpubuffer")) + ); +} diff --git a/scripts/webgpu-test.ts b/scripts/webgpu-test.ts new file mode 100644 index 0000000..c9fe038 --- /dev/null +++ b/scripts/webgpu-test.ts @@ -0,0 +1,34 @@ +import { isWebGpuDeviceLostError } from "../lib/webgpu"; + +function assert(cond: unknown, msg: string): asserts cond { + if (!cond) throw new Error(msg); +} + +{ + const issue39 = new Error( + "failed to call OrtRun(). ERROR_CODE: 1, ERROR_MESSAGE: /mnt/vss/_work/1/s/onnxruntime/core/providers/webgpu/buffer_manager.cc:553 auto onnxruntime::webgpu::BufferManager::Download(WGPUBuffer, void *, size_t)::(lambda)::operator()(wgpu::MapAsyncStatus, wgpu::StringView) const status == wgpu::MapAsyncStatus::Success was false. Failed to download data from buffer: Failed to execute 'mapAsync' on 'GPUBuffer': A valid external Instance reference no longer exists." + ); + assert(isWebGpuDeviceLostError(issue39), "issue #39 OrtRun buffer download should match"); +} + +{ + assert( + isWebGpuDeviceLostError("WebGPU device lost (2): Device was destroyed."), + "transformers.js device-destroyed message should match" + ); + assert( + isWebGpuDeviceLostError( + new Error("Failed to execute 'mapAsync' on 'GPUBuffer': Buffer was unmapped before mapping was resolved.") + ), + "unmapped buffer race should match" + ); +} + +{ + assert(!isWebGpuDeviceLostError(new Error("Could not extract audio from this file.")), "audio error"); + assert(!isWebGpuDeviceLostError("out of memory"), "generic oom should not match"); + assert(!isWebGpuDeviceLostError(null), "null"); + assert(!isWebGpuDeviceLostError(""), "empty"); +} + +console.log("webgpu-test: ok"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 42ec21c..2c43348 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -31,6 +31,7 @@ import { speechSegmentsFromFrames, type SpeechSegment, } from "@/lib/vad"; +import { isWebGpuDeviceLostError } from "@/lib/webgpu"; env.allowLocalModels = false; // Serve onnxruntime-web WASM from our own origin (offline friendly). @@ -51,6 +52,8 @@ const SPEECH_PAD_S = 0.4; */ const WHISPER_LEAD_PAD_S = 0.5; +type AsrChunk = { text: string; timestamp: [number, number | null] }; + const post = (msg: WorkerResponse, transfer: Transferable[] = []) => (self as unknown as Worker).postMessage(msg, transfer); @@ -127,7 +130,13 @@ async function isModelCached(modelId: string): Promise { } } +/** Once set, this worker stays on WASM — a lost WebGPU session cannot be reused. */ +let forceWasm = false; +/** Device the current ASR pipeline is running on. */ +let asrDevice: "webgpu" | "wasm" = "wasm"; + async function pickDevice(): Promise<"webgpu" | "wasm"> { + if (forceWasm) return "wasm"; try { const gpu = (globalThis.navigator as Navigator & { gpu?: { requestAdapter: () => Promise }; @@ -146,6 +155,7 @@ async function getAsr(choice: WhisperModel) { let promise = asrPromises.get(id); if (!promise) { const device = await pickDevice(); + asrDevice = device; const label = (await isModelCached(id)) ? "Loading speech model from cache…" : "Downloading speech model…"; @@ -156,6 +166,8 @@ async function getAsr(choice: WhisperModel) { }).catch((err) => { // WebGPU can fail on some drivers; retry once on plain WASM. if (device === "webgpu") { + forceWasm = true; + asrDevice = "wasm"; return pipeline("automatic-speech-recognition", id, { dtype: dtype.wasm, device: "wasm", @@ -170,6 +182,19 @@ async function getAsr(choice: WhisperModel) { return promise; } +/** Drop the dead WebGPU pipeline and reload the same model on WASM. */ +async function fallbackAsrToWasm(choice: WhisperModel) { + forceWasm = true; + asrDevice = "wasm"; + asrPromises.clear(); + post({ + type: "progress", + message: "GPU interrupted — continuing on CPU…", + value: null, + }); + return getAsr(choice); +} + /** * Build decoder input tokens implementing Whisper's "initial prompt" * conditioning: `<|startofprev|> …prompt… <|startoftranscript|> <|lang|> @@ -319,8 +344,6 @@ async function detectSpeechSegments( } } -type AsrChunk = { text: string; timestamp: [number, number | null] }; - /** Nominal length given to a word whose end timestamp is missing or unusable. */ const FALLBACK_WORD_S = 0.5; @@ -450,7 +473,8 @@ self.onmessage = async (event: MessageEvent) => { // Overlap Whisper + Silero downloads; diarizer warms in the background. getDiarizer().catch(() => {}); - const [transcriber, vad] = await Promise.all([getAsr(choice), getVad()]); + 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 } = @@ -482,10 +506,6 @@ self.onmessage = async (event: MessageEvent) => { // @ts-expect-error model config is untyped (transcriber.model.config.max_source_positions ?? 1500); - const tokenizer = transcriber.tokenizer as ConstructorParameters< - typeof WhisperTextStreamer - >[0]; - let speechDone = 0; let transcribed = 0; let chunkFloor = 0; @@ -550,37 +570,65 @@ self.onmessage = async (event: MessageEvent) => { const sliceDuration = slice.length / VAD_SAMPLE_RATE; const offsetS = seg.startSample / VAD_SAMPLE_RATE - WHISPER_LEAD_PAD_S; - // 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 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(); - }, - }); + // 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[]; + }; + + 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 output = await transcriber(slice, { ...asrOptions, streamer }); - const result = Array.isArray(output) ? output[0] : output; - const chunks = (result.chunks ?? []) as AsrChunk[]; rawWords.push(...wordsFromChunks(chunks, offsetS, sliceDuration, duration)); - speechDone += segmentSamples; reportProgress(0, 0); } @@ -608,7 +656,11 @@ self.onmessage = async (event: MessageEvent) => { console.error(err); post({ type: "error", - message: err instanceof Error ? err.message : "Transcription failed.", + message: isWebGpuDeviceLostError(err) + ? "Transcription was interrupted when the GPU reset (often after locking the screen). Please try again." + : err instanceof Error + ? err.message + : "Transcription failed.", }); } };