455 lines
17 KiB
TypeScript
455 lines
17 KiB
TypeScript
/**
|
||
* Inference Web Worker (Decisions.md D-003 — inference never runs on the main
|
||
* thread). Hosts the Transformers.js `automatic-speech-recognition` pipeline.
|
||
*
|
||
* M2: hardware backend detection with automatic fallback — "auto" resolves
|
||
* WebNN → WebGPU → WASM (D-004, D-017) — 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 {
|
||
deviceAttempts,
|
||
dtypeRequiresF16,
|
||
isUiToWorkerMessage,
|
||
TranscribeProgressTracker,
|
||
type DeviceCapabilities,
|
||
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/";
|
||
}
|
||
|
||
/**
|
||
* Set the WASM thread count for the device we are about to create a pipeline
|
||
* on (D-023). The server sends COOP/COEP, so SharedArrayBuffer exists and
|
||
* threads genuinely run — which cuts both ways:
|
||
* - Pure WASM inference scales well past onnxruntime-web's 4-thread default,
|
||
* so use most of the machine (measured 52 s → 9.5 s on a 14 s clip).
|
||
* - The JSEP-based accelerator EPs (WebGPU, WebNN) are NOT thread-safe:
|
||
* with numThreads > 1 (even ort's own default of ~4 once SAB exists),
|
||
* inference fails with "WebGPU validation failed. Buffer with '' label is
|
||
* invalid". They must run the wasm side single-threaded.
|
||
* Caveat: onnxruntime-web initializes its runtime on the first session
|
||
* creation and ignores later numThreads changes, so the first attempt in
|
||
* this worker decides. A wasm *fallback* after a failed accelerator attempt
|
||
* therefore may run single-threaded — correct, just slower; the next worker
|
||
* respawn (cancel does one, D-018) picks the full count up again.
|
||
*/
|
||
function desiredWasmThreads(): number {
|
||
return globalThis.crossOriginIsolated
|
||
? Math.max(1, Math.min((navigator.hardwareConcurrency || 4) - 1, 12))
|
||
: 1;
|
||
}
|
||
|
||
function configureWasmThreads(device: InferenceDevice): void {
|
||
const wasm = env.backends.onnx.wasm;
|
||
if (!wasm) return;
|
||
wasm.numThreads = device === "wasm" ? desiredWasmThreads() : 1;
|
||
}
|
||
|
||
/**
|
||
* Thrown to abort a load whose WASM fallback would be stuck single-threaded
|
||
* (D-025). The worker has already posted `load-restart`; the UI terminates
|
||
* this worker and retries in a fresh one, so nothing is reported as an error.
|
||
*/
|
||
class RestartRequested extends Error {
|
||
constructor() {
|
||
super("load aborted pending worker restart (D-025)");
|
||
}
|
||
}
|
||
|
||
// `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 device fallback retries (D-004, D-017).
|
||
*/
|
||
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, D-017): "auto" resolves WebNN → WebGPU → WASM.
|
||
// The probe only answers "does the backend exist?"; the attempt ordering and
|
||
// fallback policy live in the pure `deviceAttempts` (protocol.ts, unit-tested).
|
||
// ---------------------------------------------------------------------------
|
||
|
||
interface GpuAdapterLike {
|
||
features?: { has(name: string): boolean };
|
||
/** Pre-2024 spec location of the software-adapter flag. */
|
||
isFallbackAdapter?: boolean;
|
||
/** Current spec location (GPUAdapterInfo). */
|
||
info?: { isFallbackAdapter?: boolean };
|
||
}
|
||
|
||
/** True for software adapters (SwiftShader, llvmpipe, WARP) — GPU-shaped but
|
||
* slower than multithreaded WASM, so never worth an attempt (D-025). */
|
||
function isSoftwareAdapter(adapter: GpuAdapterLike): boolean {
|
||
return Boolean(adapter.isFallbackAdapter ?? adapter.info?.isFallbackAdapter);
|
||
}
|
||
|
||
interface NavigatorWithHardware {
|
||
gpu?: { requestAdapter(): Promise<GpuAdapterLike | null> };
|
||
ml?: unknown;
|
||
}
|
||
|
||
async function probeCapabilities(
|
||
preference: DevicePreference,
|
||
): Promise<DeviceCapabilities> {
|
||
const capabilities: DeviceCapabilities = { webnn: false, webgpu: false };
|
||
const hardware = navigator as NavigatorWithHardware;
|
||
if (preference === "auto" || preference === "webnn") {
|
||
try {
|
||
capabilities.webnn = Boolean(hardware.ml);
|
||
} catch {
|
||
// Probing threw — treat WebNN as unavailable.
|
||
}
|
||
}
|
||
if (preference === "auto" && /Firefox\//.test(navigator.userAgent)) {
|
||
// Firefox creates ort-web WebGPU sessions fine but runs them far slower
|
||
// than this machine's multithreaded WASM (field-verified 2026-07-17,
|
||
// D-026) — and a *successful* slow session is invisible to the
|
||
// failure-driven fallback chain. Never auto-pick WebGPU on Firefox; an
|
||
// explicit selection is honored (the retest hatch for future ort/Firefox).
|
||
console.info(
|
||
"[whisper] Firefox: skipping WebGPU for \"auto\" — onnxruntime-web's WebGPU backend " +
|
||
"runs much slower than multithreaded WASM on Firefox (D-026). " +
|
||
"Select WebGPU explicitly to retest.",
|
||
);
|
||
} else if (preference === "auto" || preference === "webgpu") {
|
||
try {
|
||
const adapter = (await hardware.gpu?.requestAdapter()) ?? null;
|
||
if (adapter && isSoftwareAdapter(adapter)) {
|
||
// A fallback (software-rendered) adapter would "succeed" and then run
|
||
// far slower than multithreaded WASM — never attempt it (D-025).
|
||
console.info("[whisper] WebGPU adapter is software-rendered — skipping WebGPU.");
|
||
} else {
|
||
capabilities.webgpu = adapter !== null;
|
||
// fp16 weights need the "shader-f16" feature (D-025); Firefox ships
|
||
// WebGPU without it, so createPipeline gates on this.
|
||
capabilities.webgpuF16 = Boolean(adapter?.features?.has("shader-f16"));
|
||
}
|
||
} catch {
|
||
// Probing threw — treat WebGPU as unavailable.
|
||
}
|
||
}
|
||
return capabilities;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 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 capabilities = await probeCapabilities(message.device);
|
||
// WebGPU without "shader-f16" (Firefox) cannot run fp16 weights — skip the
|
||
// doomed attempt instead of failing into a degraded-thread fallback (D-025).
|
||
if (capabilities.webgpu && !capabilities.webgpuF16) {
|
||
const webgpuDtype =
|
||
message.dtype ?? resolveDefaultDtype(message.modelId, "webgpu");
|
||
if (dtypeRequiresF16(webgpuDtype)) {
|
||
capabilities.webgpu = false;
|
||
console.info(
|
||
"[whisper] WebGPU adapter lacks shader-f16 but the model dtype needs fp16 — skipping WebGPU.",
|
||
);
|
||
}
|
||
}
|
||
// Ordered devices to try (D-004/D-017): the first is the selected device;
|
||
// if pipeline creation throws there (driver quirks, unsupported ops,
|
||
// out-of-memory), the next one is tried automatically. For "auto" the chain
|
||
// is WebNN → WebGPU → WASM (filtered by the probe); WASM always terminates it.
|
||
const attempts: InferenceDevice[] = deviceAttempts(message.device, capabilities);
|
||
|
||
let lastError: unknown;
|
||
let acceleratorAttempted = false;
|
||
for (const device of attempts) {
|
||
// A failed WebGPU/WebNN attempt has already initialized onnxruntime's
|
||
// WASM runtime single-threaded, and the thread pool is fixed at first
|
||
// init (D-023) — a WASM session created now would run on one thread for
|
||
// this worker's whole life (measured 5.5× slower). When multithreading is
|
||
// actually available, hand the fallback to a fresh worker instead (D-025).
|
||
if (device === "wasm" && acceleratorAttempted && desiredWasmThreads() > 1) {
|
||
post({
|
||
type: "load-restart",
|
||
modelId: message.modelId,
|
||
preference: message.device,
|
||
device: "wasm",
|
||
});
|
||
throw new RestartRequested();
|
||
}
|
||
configureWasmThreads(device);
|
||
const dtype = message.dtype ?? resolveDefaultDtype(message.modelId, device);
|
||
try {
|
||
const asr = await pipeline("automatic-speech-recognition", message.modelId, {
|
||
device,
|
||
dtype,
|
||
// WASM sessions cap graph optimization at "basic" (D-014): the level-2
|
||
// QDQ transform in the bundled onnxruntime-web build rejects the 8-bit
|
||
// decoders' tied embed_tokens weight ("Missing required scale …
|
||
// DequantizeLinear"). q4 models use native MatMulNBits ops, so the
|
||
// lower level costs them nothing. WebGPU keeps the default level.
|
||
...(device === "wasm" && {
|
||
session_options: { graphOptimizationLevel: "basic" as const },
|
||
}),
|
||
progress_callback: forwardProgress,
|
||
});
|
||
// What ort actually ended up with (it clamps to 1 without isolation) —
|
||
// surfaced to the UI so degraded threading is visible in the field.
|
||
const threads = env.backends.onnx.wasm?.numThreads ?? 1;
|
||
console.info(`[whisper] pipeline ready: device=${device} wasmThreads=${threads}`);
|
||
post({ type: "load-progress", status: "ready", device, threads });
|
||
return asr;
|
||
} catch (error) {
|
||
lastError = error;
|
||
if (device !== "wasm") acceleratorAttempted = true;
|
||
console.warn(
|
||
`[whisper] pipeline creation failed on ${device}: ${errorMessage(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
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/** Audio sample rate the whole app standardizes on (D-012). */
|
||
const SAMPLE_RATE = 16_000;
|
||
/** Chunked-pipeline window / overlap, in seconds (Whisper's usual 30/5). */
|
||
const CHUNK_LENGTH_S = 30;
|
||
const STRIDE_LENGTH_S = 5;
|
||
|
||
async function transcribe(message: TranscribeMessage): Promise<void> {
|
||
if (active === null) {
|
||
post({
|
||
type: "error",
|
||
stage: "transcribe",
|
||
message: "No model loaded — send a \"load\" message first.",
|
||
requestId: message.requestId,
|
||
});
|
||
return;
|
||
}
|
||
const asr = await active.promise;
|
||
|
||
const options: Record<string, unknown> = {
|
||
chunk_length_s: CHUNK_LENGTH_S,
|
||
stride_length_s: STRIDE_LENGTH_S,
|
||
return_timestamps: true,
|
||
// With return_timestamps and no explicit max_new_tokens, Transformers.js
|
||
// 4.2's Whisper generate() takes a seek-loop path that re-decodes each
|
||
// window's tail (1–3 generate passes per window). Under this pipeline's
|
||
// own 30 s/5 s chunking that is pure overhead: ~1.4× slower on long
|
||
// files, duplicated boundary text, on_finalize firing once per pass
|
||
// (breaking TranscribeProgressTracker's one-finalize-per-window
|
||
// assumption — the progress bar saturated early, then looked frozen),
|
||
// and a potential infinite loop when a window's last boundary timestamp
|
||
// is 0. Setting max_new_tokens selects the classic single-pass path;
|
||
// 440 ≈ Whisper's 448 decoder positions minus prompt headroom, far more
|
||
// than any real 30 s window produces.
|
||
max_new_tokens: 440,
|
||
};
|
||
if (message.stream ?? message.requestId === undefined) {
|
||
// File-flow pass (id-less, or routed through BrowserEngine with an
|
||
// explicit stream flag, D-010): stream partial tokens and report audio
|
||
// progress. The streamer's chunk timestamps are window-relative; the
|
||
// tracker maps them to monotone absolute seconds (see protocol.ts).
|
||
// Live passes skip both — they resolve wholesale via "result" (D-016).
|
||
const totalSeconds = message.audio.length / SAMPLE_RATE;
|
||
const tracker = new TranscribeProgressTracker(
|
||
totalSeconds,
|
||
CHUNK_LENGTH_S,
|
||
CHUNK_LENGTH_S - 2 * STRIDE_LENGTH_S,
|
||
);
|
||
const postProgress = (processedSeconds: number | null) => {
|
||
if (processedSeconds === null) return;
|
||
post({
|
||
type: "transcribe-progress",
|
||
processedSeconds,
|
||
totalSeconds,
|
||
requestId: message.requestId,
|
||
});
|
||
};
|
||
// Immediate 0/total so the UI shows the audio length before the first
|
||
// (possibly slow) token arrives.
|
||
post({
|
||
type: "transcribe-progress",
|
||
processedSeconds: 0,
|
||
totalSeconds,
|
||
requestId: message.requestId,
|
||
});
|
||
options.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, requestId: message.requestId }),
|
||
on_chunk_start: (time: number) => postProgress(tracker.timestamp(time)),
|
||
on_chunk_end: (time: number) => postProgress(tracker.timestamp(time)),
|
||
on_finalize: () => postProgress(tracker.finalizeWindow()),
|
||
},
|
||
);
|
||
}
|
||
// 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 ?? [],
|
||
requestId: message.requestId,
|
||
};
|
||
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) => {
|
||
// A restart-aborted load is not an error — the UI already got the
|
||
// load-restart message and is respawning this worker (D-025).
|
||
if (error instanceof RestartRequested) return;
|
||
post({ type: "error", stage: "load", message: errorMessage(error) });
|
||
});
|
||
break;
|
||
case "transcribe":
|
||
transcribe(data).catch((error: unknown) => {
|
||
if (error instanceof RestartRequested) return;
|
||
post({
|
||
type: "error",
|
||
stage: "transcribe",
|
||
message: errorMessage(error),
|
||
requestId: data.requestId,
|
||
});
|
||
});
|
||
break;
|
||
}
|
||
});
|