ava-real / workers /model-worker.js
fivetech's picture
Voz femenina por defecto (es_ES-sharvard speaker F); lib patched para speakerId
c5622cf
Raw History Blame Contribute Delete
13.7 kB
// @ts-check
/**
* LFM3-Omni on WebGPU, off the main thread so the avatar keeps its frame rate.
* One conversation lives here: the history, the prefix-cached session, and the
* context budget. Everything the main thread needs comes back as a stream of events.
*
* main -> worker:
* { kind: "load", model, system }
* { kind: "respond", id, audio } a user turn, 16 kHz Float32
* { kind: "respond", id, text } a user turn typed in the chat box
* { kind: "abort", id } barge-in: stop generating this turn
* { kind: "system", system } new persona / instructions
* { kind: "reset" } forget the conversation
*
* worker -> main:
* { kind: "progress", loaded?, total?, fraction?, status? }
* { kind: "ready", device, sampleRate } { kind: "failed", name, message }
* { kind: "text", id, text } { kind: "audio", id, samples }
* { kind: "finished", id, reason, usage, generated, error? }
*
* Voice: with Piper enabled (default, `?voice=` selects or disables with
* "model") replies generate as text only and each sentence is spoken by a
* native Spanish Piper voice, resampled into the model's audio rate.
*/
import {
Lfm3OmniForConditionalGeneration as Model,
ContextOverflowError,
DeviceLostError,
} from "../vendor/Lfm3OmniForConditionalGeneration.js";
const LOAD_ATTEMPTS = 4; // Chromium occasionally returns no adapter on a cold GPU process, and networks drop requests
const LOAD_RETRY_MS = 1500;
// Her replies stay in history as text only: her speech costs ~16 tokens/s and left
// no room for what was said, so a 44-turn chat now fits and she still recalls it.
// Past the budget half of it goes, well under the session's 8192 tokens.
const SEARCH = new URL(location.href).searchParams;
const HISTORY_BUDGET_TOKENS = Number(SEARCH.get("historyBudget")) || 6000;
const PRE_SPEECH_SAMPLES = 1200; // keep 50 ms before her first sound
/** @type {any} */ let model = null;
/** @type {any} */ let session = null;
let system = "";
/** @type {any[]} */ let history = [];
let idle = Promise.resolve();
/** @type {AbortController | null} */ let rereading = null;
/** @type {Map<number, AbortController>} */
const running = new Map();
const simulateRtf = Number(SEARCH.get("simulateRtf")) || 0;
const holdLoadMs = Number(SEARCH.get("holdLoadMs")) || 0; // testing: stretch out the load
const post = (msg, transfer = []) => self.postMessage(msg, { transfer });
// Piper TTS: native Spanish voice. `?voice=model` keeps the model's own voice; any
// other value is a Piper voice id (e.g. es_ES-sharvard-medium) + `?speaker=` (F=1).
let ttsSession = null;
let ttsWarm = Promise.resolve();
const ttsVoice = SEARCH.get("voice") ?? "es_ES-sharvard-medium";
const ttsSpeaker = Number(SEARCH.get("speaker")) || 1;
if (ttsVoice !== "model") {
post({ kind: "progress", status: "Preparando voz" });
ttsWarm = (async () => {
try {
const lib = await import("../vendor/piper/piper-tts-web.js");
ttsSession = await lib.TtsSession.create({ voiceId: ttsVoice, speakerId: ttsSpeaker, logger: (m) => console.log("tts:", m) });
post({ kind: "progress", status: "Voz lista" });
} catch (err) {
console.warn("voice load failed:", err);
}
})();
void ttsWarm;
}
/** Split off a sentence boundary at the front of `pendingText`. */
function takeSentence(buf) {
const m = buf.match(/[\s\S]*?[.!?…](\s|$)/);
return m ? m[0] : null;
}
function decodeWav(buf) {
const dv = new DataView(buf);
let off = 12, fmt = null, data = null;
while (off < buf.byteLength - 8) {
const id = String.fromCharCode(dv.getUint8(off), dv.getUint8(off + 1), dv.getUint8(off + 2), dv.getUint8(off + 3));
const size = dv.getUint32(off + 4, true);
if (id === "fmt ") fmt = { comp: dv.getUint16(off + 8, true), ch: dv.getUint16(off + 10, true), rate: dv.getUint32(off + 12, true), bits: dv.getUint16(off + 22, true) };
else if (id === "data") data = { off: off + 8, size };
off += 8 + size + (size % 2);
}
if (!fmt || !data) throw new Error("wav malformed");
let pcm;
if (fmt.bits === 16) {
const n = Math.floor(data.size / 2), out = new Float32Array(n);
for (let k = 0; k < n; k++) out[k] = dv.getInt16(data.off + k * 2, true) / 32768;
pcm = out;
} else if (fmt.bits === 32 && fmt.comp === 3) {
const n = Math.floor(data.size / 4), out = new Float32Array(n);
for (let k = 0; k < n; k++) out[k] = dv.getFloat32(data.off + k * 4, true);
pcm = out;
} else throw new Error(`unsupported wav bit depth ${fmt.bits}`);
return { pcm, rate: fmt.rate };
}
function trimSilence(x, rate, padSec) {
const th = 0.012;
let a = 0, b = x.length - 1;
while (a < x.length && Math.abs(x[a]) < th) a++;
while (b > a && Math.abs(x[b]) < th) b--;
const pad = Math.round(rate * padSec);
return x.slice(Math.max(0, a - pad), Math.min(x.length, b + 1 + pad));
}
function resample(x, rateIn, rateOut) {
if (!rateOut || rateIn === rateOut) return x;
const ratio = rateIn / rateOut, n = Math.max(1, Math.round(x.length / ratio)), out = new Float32Array(n);
for (let k = 0; k < n; k++) {
const p = k * ratio, j = Math.floor(p);
out[k] = x[j] + (x[Math.min(j + 1, x.length - 1)] - x[j]) * (p - j);
}
return out;
}
/** Synthesize one sentence with Piper and hand it to the main thread like model audio. */
async function speakAndPost(id, text) {
if (!text.trim()) return 0;
try {
if (!ttsSession) {
console.warn("no voice yet, staying silent for:", text);
return 0;
}
const { pcm, rate } = decodeWav(await (await ttsSession.predict(text)).arrayBuffer());
let samples = trimSilence(pcm, rate, 0.12);
samples = resample(samples, rate, model.audio_sampling_rate);
if (samples.length > 16) post({ kind: "audio", id, samples }, [samples.buffer]);
return samples.length;
} catch (err) {
console.warn("tts synthesis failed:", err);
return 0;
}
}
function messages() {
return [{ role: "system", content: system }, ...history];
}
/** Forget the oldest user turn and everything answering it. */
function dropOldestExchange() {
const next = history.findIndex((m, i) => i > 0 && m.role === "user");
if (next < 0) return false;
history.splice(0, next);
return true;
}
/** Past the budget half the history goes at once, so the next turn's one re-prefill buys many cheap turns. */
function compactHistory() {
const exchanges = history.filter((m) => m.role === "user").length;
for (let i = 0; i < Math.floor(exchanges / 2); i++) dropOldestExchange();
}
/** The reply as history keeps it: what she said, not the sound of it. */
function textOnly(message) {
if (!Array.isArray(message.content)) return message;
return { ...message, content: message.content.filter((c) => c.type === "text").map((c) => c.text).join("") };
}
/**
* Settle the cache while she's idle, re-reading a trimmed history and rewinding
* past her audio, so nobody waits for it on their next turn. Stops short of the last reply: the cache can only rewind to its end,
* and this prompt's end must still be a prefix of the next one.
*/
async function reread(signal) {
const stream = session.generate(messages().slice(0, -1), { generation_mode: "text", max_new_tokens: 1, signal });
for await (const _ of stream);
await stream.result;
}
async function load({ model: modelId, system: sys }) {
system = sys;
let lastErr;
for (let attempt = 1; attempt <= LOAD_ATTEMPTS; attempt++) {
try {
if (!navigator.gpu) {
throw Object.assign(new Error("This browser doesn't support WebGPU."), { name: "WebGPUUnavailable" });
}
post({ kind: "progress", status: "Checking your GPU" });
const check = await Model.checkAvailability(modelId);
if (!check.ok && !/No WebGPU adapter/i.test(check.reason ?? "")) {
throw Object.assign(new Error(check.reason), { name: "WebGPUUnavailable" });
}
model = await Model.from_pretrained(modelId, {
onProgress: (p) => {
if (p.kind === "bytes" && p.total > 0) post({ kind: "progress", loaded: p.loaded, total: p.total });
else if (p.status === "ready") post({ kind: "progress", fraction: 1 });
else post({ kind: "progress", status: p.message ?? p.status });
},
});
session = model.create_session();
post({ kind: "progress", status: "Getting ready" });
await session.prime(messages()).result;
if (holdLoadMs) await new Promise((r) => setTimeout(r, holdLoadMs));
post({ kind: "ready", device: model.deviceInfo(), sampleRate: model.audio_sampling_rate });
return;
} catch (err) {
lastErr = err;
console.warn(`model load attempt ${attempt} failed:`, err);
const adapterGlitch = /No WebGPU adapter/i.test(String(err?.message));
if (adapterGlitch && attempt === LOAD_ATTEMPTS) lastErr = Object.assign(new Error(err.message), { name: "WebGPUUnavailable" });
// A multi-gigabyte download over a real connection will sometimes drop a request. That's worth another go, after a growing pause; what's already cached isn't fetched again.
const networkGlitch = /failed to fetch|networkerror|network error|load failed|err_/i.test(String(err?.message));
if ((!adapterGlitch && !networkGlitch) || attempt === LOAD_ATTEMPTS) break;
await new Promise((r) => setTimeout(r, LOAD_RETRY_MS * 2 ** (attempt - 1)));
}
}
post({ kind: "failed", name: lastErr?.name ?? "Error", message: lastErr?.message ?? String(lastErr) });
}
/** One user turn: spoken (audio) or typed (text). */
async function respond({ id, audio, text }) {
if (!audio && !text) {
post({ kind: "finished", id, reason: "noop" });
return;
}
const controller = new AbortController();
running.set(id, controller);
const { signal } = controller;
rereading?.abort();
const prior = idle;
/** @type {() => void} */ let release = () => {};
idle = new Promise((r) => (release = r));
let reason = "stop";
let usage = null;
let generated = 0;
let error;
try {
await prior;
signal.throwIfAborted();
history.push(audio ? { role: "user", content: [{ type: "audio", audio }] } : { role: "user", content: text });
for (;;) {
const textMode = ttsVoice !== "model" && !audio;
const stream = session.generate(
messages(),
textMode ? { generation_mode: "text", signal } : { generation_mode: "interleaved", signal },
);
let emitted = false;
let pendingText = "";
const started = performance.now();
try {
for await (const ev of stream) {
if (ev.type === "text") {
emitted = true;
if (textMode) {
pendingText += ev.text;
for (let s; (s = takeSentence(pendingText)); ) {
pendingText = pendingText.slice(s.length);
const clean = s.trim();
if (clean) {
post({ kind: "text", id, text: clean });
generated += (await speakAndPost(id, clean)) ?? 0;
}
}
} else {
post({ kind: "text", id, text: ev.text });
}
} else if (ev.type === "audio") {
emitted = true;
let samples = ev.audio.slice();
if (generated === 0) {
// Her replies open with ~0.8 s of silence. Don't make anyone wait for it.
const start = samples.findIndex((v) => Math.abs(v) > 0.01);
if (start < 0) continue;
samples = samples.slice(Math.max(0, start - PRE_SPEECH_SAMPLES));
}
generated += samples.length;
if (simulateRtf) {
const due = started + (generated / model.audio_sampling_rate / simulateRtf) * 1000;
await new Promise((r) => setTimeout(r, Math.max(0, due - performance.now())));
}
post({ kind: "audio", id, samples }, [samples.buffer]);
}
}
if (textMode && pendingText.trim()) {
const clean = pendingText.trim();
post({ kind: "text", id, text: clean });
generated += (await speakAndPost(id, clean)) ?? 0;
pendingText = "";
}
} catch (err) {
if (!emitted && err instanceof ContextOverflowError && dropOldestExchange()) continue;
throw err;
}
const result = await stream.result;
usage = result.usage;
reason = result.finish_reason;
history.push(textOnly(result.message));
break;
}
} catch (err) {
if (err?.name === "AbortError" || signal.aborted) {
reason = "cancelled";
} else {
console.error(err);
reason = "error";
error = { name: err?.name, message: err?.message ?? String(err), fatal: err instanceof DeviceLostError };
}
} finally {
running.delete(id);
post({ kind: "finished", id, reason, usage, generated, error });
if (usage && !signal.aborted) {
if (session.count_tokens(messages()).prompt_tokens > HISTORY_BUDGET_TOKENS) compactHistory();
rereading = new AbortController();
await reread(rereading.signal).catch((err) => {
if (err?.name !== "AbortError" && !rereading?.signal.aborted) console.warn("re-read failed:", err);
});
rereading = null;
}
release();
}
}
self.onmessage = ({ data }) => {
switch (data.kind) {
case "load":
void load(data);
break;
case "respond":
void respond(data);
break;
case "abort":
running.get(data.id)?.abort();
break;
case "system":
system = data.system;
break;
case "reset":
idle = idle.then(() => {
history = [];
});
break;
}
};