whisper_web/src/worker/inference.worker.ts

/**
 * Inference Web Worker (Decisions.md D-003 — inference never runs on the main
 * thread). Hosts the Transformers.js `automatic-speech-recognition` pipeline.
 *
 * M2: WebGPU detection with automatic WASM fallback (D-004), catalog-driven
 * dtype defaults (D-006), model switching with pipeline disposal, and
 * language/task options for multilingual models.
 */

import {
  env,
  pipeline,
  WhisperTextStreamer,
  type AutomaticSpeechRecognitionPipeline,
  type ProgressInfo,
} from "@huggingface/transformers";
import { resolveDefaultDtype } from "../lib/models";
import {
  isUiToWorkerMessage,
  type DevicePreference,
  type InferenceDevice,
  type LoadMessage,
  type ResultMessage,
  type TranscribeMessage,
  type TranscriptChunk,
  type WorkerToUiMessage,
} from "./protocol";

// Load onnxruntime's .wasm binaries from our own origin instead of the
// jsdelivr CDN default — required for offline use (M3). The server exposes
// them at /ort/* (src/index.ts; copied into dist/ort by scripts/build.ts)
// and the service worker runtime-caches them cache-first (src/pwa/sw.ts).
if (env.backends.onnx.wasm) {
  env.backends.onnx.wasm.wasmPaths = "/ort/";
}

// `lib: ["DOM"]` types `self` as Window; give it the worker-scope shape we use.
interface DedicatedWorkerScope {
  postMessage(message: WorkerToUiMessage): void;
  addEventListener(
    type: "message",
    listener: (event: MessageEvent<unknown>) => void,
  ): void;
}
const scope = self as unknown as DedicatedWorkerScope;

function post(message: WorkerToUiMessage): void {
  scope.postMessage(message);
}

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

/**
 * Forward Transformers.js progress events to the UI as load-progress messages.
 * The library's own "ready" event is swallowed — the worker posts its own
 * `ready` (with the selected device) once the pipeline actually resolved,
 * which also covers the WebGPU → WASM fallback retry (D-004).
 */
function forwardProgress(info: ProgressInfo): void {
  switch (info.status) {
    case "initiate":
    case "download":
    case "done":
      post({ type: "load-progress", status: info.status, file: info.file });
      break;
    case "progress":
      post({
        type: "load-progress",
        status: "progress",
        file: info.file,
        loaded: info.loaded,
        total: info.total,
        progress: info.progress,
      });
      break;
    default:
      // "ready" (posted by the worker itself, with the device) and
      // "progress_total" — the UI derives totals per-file itself.
      break;
  }
}

// ---------------------------------------------------------------------------
// Device selection (D-004): WebGPU when available, automatic WASM fallback
// ---------------------------------------------------------------------------

interface NavigatorWithGpu {
  gpu?: { requestAdapter(): Promise<unknown | null> };
}

async function detectDevice(preference: DevicePreference): Promise<InferenceDevice> {
  if (preference === "wasm") return "wasm";
  try {
    const adapter = await (navigator as NavigatorWithGpu).gpu?.requestAdapter();
    return adapter ? "webgpu" : "wasm";
  } catch {
    return "wasm";
  }
}

// ---------------------------------------------------------------------------
// Pipeline lifecycle (load, switch, dispose)
// ---------------------------------------------------------------------------

interface ActivePipeline {
  /** Identity of this load request; a different key means "switch models". */
  key: string;
  promise: Promise<AutomaticSpeechRecognitionPipeline>;
}

let active: ActivePipeline | null = null;

function loadKey(message: LoadMessage): string {
  return JSON.stringify({
    modelId: message.modelId,
    device: message.device,
    dtype: message.dtype ?? null,
  });
}

async function createPipeline(
  message: LoadMessage,
): Promise<AutomaticSpeechRecognitionPipeline> {
  const detected = await detectDevice(message.device);
  // If WebGPU was selected but pipeline creation throws (driver quirks,
  // unsupported ops, out-of-memory), fall back to WASM automatically (D-004).
  const attempts: InferenceDevice[] =
    detected === "webgpu" ? ["webgpu", "wasm"] : ["wasm"];

  let lastError: unknown;
  for (const device of attempts) {
    const dtype = message.dtype ?? resolveDefaultDtype(message.modelId, device);
    try {
      const asr = await pipeline("automatic-speech-recognition", message.modelId, {
        device,
        dtype,
        progress_callback: forwardProgress,
      });
      post({ type: "load-progress", status: "ready", device });
      return asr;
    } catch (error) {
      lastError = error;
    }
  }
  throw lastError;
}

function load(message: LoadMessage): Promise<AutomaticSpeechRecognitionPipeline> {
  const key = loadKey(message);
  // Idempotent: repeated "load" messages for the same model/device/dtype reuse
  // the in-flight or completed pipeline.
  if (active?.key === key) return active.promise;

  const previous = active;
  const promise = (async () => {
    if (previous) {
      // Dispose the old pipeline (frees WASM memory / GPU buffers) before
      // creating the new one, so both never coexist.
      try {
        const old = await previous.promise;
        await old.dispose();
      } catch {
        // The previous load itself failed — nothing to dispose.
      }
    }
    return createPipeline(message);
  })();
  active = { key, promise };
  promise.catch(() => {
    // Allow a clean retry after a failed load; ignore stale failures that
    // were already superseded by a newer load request.
    if (active?.key === key) active = null;
  });
  return promise;
}

// ---------------------------------------------------------------------------
// Transcription
// ---------------------------------------------------------------------------

async function transcribe(message: TranscribeMessage): Promise<void> {
  if (active === null) {
    post({
      type: "error",
      stage: "transcribe",
      message: "No model loaded — send a \"load\" message first.",
    });
    return;
  }
  const asr = await active.promise;

  const streamer = new WhisperTextStreamer(
    // The ASR pipeline for Whisper models always carries a WhisperTokenizer.
    asr.tokenizer as ConstructorParameters<typeof WhisperTextStreamer>[0],
    {
      skip_prompt: true,
      callback_function: (text: string) => post({ type: "partial", text }),
    },
  );

  const options: Record<string, unknown> = {
    chunk_length_s: 30,
    stride_length_s: 5,
    return_timestamps: true,
    streamer,
  };
  // Only meaningful for multilingual models; the UI omits these for `.en`
  // variants. "auto" / undefined lets Whisper detect the language.
  if (message.language !== undefined && message.language !== "auto") {
    options.language = message.language;
  }
  if (message.task !== undefined) {
    options.task = message.task;
  }

  const output = await asr(message.audio, options);

  const single = (Array.isArray(output) ? output[0] : output) as {
    text?: string;
    chunks?: TranscriptChunk[];
  };
  const result: ResultMessage = {
    type: "result",
    text: (single.text ?? "").trim(),
    chunks: single.chunks ?? [],
  };
  post(result);
}

scope.addEventListener("message", (event) => {
  const data = event.data;
  if (!isUiToWorkerMessage(data)) {
    post({
      type: "error",
      stage: "transcribe",
      message: "Malformed message sent to inference worker.",
    });
    return;
  }
  switch (data.type) {
    case "load":
      load(data).catch((error: unknown) => {
        post({ type: "error", stage: "load", message: errorMessage(error) });
      });
      break;
    case "transcribe":
      transcribe(data).catch((error: unknown) => {
        post({ type: "error", stage: "transcribe", message: errorMessage(error) });
      });
      break;
  }
});