open-standards/aac/aac-rs/third_party/onnxruntime-web/worker.js
Transaction Science 66218e2b83 fix(aac listener): its worker owns its GPU device, and gives it back before it ends
In WebKit, ending a worker whose WebGPU device onnxruntime-web's JSPI build
still holds ends the page with it (found by the First letters work, 215f65a4;
seen here with Playwright's WebKit, 2026-10-02). On an iPad that is the end of
every listen: the listener lets go as listening ends (the models take turns),
as the sentence helper closes, after a failed run, and when it sits unused.

The worker now asks the adapter for a device of its own (every limit and
feature it offers) and hands it to the runtime for all three sessions.
Unloading closes any listening, releases the sessions and destroys the
device, which gives back all of its graphics memory; a load that fails
destroys it too. listen.js ends a worker only once it has answered that
unload (retire, used by dropWorker, through which every path that lets the
listener go passes); one that never answers is left rather than ended.

Its bytes changed, so the runtime directory is onnxruntime-web-1.30.0-aac4,
with the worker re-pinned in RUNTIME_FILES and the directory in vendor.sh.
The bundle and the WebAssembly are unchanged.

Checked in Playwright's WebKit with the app's own code loading three
sessions on WebGPU: closing the row (twice) and a failed run leave the page
alive; the same steps with the previous worker end the page.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-02 00:47:11 -04:00

623 lines
24 KiB
JavaScript

// The listener's worker for the AAC reference app: on-device speech
// recognition that turns what the person's partner says into text. NVIDIA
// Nemotron 3.5 ASR Streaming 0.6B writes the words down; Silero VAD finds
// where speech starts and ends. Both run on onnxruntime-web's WebGPU
// execution provider. Audio in, text out.
//
// Nothing leaves the device, and nothing is kept: the audio arrives in small
// blocks while the person has asked to listen, is turned into text 320 ms at
// a time on the graphics chip, and is dropped. The worker never fetches: the
// runtime (onnxruntime-web's bundle and its WebAssembly) and the model's
// graphs and weights come from the model store (Cache API "aac-models-v1",
// each file SHA-256 checked when it was stored), so it works offline.
//
// Graphics chip only. Every session is created with the processor fallback
// switched off (session.disable_cpu_ep_fallback): a graph a device cannot run
// wholly on its graphics chip fails to load, rather than running on the
// processor.
//
// The worker has a GPU device of its own (every limit and feature the
// adapter offers, as the runtime's own would), given to the runtime for all
// three sessions. Unloading releases the sessions and destroys the device,
// which gives all of its graphics memory back; only then may the page end
// the worker. (In WebKit, ending a worker whose WebGPU device the runtime's
// JSPI build still holds ends the page with it: seen with Playwright's
// WebKit, 2026-10-02.)
//
// A module worker, started by src/bindings/listen.js from a Blob URL. It
// imports onnxruntime-web's JSPI build (JavaScript Promise Integration: the
// listener is offered only in browsers that have it) from a Blob URL too, and
// hands it the WebAssembly's bytes, so the runtime never looks for a file by
// name. JSPI lets the runtime read the model's weights one tensor at a time,
// so the whole model is never held in memory. One thread: no model compute
// on the processor.
//
// The graphs and the streaming contract are the asr-proto prototype's
// (tools/asr/README.md): each 320 ms window (5120 samples at 16 kHz, with
// the 201 before it and 40 after it) is one run of the encoder (log-mel front
// end and cache-aware encoder; its caches stay on the GPU), then greedy RNN-T
// decoding with one fused step graph (prediction net, joint, ArgMax for the
// window's four frames) that reads back only token ids.
//
// Messages in (each with an `id`, answered in order with {id, ok, ...} or
// {id, ok: false, error}):
// {type: "init", bundle, wasm} model-store URLs of the runtime
// {type: "load", key, model} model: {vad, encoder, step: {url,
// data: [{name, url}]}, vocab: url}
// {type: "check", slot} two silent windows through every graph
// {type: "start", slot} listen for one utterance (slot: the
// language's prompt slot)
// {type: "audio", samples} 16 kHz mono Float32Array (unanswered)
// {type: "finish"} end it now -> {text}
// {type: "stop"} drop it, unheard
// {type: "unload"} releases the model and destroys the
// device; the worker can then be
// ended, or load again
// Events out: {event: "speech"} someone started speaking; {event: "words",
// text} what is heard so far; {event: "quiet"} a sound with no words in it
// ended (a cough, a door): still listening, as if nothing had been heard;
// {event: "final", text} they stopped (the words in the end; listening is
// over); {event: "error", text}.
//
// Vendored into vendor/onnxruntime-web-<version>-<revision>/ as asr-worker.js
// by vendor.sh, unchanged. Its bytes are pinned in
// crates/aac-web/src/listen.rs (RUNTIME_FILES): change this file only
// together with a new ORT_WEB_VENDOR_REV and its new size and hash there.
const CACHE_NAME = "aac-models-v1";
// The model (soniqo's export of Nemotron 3.5 ASR, as edited by asr-proto).
const BLANK = 13087;
const MAX_SYMBOLS = 10; // tokens per frame at most (greedy RNN-T)
const WIN = 5120; // 320 ms
const LEFT = 201; // samples before a window the front end reads
const IN = 5361; // samples one encoder run reads
const PREEMPH = 0.97;
// When speech starts and ends (asr-proto, Silero VAD v6.2.3, one call per
// window of 10 frames of 512 samples): it starts at a frame of probability
// 0.5 or more, and the two windows before it go in first (640 ms, more than
// the 560 ms the model looks back); it ends after 1 s of frames below 0.35.
// (500 ms, the prototype's rule, cut people off at a pause between two
// sentences: tools/asr/README.md.)
const ON = 0.5;
const OFF = 0.35;
const SILENCE_MS = 1000;
const FRAME_MS = 32;
const PRE_ROLL = 2;
// One utterance is at most this long: the field it goes in holds 300
// characters.
const MAX_SPEECH_MS = 30000;
const ENC_OUT = ["encoded_output", "encoded_length", "new_pre_cache", "new_cache_last_channel",
"new_cache_last_time", "new_cache_last_channel_len"];
let ort = null;
let S = null; // {vad, enc, step} sessions
let device = null; // the GPU device they run on, this worker's own
let key = "";
let vocab = [];
let special = new Set();
let listening = null; // the Utterance being heard
let generation = 0; // counts utterances: one that ended says nothing more
function describe(err) {
return String(err && err.message ? err.message : err);
}
async function stored(cache, url) {
const res = await cache.match(url);
if (!res) throw new Error(`the listener is not on this device (${url.split("/").pop()})`);
return res;
}
async function init(msg) {
if (ort) return {};
if (typeof caches === "undefined") throw new Error("this browser keeps no model store in workers");
const cache = await caches.open(CACHE_NAME);
const [code, wasm] = await Promise.all([
stored(cache, msg.bundle).then((r) => r.blob()),
stored(cache, msg.wasm).then((r) => r.arrayBuffer()),
]);
const url = URL.createObjectURL(new Blob([code], { type: "text/javascript" }));
let runtime;
try {
runtime = await import(url);
} finally {
URL.revokeObjectURL(url);
}
runtime.env.wasm.wasmBinary = wasm; // never fetched by name
runtime.env.wasm.numThreads = 1;
runtime.env.wasm.proxy = false;
runtime.env.logLevel = "error";
ort = runtime;
return {};
}
// ─── Loading ───────────────────────────────────────────────────────────────
// The part of the model the runtime is reading now. It reads the weights one
// tensor after another, in file order, so it is mostly the same part again:
// one part (64 MiB at most) is held at a time, never the whole model.
let held = { url: "", blob: null };
function partBlob(cache, url) {
if (held.url !== url) held = { url, blob: stored(cache, url).then((r) => r.blob()) };
return held.blob;
}
// A weights file as the runtime reads it: a Blob it slices one tensor at a
// time, whose bytes stay in the model store until then.
class StoredPart extends Blob {
constructor(cache, url, bytes) {
super([]);
this.cache = cache;
this.url = url;
this.bytes = bytes;
}
get size() {
return this.bytes;
}
slice(start, end) {
const { cache, url } = this;
return {
stream: () => new ReadableStream({
async start(controller) {
const blob = await partBlob(cache, url);
controller.enqueue(new Uint8Array(await blob.slice(start, end).arrayBuffer()));
controller.close();
},
}),
};
}
}
async function load(msg) {
if (S && key === msg.key) return {};
await unload();
const cache = await caches.open(CACHE_NAME);
const m = msg.model;
const graph = async (g) => new Uint8Array(await (await stored(cache, g.url)).arrayBuffer());
// Each tensor is read from the model store when the runtime needs it.
const external = (g) => (g.data || []).map((d) => ({ path: d.name, data: new StoredPart(cache, d.url, d.bytes) }));
const options = (outputs) => ({
// Options reach the WebGPU provider only when it is named as an
// object; the device is this worker's own.
executionProviders: [{ name: "webgpu", device }],
graphOptimizationLevel: "all",
preferredOutputLocation: outputs,
logSeverityLevel: 3,
// Graphics chip or nothing: no node may fall back to the processor.
extra: { session: { disable_cpu_ep_fallback: "1" } },
});
const sessions = {};
try {
device = await ownDevice();
// The VAD first: onnxruntime-web 1.30 can break a session created
// before another one with other output locations (asr-proto, step 7).
sessions.vad = await ort.InferenceSession.create(await graph(m.vad),
options({ speech_probs: "cpu", hn: "gpu-buffer", cn: "gpu-buffer" }));
sessions.enc = await ort.InferenceSession.create(await graph(m.encoder), {
...options(Object.fromEntries(ENC_OUT.map((n) => [n, "gpu-buffer"]))),
externalData: external(m.encoder),
});
sessions.step = await ort.InferenceSession.create(await graph(m.step), {
...options({ token_out: "cpu", h_out: "gpu-buffer", c_out: "gpu-buffer" }),
externalData: external(m.step),
});
const pieces = await (await stored(cache, m.vocab)).json();
vocab = [];
for (const [i, p] of Object.entries(pieces)) vocab[Number(i)] = p;
// <unk> and the language tags (<en-US> ...) the model writes after a
// sentence's last punctuation even when told the language.
special = new Set(vocab.flatMap((p, i) => (p === "<unk>" || /^<[a-z]{2}-[A-Z]{2}>$/.test(p) ? [i] : [])));
} catch (e) {
await release(sessions);
throw e;
} finally {
held = { url: "", blob: null };
}
S = sessions;
key = msg.key;
return {};
}
// A GPU device for this worker alone: every limit and feature the adapter
// offers (the runtime asks for many of them: storage buffers per shader,
// workgroup sizes, half precision).
async function ownDevice() {
const adapter = navigator.gpu && (await navigator.gpu.requestAdapter());
if (!adapter) throw new Error("this browser offers no graphics chip to the listener");
const requiredLimits = {};
for (const k in adapter.limits) if (typeof adapter.limits[k] === "number") requiredLimits[k] = adapter.limits[k];
return adapter.requestDevice({ requiredFeatures: [...adapter.features], requiredLimits });
}
// The sessions, then the device: all of its graphics memory goes back.
async function release(sessions) {
for (const s of Object.values(sessions || {})) await s.release().catch(() => {});
if (device) device.destroy();
device = null;
}
async function unload() {
generation++;
if (listening) listening.close();
listening = null;
const old = S;
S = null;
key = "";
await release(old);
return {};
}
function detok(ids) {
let s = "";
for (const i of ids) if (!special.has(i) && vocab[i]) s += vocab[i];
return s.replaceAll("▁", " ").replace(/\s+/g, " ").trim();
}
// ─── One stream of speech ──────────────────────────────────────────────────
// The encoder's caches and the prediction net's state, carried from window to
// window on the GPU (only the first window uploads zeros).
class Stream {
constructor(slot) {
this.mask = new Float32Array(128);
this.mask[slot] = 1;
this.langMask = new ort.Tensor("float32", this.mask, [1, 128]);
this.state = {
pre_cache: new ort.Tensor("float32", new Float32Array(128 * 9), [1, 128, 9]),
cache_last_channel: new ort.Tensor("float32", new Float32Array(24 * 56 * 1024), [24, 1, 56, 1024]),
cache_last_time: new ort.Tensor("float32", new Float32Array(24 * 1024 * 8), [24, 1, 1024, 8]),
cache_last_channel_len: new ort.Tensor("int32", new Int32Array(1), [1]),
};
this.stateGpu = false;
this.h = new ort.Tensor("float16", new Uint16Array(1280), [2, 1, 640]);
this.c = new ort.Tensor("float16", new Uint16Array(1280), [2, 1, 640]);
this.hGpu = false;
this.last = BLANK;
this.ids = [];
// Samples [5120k - 201, ...) for the next window k; it starts with
// the 201 (silent) samples before the stream.
this.buf = new Float32Array(1 << 16);
this.len = LEFT;
this.k = 0;
this.total = 0;
this.lastSample = 0;
}
append(x) {
if (this.len + x.length > this.buf.length) {
const nb = new Float32Array(Math.max(this.buf.length * 2, this.len + x.length));
nb.set(this.buf.subarray(0, this.len));
this.buf = nb;
}
this.buf.set(x, this.len);
this.len += x.length;
}
// Takes 16 kHz samples; decodes every window whose samples (and the 40
// after it) are in.
async push(x) {
if (x.length) {
this.append(x);
this.total += x.length;
this.lastSample = x[x.length - 1];
}
while (this.len >= IN) await this.window();
}
// The end: the open window is padded so its pre-emphasis is exactly zero
// (x[L-1]·0.97^i), decoded, and the text returned.
async end() {
const windows = Math.max(1, Math.ceil(this.total / WIN));
while (this.k < windows) {
const need = IN - this.len;
if (need > 0) {
const pad = new Float32Array(need);
let v = this.lastSample;
for (let i = 0; i < need; i++) {
v *= PREEMPH;
pad[i] = v;
}
this.lastSample = v;
this.append(pad);
}
await this.window();
}
this.dispose();
return detok(this.ids);
}
async window() {
const w = this.buf.slice(0, IN);
this.buf.copyWithin(0, WIN, this.len);
this.len -= WIN;
this.k++;
const eo = await S.enc.run({ audio: new ort.Tensor("float32", w, [1, 1, IN]), language_mask: this.langMask, ...this.state });
eo.encoded_length.dispose();
if (this.stateGpu) for (const v of Object.values(this.state)) v.dispose();
this.state = {
pre_cache: eo.new_pre_cache,
cache_last_channel: eo.new_cache_last_channel,
cache_last_time: eo.new_cache_last_time,
cache_last_channel_len: eo.new_cache_last_channel_len,
};
this.stateGpu = true;
// Greedy RNN-T over the window's four frames. The step scores all four
// with the current prediction state: frames before the first one that
// is not blank are blank under it, so a window costs 1 + tokens runs.
try {
let t = 0;
let sym = 0;
while (t < 4) {
const r = await S.step.run({
encoder_out: eo.encoded_output,
token: new ort.Tensor("int32", Int32Array.of(this.last), [1, 1]),
h: this.h,
c: this.c,
});
const toks = r.token_out.data;
let f = -1;
for (let i = t; i < 4; i++) {
if (Number(toks[i]) !== BLANK) {
f = i;
break;
}
}
if (f < 0) {
r.h_out.dispose();
r.c_out.dispose();
break;
}
if (f !== t) {
t = f;
sym = 0;
}
const tok = Number(toks[f]);
this.ids.push(tok);
this.last = tok;
if (this.hGpu) {
this.h.dispose();
this.c.dispose();
}
this.h = r.h_out;
this.c = r.c_out;
this.hGpu = true;
if (++sym >= MAX_SYMBOLS) {
t++;
sym = 0;
}
}
} finally {
eo.encoded_output.dispose();
}
}
dispose() {
if (this.stateGpu) for (const v of Object.values(this.state)) v.dispose();
this.stateGpu = false;
if (this.hGpu) {
this.h.dispose();
this.c.dispose();
}
this.hGpu = false;
// No sound is kept once the stream ends.
this.buf = new Float32Array(0);
this.len = 0;
}
}
// ─── One utterance: the VAD decides when the stream runs ───────────────────
class Utterance {
constructor(slot, emit) {
this.slot = slot;
this.emit = emit;
this.h = new ort.Tensor("float32", new Float32Array(128), [1, 1, 128]);
this.c = new ort.Tensor("float32", new Float32Array(128), [1, 1, 128]);
this.hGpu = false;
this.ctx = new Float32Array(64); // the 64 samples before the VAD frame
this.pending = new Float32Array(0); // samples not yet in a whole window
this.history = []; // the last PRE_ROLL windows, for the pre-roll
this.stream = null;
this.silence = 0;
this.speech = 0;
this.words = "";
this.done = false;
}
async vad(w) {
const inp = new Float32Array(10 * 576);
for (let f = 0; f < 10; f++) {
inp.set(f === 0 ? this.ctx : w.subarray(f * 512 - 64, f * 512), f * 576);
inp.set(w.subarray(f * 512, (f + 1) * 512), f * 576 + 64);
}
this.ctx = w.slice(WIN - 64);
const r = await S.vad.run({ input: new ort.Tensor("float32", inp, [10, 576]), h: this.h, c: this.c });
if (this.hGpu) {
this.h.dispose();
this.c.dispose();
}
this.h = r.hn;
this.c = r.cn;
this.hGpu = true;
return Array.from(r.speech_probs.data);
}
async push(x) {
if (this.done) return;
if (this.stream) await this.stream.push(x); // the stream takes audio as it comes
const p = new Float32Array(this.pending.length + x.length);
p.set(this.pending);
p.set(x, this.pending.length);
this.pending = p;
while (this.pending.length >= WIN && !this.done) {
const w = this.pending.slice(0, WIN);
this.pending = this.pending.slice(WIN);
await this.window(w);
}
}
async window(w) {
const probs = await this.vad(w);
if (!this.stream) {
if (probs.some((q) => q >= ON)) {
// Speech: a stream from the pre-roll, this window, and what
// came after it already.
this.stream = new Stream(this.slot);
this.silence = 0;
this.speech = 0;
this.emit({ event: "speech" });
for (const hw of this.history) await this.stream.push(hw);
await this.stream.push(w);
if (this.pending.length) await this.stream.push(this.pending);
}
} else {
const words = detok(this.stream.ids);
if (words !== this.words) {
this.words = words;
this.emit({ event: "words", text: words });
}
}
if (this.stream) {
for (const q of probs) this.silence = q < OFF ? this.silence + FRAME_MS : 0;
this.speech += WIN / 16;
if (this.silence >= SILENCE_MS || this.speech >= MAX_SPEECH_MS) {
const text = await this.endStream();
if (text) {
this.close();
this.emit({ event: "final", text });
return;
}
// A sound with no words in it (a cough, a door) is not the
// end: go on listening, as if nothing had been heard.
this.emit({ event: "quiet" });
}
}
this.history.push(w);
if (this.history.length > PRE_ROLL) this.history.shift();
}
async endStream() {
const s = this.stream;
this.stream = null;
this.words = "";
return s ? s.end() : "";
}
// What was heard, now.
async finish() {
const text = this.done ? "" : await this.endStream();
this.close();
return text;
}
close() {
this.done = true;
if (this.stream) {
this.stream.dispose();
this.stream = null;
}
if (this.hGpu) {
this.h.dispose();
this.c.dispose();
}
this.hGpu = false;
// No sound is kept once listening ends.
this.pending = new Float32Array(0);
this.history = [];
this.ctx = new Float32Array(64);
}
}
// ─── Messages ──────────────────────────────────────────────────────────────
function emitter(gen) {
return (event) => {
if (gen === generation) self.postMessage(event);
};
}
function loaded() {
if (!S) throw new Error("the listener is not loaded");
}
async function check(msg) {
loaded();
// Every graph once, on two windows of silence: the first run compiles
// the graphics chip's programs, and a device that cannot run them says
// so now.
const u = new Utterance(msg.slot, () => {});
await u.push(new Float32Array(WIN));
u.close();
const s = new Stream(msg.slot);
await s.push(new Float32Array(2 * WIN));
await s.end();
return {};
}
async function start(msg) {
loaded();
if (listening) listening.close();
generation++;
listening = new Utterance(msg.slot, emitter(generation));
return {};
}
async function finish() {
const u = listening;
listening = null;
generation++;
return { text: u ? await u.finish() : "" };
}
async function stop() {
if (listening) listening.close();
listening = null;
generation++;
return {};
}
const HANDLERS = { init, load, check, start, finish, stop, unload };
// One thing at a time, in the order it came: audio, and the requests between
// it (audio sent before a stop is heard before it).
let queue = Promise.resolve();
self.onmessage = (event) => {
const msg = event.data || {};
if (msg.type === "audio") {
queue = queue.then(async () => {
const u = listening;
if (!u || !S) return;
try {
await u.push(msg.samples);
} catch (e) {
if (listening === u) {
u.close();
listening = null;
generation++;
self.postMessage({ event: "error", text: describe(e) });
}
}
});
return;
}
const handler = HANDLERS[msg.type];
queue = queue.then(async () => {
if (!handler) {
self.postMessage({ id: msg.id, ok: false, error: `unknown message ${msg.type}` });
return;
}
try {
const reply = await handler(msg);
self.postMessage({ id: msg.id, ok: true, ...reply });
} catch (e) {
self.postMessage({ id: msg.id, ok: false, error: describe(e) });
}
});
};