Koko-TTS ONNX
ONNX export of saki22/koko-tts (~24.6M param flow-matching TTS, 128 LibriTTS-R speakers, English) together with the charactr/vocos-mel-24khz vocoder (MIT), split into three graphs for running in the browser with onnxruntime-web (WebGPU).
Demo: mrfakename/koko-tts-webgpu
| File | Inputs → outputs | Notes |
|---|---|---|
onnx/text_encoder.onnx |
input_ids [1,T] int64, spk_id [1] int64 → log_dur [1,T], mu_tok [1,T,100], spk_emb [1,256] |
fp32 (RoPE positions need fp32). mu_tok is to_mu applied per token, so expanding it by durations equals the original mu. |
onnx/flow_step.onnx, onnx/flow_step_fp16.onnx |
x [1,100,F], mu [1,100,F], spk_emb [1,256], t [1], dt [1], cfg [1] → x_next [1,100,F] |
One Euler step of the Matcha-style UNet with classifier-free guidance (cond + uncond batched). fp16 variant keeps the time embedding and the state update in fp32. |
onnx/vocoder.onnx, onnx/vocoder_fp16.onnx |
mel [1,100,F] → audio [1,256*(F-1)] |
Vocos backbone + head; the iSTFT (n_fft 1024, hop 256, center padding) is folded into a ConvTranspose with an inverse-DFT basis, so it runs entirely on the GPU. |
Inference recipe
- Tokenize: collapse whitespace to a single space, map every character to its id in
tokenizer.json(<unk>=1). No BOS/EOS. - Run
text_encoder; durations =max(1, round(exp(log_dur) * length_scale)); repeat each row ofmu_tokby its duration and transpose tomu [1,100,F]. x = randn(1,100,F) * temperature; time gridu = linspace(0,1,n_steps+1),t = u + sway*(cos(pi/2*u) - 1 + u)(sway = -1).- For each step:
x = flow_step(x, mu, spk_emb, t[i], t[i+1]-t[i], cfg). audio = vocoder(x)at 24 kHz.
Defaults: 16 steps, temperature 0.667, cfg 1.5, length_scale 1.0 (<1 faster, >1 slower).
Verification
Exported graphs were checked against the PyTorch reference: patched attention max diff 7e-6, vocoder (incl. iSTFT) max diff 2e-6 vs Vocos.decode, fp16 vs fp32 mel cosine 1.000; generated speech transcribes back correctly with Whisper.
License
Apache-2.0, following the original model. The bundled Vocos vocoder weights are MIT-licensed (charactr).