import * as ort from './ort/ort.mjs'; import { F5Engine } from './engine/f5engine.js'; const SAMPLE_RATE = 24000; const MEL_BINS = 100; const CFG_STRENGTH = 2; let cfgStrength = CFG_STRENGTH; const N_FFT = 1024; const HOP = 256; const SPEC_BINS = N_FFT / 2 + 1; const DEFAULT_BASE = '/models/f5tts'; const DEFAULT_STEPS = 32; const MAX_REF_FRAMES = 768; const MIN_GEN_FRAMES = 16; const MAX_GEN_FRAMES = 640; const TARGET_RMS = 0.1; let refTextEncoder = null; let engine = null; let engineDevice = null; ort.env.wasm.wasmPaths = '/ort/'; ort.env.wasm.numThreads = Math.min(8, Math.max(1, (self.navigator?.hardwareConcurrency || 4) - 1)); ort.env.logLevel = 'warning'; let base = DEFAULT_BASE; let manifest = null; let ditSession = null; let vocoderSession = null; let tuning = null; let vocab = null; let pinyin = null; let caps = null; const voiceCache = new Map(); const post = (msg, transfer) => self.postMessage(msg, transfer || []); const reply = (msg, nonce) => post({ ...msg, __nonce: nonce }); function progress(stage, data = {}) { post({ type: 'progress', stage, ...data }); } async function fetchJson(url) { const res = await fetch(url); if (!res.ok) throw new Error(`asset missing: ${url} (${res.status})`); return res.json(); } async function probeCaps() { if (caps) return caps; const result = { webgpu: false }; try { const gpu = navigator.gpu; if (!gpu) return (caps = result); const adapter = await gpu.requestAdapter({ powerPreference: 'high-performance' }); if (!adapter) return (caps = result); const info = adapter.info || {}; const has = (f) => { try { return adapter.features.has(f); } catch { return false; } }; caps = { webgpu: true, vendor: info.vendor || '', architecture: info.architecture || '', description: info.description || '', isFallback: adapter.isFallbackAdapter === true, shaderF16: has('shader-f16'), subgroups: has('subgroups'), timestampQuery: has('timestamp-query'), crossOriginIsolated: self.crossOriginIsolated === true, numThreads: ort.env.wasm.numThreads, }; } catch { caps = result; } return caps; } async function loadTuning() { if (tuning) return tuning; const defaults = { llm: { maxTokens: 256, temperature: 0.7, prefillChunk: 512 }, tts: { flowSteps: DEFAULT_STEPS }, session: { preferPrecision: 'auto' } }; try { const res = await fetch('/tuning.json'); const type = res.headers.get('content-type') || ''; tuning = res.ok && type.includes('json') ? { ...defaults, ...(await res.json()) } : defaults; } catch { tuning = defaults; } return tuning; } // Download a graph with byte progress (falls back to plain fetch when streams are unavailable). async function fetchModel(url, label) { const res = await fetch(url); if (!res.ok) throw new Error(`failed to download ${label}: ${url} (${res.status})`); const total = Number(res.headers.get('content-length') || 0); if (!res.body || !res.body.getReader) return { data: new Uint8Array(await res.arrayBuffer()), total: 0 }; const reader = res.body.getReader(); const chunks = []; let loaded = 0; let lastPost = 0; for (;;) { const { done, value } = await reader.read(); if (done) break; chunks.push(value); loaded += value.length; const now = performance.now(); if (now - lastPost > 100) { progress('download', { label, loaded, total }); lastPost = now; } } progress('download', { label, loaded, total }); const data = new Uint8Array(loaded); let offset = 0; for (const chunk of chunks) { data.set(chunk, offset); offset += chunk.length; } return { data, total: loaded }; } async function createSession(url, label, executionProviders) { const started = performance.now(); progress('session', { label, message: `fetching ${label}` }); const { data, total } = await fetchModel(url, label); progress('session', { label, message: `initializing ${label} (${(total / 1e6).toFixed(0)} MB)`, loaded: total, total }); const session = await ort.InferenceSession.create(data, { executionProviders, graphOptimizationLevel: executionProviders.includes('webgpu') ? 'basic' : 'all', }); progress('session', { label, message: `${label} ready in ${((performance.now() - started) / 1000).toFixed(1)}s`, done: true }); return session; } async function initEngine() { const want = (tuning?.tts?.backend ?? 'engine') === 'engine'; if (!want) return; const info = await probeCaps(); if (!info.webgpu || !info.shaderF16) return; try { const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' }); if (!adapter) return; const feats = ['shader-f16']; let sgmat = false; try { sgmat = adapter.features.has('chromium-experimental-subgroup-matrix'); } catch {} if (sgmat) feats.push('chromium-experimental-subgroup-matrix'); const device = await adapter.requestDevice({ requiredFeatures: feats, requiredLimits: { maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize, maxBufferSize: adapter.limits.maxBufferSize, maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize, }, }); engineDevice = device; engine = await F5Engine.create(device, { gemm: tuning?.tts?.gemm || (sgmat ? 'sgmat' : 'simt'), branches: 1, metaUrl: `${base}/f5w_q4.json`, binUrl: `${base}/f5w_q4.bin`, }); progress('session', { message: 'DiT kernels ready (WebGPU)' }); } catch (err) { engine = null; engineDevice = null; progress('fallback', { message: `engine unavailable, using ORT: ${err?.message || err}` }); } } async function loadGraphs() { if (engine) { if (!vocoderSession) { const g = manifest.graphs || {}; vocoderSession = await createSession(`${base}/${g.vocoder_fp16 || g.vocoder}`, 'vocoder · WebGPU', ['webgpu', 'wasm']); } return; } if (ditSession && vocoderSession) return; const graphs = manifest.graphs || {}; const fp16Available = !!(graphs.dit_step_fp16 && graphs.vocoder_fp16); const info = await probeCaps(); const tryGpu = info.webgpu && fp16Available; if (tryGpu) { progress('session', { message: 'WebGPU detected — loading fp16 graphs' }); try { ditSession = await createSession(`${base}/${graphs.dit_step_fp16}`, 'dit · fp16 · WebGPU', ['webgpu', 'wasm']); } catch (err) { progress('fallback', { message: `webgpu dit failed, falling back to wasm: ${err?.message || err}` }); } try { vocoderSession = await createSession(`${base}/${graphs.vocoder_fp16}`, 'vocoder · fp16 · WebGPU', ['webgpu', 'wasm']); } catch (err) { progress('fallback', { message: `webgpu vocoder failed, falling back to wasm: ${err?.message || err}` }); } } if (!ditSession) { ditSession = await createSession(`${base}/${graphs.dit_step}`, 'dit · int8 · WASM', ['wasm']); } if (!vocoderSession) { vocoderSession = await createSession(`${base}/${graphs.vocoder}`, 'vocoder · int8 · WASM', ['wasm']); } } async function warmup() { const frames = tuning?.tts?.warmupFrames ?? manifest.vocoder_chunk_frames ?? 375; if (!engine) await runAndDispose(ditSession, () => ditInputs(null, frames, 16, null, null, null)); await runAndDispose(vocoderSession, () => ({ mel: f16Tensor(new Float32Array(MEL_BINS * frames), [1, MEL_BINS, frames]) })); } async function runAndDispose(session, feeds) { const inputs = feeds(); try { const outputs = await session.run(inputs); disposeAll(outputs); } finally { disposeAll(inputs); } } // --------------------------------------------------------------------------- // Text → token ids (ASCII + numbered pinyin, no CJK glyphs in the vocab) // --------------------------------------------------------------------------- const isCjk = (ch) => { const cp = ch.codePointAt(0); return cp >= 12544 && cp <= 40959; }; const asciiWord = /[A-Za-z0-9]+/y; const QUOTE_MAP = { ';': ',', ',': ',', '“': '"', '”': '"', '‘': "'", '’': "'" }; function splitText(text, pinyinMap) { const normalized = Array.from(String(text ?? ''), (ch) => QUOTE_MAP[ch] ?? ch).join(''); const tokens = []; let i = 0; while (i < normalized.length) { if (isCjk(normalized[i])) { while (i < normalized.length && isCjk(normalized[i])) { const ch = normalized[i]; tokens.push(' '); tokens.push(pinyinMap[ch] ?? ch); i += 1; } continue; } asciiWord.lastIndex = i; const word = asciiWord.exec(normalized); if (word && word.index === i && word[0].length > 1) { const prev = tokens.length ? tokens[tokens.length - 1] : ''; if (tokens.length && !` :'"`.includes(prev)) tokens.push(' '); tokens.push(...Array.from(word[0])); i += word[0].length; continue; } tokens.push(normalized[i]); i += 1; } return tokens; } async function loadVocab() { if (!vocab) vocab = await fetchJson(`${base}/vocab.json`); return vocab; } async function loadPinyin() { if (!pinyin) pinyin = await fetchJson(`${base}/pinyin.json`); return pinyin; } async function tokenize(text) { const [vocabData, pinyinMap] = await Promise.all([loadVocab(), loadPinyin()]); const tokens = splitText(text, pinyinMap); const ids = new Int32Array(tokens.length); const unknown = []; for (let i = 0; i < tokens.length; i++) { const token = tokens[i]; if (token === ' ' || vocabData.index[token] !== undefined) { ids[i] = vocabData.index[token]; continue; } if (isCjk(token[0]) && !unknown.includes(token[0])) unknown.push(token[0]); ids[i] = 0; } return { ids, unknown }; } // --------------------------------------------------------------------------- // Voice reference (mel spectrogram, fp16) // --------------------------------------------------------------------------- async function loadVoice(voiceId) { const id = voiceId || 'xiaochun'; if (voiceCache.has(id)) return voiceCache.get(id); let meta; let raw; try { meta = await fetchJson(`/refs/${id}.json`); raw = await (await fetch(`/${meta.mel_file}`)).arrayBuffer(); } catch { meta = await fetchJson(`${base}/refs/${id}.json`); raw = await (await fetch(`${base}/${meta.mel_file}`)).arrayBuffer(); } const halves = new Float16Array(raw); const mel = new Float32Array(halves.length); for (let i = 0; i < halves.length; i++) mel[i] = halves[i]; const [bins, frames] = meta.mel_shape; const voice = { id: meta.id, nameZh: meta.name_zh, mel, bins, frames, refText: meta.ref_text, seconds: meta.seconds }; voiceCache.set(id, voice); return voice; } // --------------------------------------------------------------------------- // Flow matching (Euler + sway schedule + classifier-free guidance) // --------------------------------------------------------------------------- function flowSchedule(steps, sway = -1) { const t = new Float64Array(steps + 1); for (let i = 0; i <= steps; i++) t[i] = i / steps; if (sway != null) { for (let i = 0; i <= steps; i++) t[i] += sway * (Math.cos((Math.PI / 2) * t[i]) - 1 + t[i]); } return t; } function gaussianNoise(frames, bins, seed = 1234) { let state = seed >>> 0; const next = () => { state ^= state << 13; state >>>= 0; state ^= state >> 17; state ^= state << 5; state >>>= 0; return state / 4294967296; }; const out = new Float32Array(frames * bins); for (let i = 0; i < out.length; i += 2) { const u = Math.max(next(), 1e-7); const v = next(); const r = Math.sqrt(-2 * Math.log(u)); out[i] = r * Math.cos(2 * Math.PI * v); if (i + 1 < out.length) out[i + 1] = r * Math.sin(2 * Math.PI * v); } return out; } function flowSample(stepFn, frames, bins, { steps = DEFAULT_STEPS, sway = -1, seed = 1234, onStep } = {}) { const schedule = flowSchedule(steps, sway); const x = gaussianNoise(frames, bins, seed); const dt = new Float32Array(steps); for (let i = 0; i < steps; i++) dt[i] = schedule[i + 1] - schedule[i]; return (async () => { for (let i = 0; i < steps; i++) { const velocity = await stepFn(x, schedule[i]); for (let j = 0; j < x.length; j++) x[j] += dt[i] * velocity[j]; onStep?.(i + 1, steps); } return x; })(); } // --------------------------------------------------------------------------- // Model inference plumbing (fp16 tensors) // --------------------------------------------------------------------------- function f16Array(src) { const out = new Float16Array(src.length); for (let i = 0; i < src.length; i++) out[i] = src[i]; return out; } function f16Tensor(data, dims) { return new ort.Tensor('float16', f16Array(data), dims); } function toFloat32(data) { if (data instanceof Float16Array) return Float32Array.from(data); if (data instanceof Uint16Array) { const out = new Float32Array(data.length); for (let i = 0; i < data.length; i++) { const bits = data[i]; const sign = (bits & 0x8000) ? -1 : 1; const exponent = (bits >> 10) & 0x1f; const mantissa = bits & 0x3ff; if (exponent === 0) out[i] = sign * Math.pow(2, -14) * (mantissa / 1024); else if (exponent === 0x1f) out[i] = mantissa ? NaN : sign * Infinity; else out[i] = sign * Math.pow(2, exponent - 15) * (1 + mantissa / 1024); } return out; } return data; } function ditInputs(x, frames, textLen, cond, textIds, time) { const condData = cond ?? new Float32Array(frames * MEL_BINS); const ids = BigInt64Array.from({ length: textLen }, (_, i) => BigInt(textIds?.[i] ?? -1)); return { x: f16Tensor(x ?? new Float32Array(frames * MEL_BINS), [1, frames, MEL_BINS]), cond: f16Tensor(condData, [1, frames, MEL_BINS]), text_ids: new ort.Tensor('int64', ids, [1, textLen]), time: f16Tensor(Float32Array.from([time ?? 0]), [1]), mask: new ort.Tensor('bool', new Uint8Array(frames).fill(1), [1, frames]), }; } function disposeAll(tensors) { for (const tensor of Object.values(tensors)) tensor?.dispose?.(); } function padReference(voice, refFrames, frames) { const out = new Float32Array(frames * MEL_BINS); for (let f = 0; f < refFrames; f++) { for (let b = 0; b < MEL_BINS; b++) { out[f * MEL_BINS + b] = voice.mel[b * voice.frames + f]; } } return out; } async function ditStep(x, time, voice, textIds, textLen, frames, refFrames) { const condFeeds = ditInputs(x, frames, textLen, padReference(voice, refFrames, frames), textIds, time); const gen = await ditSession.run(condFeeds); const pred = toFloat32(gen.velocity.data); disposeAll(gen); disposeAll(condFeeds); if (cfgStrength === 0) return pred; const nullFeeds = ditInputs(x, frames, textLen, null, null, time); const uncond = await ditSession.run(nullFeeds); const nullPred = toFloat32(uncond.velocity.data); disposeAll(uncond); disposeAll(nullFeeds); const velocity = new Float32Array(pred.length); for (let i = 0; i < pred.length; i++) velocity[i] = pred[i] + (pred[i] - nullPred[i]) * cfgStrength; return velocity; } // --------------------------------------------------------------------------- // Vocoder tail: magnitude/phase planes → inverse STFT (1024/256, hann, centered) // --------------------------------------------------------------------------- let twiddleCos = null; let twiddleSin = null; let hannWindow = null; function initIstftTables() { if (twiddleCos) return; twiddleCos = new Float64Array(N_FFT / 2); twiddleSin = new Float64Array(N_FFT / 2); for (let i = 0; i < N_FFT / 2; i++) { twiddleCos[i] = Math.cos((2 * Math.PI * i) / N_FFT); twiddleSin[i] = Math.sin((2 * Math.PI * i) / N_FFT); } hannWindow = new Float64Array(N_FFT); for (let i = 0; i < N_FFT; i++) hannWindow[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / N_FFT); } function ifft(re, im, out) { initIstftTables(); const reFull = new Float64Array(N_FFT); const imFull = new Float64Array(N_FFT); for (let i = 0; i < SPEC_BINS; i++) { reFull[i] = re[i]; imFull[i] = im[i]; } for (let i = SPEC_BINS; i < N_FFT; i++) { reFull[i] = reFull[N_FFT - i]; imFull[i] = -imFull[N_FFT - i]; } for (let i = 1, j = 0; i < N_FFT; i++) { let bit = N_FFT >> 1; for (; j & bit; bit >>= 1) j ^= bit; j ^= bit; if (i < j) { let tmp = reFull[i]; reFull[i] = reFull[j]; reFull[j] = tmp; tmp = imFull[i]; imFull[i] = imFull[j]; imFull[j] = tmp; } } for (let size = 2; size <= N_FFT; size <<= 1) { const half = size >> 1; const step = N_FFT / size; for (let start = 0; start < N_FFT; start += size) { for (let k = 0; k < half; k++) { const angle = k * step; const wr = twiddleCos[angle]; const wi = twiddleSin[angle]; const a = start + k; const b = a + half; const tr = reFull[b] * wr - imFull[b] * wi; const ti = reFull[b] * wi + imFull[b] * wr; reFull[b] = reFull[a] - tr; imFull[b] = imFull[a] - ti; reFull[a] += tr; imFull[a] += ti; } } } for (let i = 0; i < N_FFT; i++) out[i] = reFull[i] / N_FFT; } function istft(planes, frames) { initIstftTables(); const length = (frames - 1) * HOP + N_FFT; const acc = new Float64Array(length); const norm = new Float64Array(length); const re = new Float64Array(SPEC_BINS); const im = new Float64Array(SPEC_BINS); const frame = new Float64Array(N_FFT); for (let t = 0; t < frames; t++) { for (let b = 0; b < SPEC_BINS; b++) { const mag = Math.min(Math.exp(planes[b * frames + t]), 100); const phase = planes[SPEC_BINS * frames + b * frames + t]; re[b] = mag * Math.cos(phase); im[b] = mag * Math.sin(phase); } ifft(re, im, frame); const offset = t * HOP; for (let i = 0; i < N_FFT; i++) { acc[offset + i] += frame[i] * hannWindow[i]; norm[offset + i] += hannWindow[i] * hannWindow[i]; } } const trim = N_FFT >> 1; const out = new Float32Array(length - 2 * trim); for (let i = 0; i < out.length; i++) { const d = norm[i + trim]; out[i] = d > 1e-11 ? acc[i + trim] / d : 0; } return out; } async function decodeMel(mel, frames) { const chunkFrames = manifest.vocoder_chunk_frames || 375; const chunks = []; for (let start = 0; start < frames; start += chunkFrames) { const len = Math.min(chunkFrames, frames - start); const input = new Float32Array(MEL_BINS * len); for (let b = 0; b < MEL_BINS; b++) { for (let f = 0; f < len; f++) input[b * len + f] = mel[(start + f) * MEL_BINS + b]; } const feeds = { mel: f16Tensor(input, [1, MEL_BINS, len]) }; const outputs = await vocoderSession.run(feeds); const planes = toFloat32(outputs.stft_planes.data); disposeAll(outputs); disposeAll(feeds); chunks.push(istft(planes, len)); } const total = chunks.reduce((sum, chunk) => sum + chunk.length, 0); const pcm = new Float32Array(total); let offset = 0; for (const chunk of chunks) { pcm.set(chunk, offset); offset += chunk.length; } return pcm; } // --------------------------------------------------------------------------- // Message handling // --------------------------------------------------------------------------- async function init(message) { if (message.base) base = message.base; await loadTuning(); manifest = await fetchJson(`${base}/manifest.json`); await initEngine(); await loadGraphs(); progress('warmup', { message: 'warming up kernels' }); await warmup(); await loadVoice('xiaochun').catch(() => {}); post({ type: 'caps', caps: await probeCaps() }); } const refIdCache = new Map(); async function getRefIds(voice) { if (!voice.refText) return []; if (!refIdCache.has(voice.id)) { const { ids } = await tokenize(voice.refText); refIdCache.set(voice.id, Array.from(ids)); } return refIdCache.get(voice.id); } async function synthesize(message) { if ((!engine && !ditSession) || !vocoderSession) throw new Error('worker not initialized'); const started = performance.now(); const { ids: genIds, unknown } = await tokenize(message.text); if (!genIds.length) throw new Error(`no vocab tokens for: ${String(message.text).slice(0, 20)}`); const voice = await loadVoice(message.voiceId); const refIds = await getRefIds(voice); const ids = [...refIds, ...genIds]; const refFrames = Math.min(voice.frames, MAX_REF_FRAMES); refTextEncoder ??= new TextEncoder(); const refBytes = Math.max(refTextEncoder.encode(voice.refText || '').length, 1); const genBytes = Math.max(refTextEncoder.encode(message.text).length, 1); let genFrames = Math.round((refFrames / refBytes) * genBytes); genFrames = Math.max(Math.min(genFrames, MAX_GEN_FRAMES), MIN_GEN_FRAMES); const frames = refFrames + genFrames; const textLen = Math.min(ids.length, frames); cfgStrength = tuning?.tts?.cfgStrength ?? manifest.sampler?.cfg_strength ?? CFG_STRENGTH; const steps = tuning?.tts?.flowSteps ?? manifest.sampler?.steps ?? DEFAULT_STEPS; progress('synth', { message: 'flow matching', step: 0, steps }); let mel; if (engine) { engine.setup(frames); await engine.setCond(padReference(voice, refFrames, frames), ids); engine.setX(gaussianNoise(frames, MEL_BINS, message.seed ?? 1234)); await engineDevice.queue.onSubmittedWorkDone(); const sched = flowSchedule(steps, manifest.sampler?.sway_sampling_coef ?? -1); for (let i = 0; i < steps; i++) { progress('synth', { step: i + 1, steps, message: `flow step ${i + 1}/${steps}` }); await engine.step(sched[i], sched[i + 1] - sched[i]); } mel = await engine.readX(); } else { mel = await flowSample( (x, time) => ditStep(x, time, voice, ids, textLen, frames, refFrames), frames, MEL_BINS, { steps, sway: manifest.sampler?.sway_sampling_coef ?? -1, seed: message.seed ?? 1234, onStep: (step, total) => progress('synth', { step, steps: total, message: `flow step ${step}/${total}` }), }, ); } progress('synth', { message: 'vocoding', step: steps, steps }); const pcm = await decodeMel(mel.subarray(refFrames * MEL_BINS), genFrames); let sum = 0; for (let i = 0; i < pcm.length; i++) sum += pcm[i] * pcm[i]; const rms = Math.sqrt(sum / Math.max(pcm.length, 1)); if (rms > 1e-6) { const scale = TARGET_RMS / rms; for (let i = 0; i < pcm.length; i++) pcm[i] *= scale; } post({ type: 'audio', pcm, sampleRate: SAMPLE_RATE, seconds: pcm.length / SAMPLE_RATE, __nonce: message.__nonce }, [pcm.buffer]); reply({ type: 'done', steps, frames, genFrames, msTotal: performance.now() - started, voiceId: voice.id, unknownChars: unknown, }, message.__nonce); } self.onmessage = async (event) => { const message = event.data || {}; const nonce = message.__nonce; try { if (message.type === 'init') { await init(message); reply({ type: 'ready', refs: [{ id: 'xiaochun', nameZh: '小莼(纯儿)' }] }, nonce); return; } if (message.type === 'synthesize') { await synthesize(message); return; } } catch (err) { reply({ type: 'error', message: String(err?.message || err) }, nonce); } };