msgflux/Kokoro-82M-streaming-onnx

Kokoro-82M Streaming ONNX

Fixed-bucket ONNX exports of hexgrad/Kokoro-82M for local TTS with FastKokoro and ONNX Runtime.

This release uses a b96 fixed bucket, opset 17, length-aware duration prediction, FP32 ALBERT, decoder-only FP16, standard ONNX attention, and a portable FP32 polynomial replacement for the vocoder's atan2.

Files

Output Geometry

Kokoro's native decoder tensor contains 600 samples per alignment frame. Multilingual listening tests found that 480 samples per predicted duration frame tracks the useful speech boundary more reliably; retaining the full native tail can expose stochastic vocoder noise.

• —b96: 200 alignment frames, 480 samples/frame, 108,000 output samples.
• —input_lengths <= 32 uses a 4,800-sample margin.
• —input_lengths <= 64 uses an 8,400-sample margin.
• —Longer inputs use a 12,000-sample margin.

The selection happens inside the ONNX graph from input_lengths; margin is not a runtime configuration input. Margins shift the fixed mask-position vector into negative indices instead of being added to the predicted active-sample count. input_lengths includes the start and end pad positions. Speed values from 1.0 through 2.0 are supported.

Inputs and Outputs

Runtime Compatibility

The final graph loads with ONNX Runtime 1.16.3, 1.17.3, and 1.18.1. It runs with CPU EP, CUDA EP, and TensorRT EP. The published file uses only the standard ONNX opset 17 domain, so it does not depend on the opset 24 Attention operator or com.microsoft.Attention.

TensorRT 10.11 on a GTX 1650/SM75 compiled the checkpoint into one engine with no CUDA or CPU node fallback.

Latency

Model-call latency was measured on a GTX 1650 (SM75) after five warmups and across 25 measured iterations.

TensorRT numbers are cache-hit; persist the engine and timing cache in production.

Quickstart

Install FastKokoro with an ONNX Runtime backend:

uv pip install "fastkokoro[gpu]"
# CPU-only alternative:
uv pip install "fastkokoro[cpu]"

Start the default b96 server:

fastkokoro

Generate English speech:

curl http://localhost:8880/v1/audio/speech \
  -H "Content-Type: application/json" \
  -d '{
    "model": "kokoro",
    "voice": "af_heart",
    "lang": "en-us",
    "input": "Hello! This audio was generated by FastKokoro.",
    "response_format": "wav"
  }' \
  --output speech.wav

TensorRT with CUDA and CPU fallback:

FASTKOKORO_ONNX_PROVIDERS=TensorrtExecutionProvider,CUDAExecutionProvider,CPUExecutionProvider
FASTKOKORO_ONNX_PROVIDER_OPTIONS='{"TensorrtExecutionProvider":{"trt_engine_cache_enable":"True","trt_engine_cache_path":"/models/trt-cache","trt_timing_cache_enable":"True","trt_timing_cache_path":"/models/trt-cache"}}'

Export Recipes

The model was exported with PyTorch 2.5.1, Transformers 4.48.3, ONNX 1.21.0, NumPy 1.26.4, Hugging Face Hub 0.36.2, Loguru 0.7.3, and Misaki 0.9.4. ALBERT stays FP32; use --precision decoder-fp16, not global FP16.

b96 geometry:

B=96
ALIGN=$((2 * B + 8))
TAIL_MARGIN=12000
SAMPLES=$((ALIGN * 480 + TAIL_MARGIN))
SNAPSHOT="$HOME/.cache/huggingface/hub/models--hexgrad--Kokoro-82M/snapshots/f3ff3571791e39611d31c381e3a41a3af07b4987"

Exporter command:

uv run \
  --with torch==2.5.1 \
  --with transformers==4.48.3 \
  --with onnx==1.21.0 \
  --with numpy==1.26.4 \
  --with huggingface-hub==0.36.2 \
  --with loguru==0.7.3 \
  --with 'misaki[en]==0.9.4' \
  python scripts/export_kokoro_torch_ttfc.py \
  --kokoro-repo demo-output/reexport/hexgrad-kokoro \
  --config "$SNAPSHOT/config.json" \
  --checkpoint "$SNAPSHOT/kokoro-v1_0.pth" \
  --output "$EXPORTED_MODEL" \
  --bucket "$B" \
  --fixed-alignment-frames "$ALIGN" \
  --fixed-output-samples "$SAMPLES" \
  --output-samples-per-frame 480 \
  --output-tail-margin-samples "$TAIL_MARGIN" \
  --output-short-tail-margin-samples 4800 \
  --output-short-tail-margin-max-tokens 32 \
  --output-medium-tail-margin-samples 8400 \
  --output-medium-tail-margin-max-tokens 64 \
  --precision decoder-fp16 \
  --opset 17 \
  --legacy-export \
  --length-aware \
  --patch-fixed-lstm \
  --patch-fixed-lstm-scope duration \
  --patch-scatterless-sine-source \
  --patch-split-adain \
  --patch-albert-sdpa-bool-mask-scale \
  --fold-constant-reciprocals \
  --device cuda

Final portable graph processing:

uv run --with onnxsim==0.6.5 --with onnx==1.21.0 --with numpy==1.26.4 \
  python scripts/optimize_kokoro_onnx.py \
    --input "$EXPORTED_MODEL" \
    --output "$FINAL_MODEL" \
    --simplify \
    --atan2 portable

An experimental fusion to com.microsoft.Attention was discarded. It improved CUDA latency by only 2-3%, while TensorRT 10.11 rejected the 12 fused nodes and fragmented execution on SM75. The release optimizer does not expose or apply that transformation.

Only the duration-prediction LSTMs use real sequence lengths. Multilingual listening tests found that retaining fixed-width context in the acoustic text encoder and shared F0/noise LSTM avoids harsh output in some voices while preserving short British English endings.

Checksums

Voices and Languages

See voices.txt for the exact voice ordering.

References

• —StyleTTS 2
• —ISTFTNet