Fix transcription failure after Windows screen lock (#39)

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"))
+  );
+}

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");

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<boolean> {
   }
 }
 
+/** 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<unknown | null> };
@@ -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<WorkerRequest>) => {
 
     // 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<WorkerRequest>) => {
       // @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<WorkerRequest>) => {
       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<WorkerRequest>) => {
     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.",
     });
   }
 };