Spaces:
Running
Running
| // AntSpeaker (MECT-B2) speaker embeddings + Silero VAD, fully in the browser. | |
| // Embedding model on WebGPU (onnxruntime-web), Kaldi fbank front end in JS, VAD on WASM. | |
| import { fbank } from "./fbank.js"; | |
| const ORT_VERSION = "1.30.0"; | |
| const ORT_BASE = `https://cdn.jsdelivr.net/npm/onnxruntime-web@${ORT_VERSION}/dist/`; | |
| const REPO = "https://huggingface.co/mrfakename/AntSpeaker-ONNX/resolve/main/"; | |
| const VAD_URL = "https://huggingface.co/onnx-community/silero-vad/resolve/main/onnx/model.onnx"; | |
| const CACHE_NAME = "antspeaker-onnx-v1"; | |
| const PART = 8 * 1024 * 1024; | |
| const PARALLEL = 6; | |
| let ort = null, enc = null, vad = null, backend = "wasm", precision = "fp32"; | |
| const post = (type, data = {}) => self.postMessage({ type, ...data }); | |
| async function headSize(url) { | |
| const r = await fetch(url, { method: "HEAD" }); | |
| const n = Number(r.headers.get("content-length") || 0); | |
| return n; | |
| } | |
| async function fetchRange(url, start, end, onDelta) { | |
| for (let attempt = 0; ; attempt++) { | |
| try { | |
| const res = await fetch(url, { headers: { Range: `bytes=${start}-${end - 1}` } }); | |
| if (res.status !== 206 && !(res.status === 200 && start === 0)) throw new Error(`HTTP ${res.status}`); | |
| const out = new Uint8Array(end - start); | |
| const reader = res.body.getReader(); | |
| let off = 0; | |
| for (;;) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| const n = Math.min(value.byteLength, out.length - off); | |
| out.set(value.subarray(0, n), off); | |
| off += n; onDelta(n); | |
| if (off >= out.length) { reader.cancel().catch(() => {}); break; } | |
| } | |
| if (off !== out.length) throw new Error("short read"); | |
| return out; | |
| } catch (e) { | |
| if (attempt >= 3) throw e; | |
| await new Promise((r) => setTimeout(r, 400 * (attempt + 1))); | |
| } | |
| } | |
| } | |
| async function fetchFile(url, size, onBytes) { | |
| let cache = null; | |
| try { | |
| cache = await caches.open(CACHE_NAME); | |
| const hit = await cache.match(url); | |
| if (hit) { const b = new Uint8Array(await hit.arrayBuffer()); onBytes(b.byteLength, b.byteLength, true); return b; } | |
| } catch (_) { cache = null; } | |
| let got = 0, buf; | |
| if (size > PART) { | |
| buf = new Uint8Array(size); | |
| const parts = []; | |
| for (let s = 0; s < size; s += PART) parts.push([s, Math.min(size, s + PART)]); | |
| let next = 0; | |
| const lane = async () => { | |
| while (next < parts.length) { | |
| const [a, b] = parts[next++]; | |
| let pg = 0; | |
| try { buf.set(await fetchRange(url, a, b, (n) => { got += n; pg += n; onBytes(got, size); }), a); } | |
| catch (e) { got -= pg; throw e; } | |
| } | |
| }; | |
| await Promise.all(Array.from({ length: PARALLEL }, lane)); | |
| } else { | |
| const res = await fetch(url); | |
| if (!res.ok) throw new Error(`HTTP ${res.status} for ${url}`); | |
| buf = new Uint8Array(await res.arrayBuffer()); | |
| onBytes(buf.byteLength, buf.byteLength); | |
| } | |
| if (cache) { try { await cache.put(url, new Response(buf)); } catch (_) {} } | |
| return buf; | |
| } | |
| async function load({ webgpu = true, fp16 = null }) { | |
| const t0 = performance.now(); | |
| post("status", { text: "Loading runtime" }); | |
| ort = await import(ORT_BASE + "ort.webgpu.bundle.min.mjs"); | |
| ort.env.wasm.wasmPaths = ORT_BASE; | |
| ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.min(4, navigator.hardwareConcurrency || 1) : 1; | |
| let gpuOk = false, hasF16 = false; | |
| if (webgpu && self.navigator.gpu) { | |
| try { | |
| const ad = await navigator.gpu.requestAdapter(); | |
| if (ad) { gpuOk = true; hasF16 = ad.features.has("shader-f16"); } | |
| } catch (_) {} | |
| } | |
| const useF16 = fp16 === null ? gpuOk && hasF16 : fp16; | |
| precision = useF16 ? "fp16" : "fp32"; | |
| const encUrl = REPO + (useF16 ? "onnx/model_fp16.onnx" : "onnx/model.onnx"); | |
| post("status", { text: "Downloading model" }); | |
| const [encSize, vadSize] = await Promise.all([headSize(encUrl).catch(() => 0), headSize(VAD_URL).catch(() => 0)]); | |
| const prog = { enc: [0, encSize || 1], vad: [0, vadSize || 1] }; | |
| const rep = () => post("progress", { got: prog.enc[0] + prog.vad[0], total: prog.enc[1] + prog.vad[1] }); | |
| const [encBuf, vadBuf] = await Promise.all([ | |
| fetchFile(encUrl, encSize, (g, t) => { prog.enc = [g, t]; rep(); }), | |
| fetchFile(VAD_URL, vadSize, (g, t) => { prog.vad = [g, t]; rep(); }), | |
| ]); | |
| const tDl = performance.now(); | |
| post("status", { text: "Compiling" }); | |
| if (gpuOk) { | |
| try { | |
| enc = await ort.InferenceSession.create(encBuf, { executionProviders: ["webgpu"], graphOptimizationLevel: "all" }); | |
| backend = "webgpu"; | |
| } catch (e) { | |
| post("warn", { text: "WebGPU session failed, falling back to WASM: " + e.message }); | |
| enc = null; | |
| } | |
| } | |
| if (!enc) { | |
| enc = await ort.InferenceSession.create(encBuf, { executionProviders: ["wasm"] }); | |
| backend = "wasm"; | |
| } | |
| vad = await ort.InferenceSession.create(vadBuf, { executionProviders: ["wasm"] }); | |
| // warm up the usual shapes (1.5 s windows, batch 1 and 8) | |
| post("status", { text: "Warming up" }); | |
| await runEmbed([new Float32Array(24000)]); | |
| const tReady = performance.now(); | |
| post("ready", { backend, precision, downloadMs: tDl - t0, totalMs: tReady - t0 }); | |
| } | |
| // clips: Float32Array[] at 16 kHz. Equal frame counts are batched together. | |
| async function runEmbed(clips, batch = 8) { | |
| const feats = clips.map((c) => fbank(c)); | |
| const out = new Array(clips.length); | |
| const groups = new Map(); | |
| feats.forEach((f, i) => { if (!groups.has(f.frames)) groups.set(f.frames, []); groups.get(f.frames).push(i); }); | |
| for (const [T, idx] of groups) { | |
| for (let s = 0; s < idx.length; s += batch) { | |
| const ids = idx.slice(s, s + batch); | |
| const B = ids.length; | |
| const data = new Float32Array(B * T * 80); | |
| ids.forEach((id, k) => data.set(feats[id].data, k * T * 80)); | |
| const r = await enc.run({ fbank: new ort.Tensor("float32", data, [B, T, 80]) }); | |
| const e = r.embedding.data; | |
| ids.forEach((id, k) => { | |
| const v = Float32Array.from(e.subarray(k * 192, (k + 1) * 192)); | |
| let n = 0; for (const x of v) n += x * x; n = Math.sqrt(n) || 1; | |
| for (let j = 0; j < 192; j++) v[j] /= n; | |
| out[id] = v; | |
| }); | |
| } | |
| } | |
| return out; | |
| } | |
| async function runVad(pcm) { | |
| let state = new ort.Tensor("float32", new Float32Array(2 * 128), [2, 1, 128]); | |
| const sr = new ort.Tensor("int64", BigInt64Array.from([16000n]), []); | |
| const probs = new Float32Array(Math.floor(pcm.length / 512)); | |
| const inp = new Float32Array(576); | |
| for (let i = 0; i < probs.length; i++) { | |
| inp.copyWithin(0, 512, 576); // last 64 samples of previous chunk = context | |
| if (i === 0) inp.fill(0, 0, 64); | |
| inp.set(pcm.subarray(i * 512, i * 512 + 512), 64); | |
| const r = await vad.run({ input: new ort.Tensor("float32", inp.slice(), [1, 576]), state, sr }); | |
| probs[i] = r.output.data[0]; | |
| state = r.stateN; | |
| if (i % 200 === 0) post("vadProgress", { done: i, total: probs.length }); | |
| } | |
| return probs; | |
| } | |
| // onnxruntime sessions must not run concurrently: serialise all jobs. | |
| let queue = Promise.resolve(); | |
| self.onmessage = (ev) => { queue = queue.then(() => handle(ev)); }; | |
| async function handle(ev) { | |
| const { type, id } = ev.data; | |
| try { | |
| if (type === "load") await load(ev.data); | |
| else if (type === "embed") { | |
| const t0 = performance.now(); | |
| const embs = await runEmbed(ev.data.clips, ev.data.batch || 8); | |
| post("result", { id, embs, ms: performance.now() - t0 }); | |
| } else if (type === "vad") { | |
| const t0 = performance.now(); | |
| const probs = await runVad(ev.data.pcm); | |
| post("result", { id, probs, ms: performance.now() - t0 }); | |
| } | |
| } catch (e) { | |
| post("error", { id, text: e && e.message ? e.message : String(e) }); | |
| } | |
| } | |