Pocket-TTS-ONNX-WebGPU

Kyutai's Pocket TTS English text-to-speech model as ONNX graphs for running in a web browser with ONNX Runtime Web's WebGPU backend, as used by Quick Inference. Mono 24 kHz, made and played one 80 ms frame at a time.

Attribution

• Model: Pocket TTS by Kyutai (Manu Orsini, Simon Rouard, Gabriel De Marmiesse, Václav Volhejn, Neil Zeghidour, Alexandre Défossez), from kyutai/pocket-tts (languages/english), licensed under CC BY 4.0.
• Voice voice_alba.bin: derived from Alba MacKenna's "casual" recording in kyutai/tts-voices (alba-mackenna/casual.wav), CC BY 4.0.

These files are modified versions of the originals (see below). Kyutai does not endorse them.

Changes from the original

• Exported to ONNX from a PyTorch reimplementation, checked against Kyutai's pocket-tts package: in float32 the graphs give the same latents, end-of-speech logits and audio to about 1e-5.
• Quantized without calibration (round to nearest, symmetric, blocks of 32): the flow LM's matrix multiplications to 4-bit MatMulNBits, the text embedding table to GatherBlockQuantized, Mimi's matrix multiplications to 4-bit and its convolution weights to int8 DequantizeLinear. Activations stay float32. 64 MB in all.
• The voice is included as its prompt for the flow LM, float32 [127, 1024]: the model's learnt start-of-voice vector, then the speaker projection of the Mimi encoding of the clip.

Their outputs are close to, but not identical to, the original model's. On eight test sentences, Whisper small.en heard it with a 7.95% word error rate, the same as the original.

Files

• onnx/text_embed.onnx: token_ids int64 [1, N] → text_emb [1, N, 1024].
• onnx/flow_lm.onnx: emb [1, T, 1024], noise [1, 32] and past_key_i / past_value_i [1, 16, P, 64] (i = 0…5) → latent [1, 1, 32], eos_logit [1, 1], next_emb [1, 1, 1024] (the next frame's input) and present_key_i / present_value_i.
• onnx/mimi_decoder.onnx: latent and its streaming state (names and shapes in pocket_tts.json, all zeros at the start; each <name>_out is the next call's <name>) → pcm [1, 1, 1920].
• voice_alba.bin, bos_emb.bin (the first frame's input, float32 [1, 1, 1024]).
• pocket_tts.json: constants (end-of-speech threshold −4.0, at least 6 frames, temperature 0.3, 5 ms fade-in) and Mimi's state.
• tokenizer.json: the original's vocabulary.

Generation

1. Prepare the text as Pocket TTS does (no number normalization): straight apostrophes, single spaces, a capital first letter, and a final . unless it already ends a sentence. Split it into chunks of whole sentences of at most 50 tokens. Frames to keep after the end: 5 for four words or fewer, else 3. Frame budget: ⌈(tokens / 3 + 2) × 12.5⌉.
2. Run flow_lm once on the voice (emb = the voice, zero noise, empty cache) and keep its cache.
3. Per chunk: run flow_lm on text_embed(tokens) with zero noise and a copy of the voice cache. Then, from bos_emb, one frame at a time: noise from a Gaussian with standard deviation √0.3, run flow_lm, feed next_emb back. The end is the first frame from the sixth on whose eos_logit exceeds −4.0; stop at end + frames to keep (that frame is not kept), or at the budget. Decode each kept latent with mimi_decoder and play it; fade the chunk's first 5 ms in.

Licence

CC BY 4.0 (LICENSE), as the original model and voice.