Running a local neural TTS model, and the artifact that changed the runtime

Running Kokoro-82M locally and in-process for on-device voice: no cloud, no native ML runtime. Based on a shipping integration, including the finding that made it change runtimes.

The headline decisions:

• Kokoro-82M (StyleTTS2 family), Apache 2.0, so commercial use is permitted. Around 82M parameters, and it punches well above that.
• HuggingFace Candle, pure Rust — not ONNX. The reason is a specific audio artifact, and it is the crux of this document.
• CPU-only, made fast with platform BLAS plus two hot-path fixes, roughly 3x combined, reaching about 2.3x realtime for five seconds of audio.
• espeak-free grapheme-to-phoneme with a dictionary fallback for out-of-vocabulary words.

The model, concretely

A StyleTTS2-family graph: a PL-BERT phoneme encoder, a linear projection into model hidden size, a prosody predictor for duration, pitch and noise, a text encoder, and an ISTFTNet decoder feeding an inverse STFT with a harmonic sine-source module. Output is 24 kHz.

Three files ship: a config carrying hyperparameters and a 178-token phoneme vocabulary, a PyTorch checkpoint of roughly 327 MB, and per-voice style packs of about half a megabyte each.

Voice packs are tiny and swappable, and this shapes the architecture. Each pack is a [512, 256] tensor — 512 context positions by 128 decoder-style plus 128 predictor-style dimensions. At synthesis you select the row by input length, then split it into the two style vectors.

The consequence worth internalising: the checkpoint is voice-independent. Switching voice selects a different small tensor; it does not reload the model. Load once and keep it warm.

Why not ONNX — the metallic artifact

This is the reason the document exists.

The integration originally ran on a pure-Rust ONNX runtime, from a custom epsilon-patched export, and carried a subtle metallic quality in the voice.

The root cause was not a single bad operator. It was a distributed numerical difference: the two runtimes diverged in f32 across the decoder’s harmonic source module, tilting high-frequency content upward. No individual operation was outside tolerance. The error accumulated.

Loading the same weights into a hand-ported Candle graph rendered cleanly and matched the reference implementation by ear, with high-frequency divergence around 0.0097. The fix was not to patch an operator or tune an epsilon — it was to change runtime so the original weights computed faithfully. The entire patched-export apparatus was then retired.

Two takeaways that generalise well beyond speech synthesis:

1. Numerical divergence between ML runtimes can be perceptible and distributed. A model can be correct operator-by-operator within tolerance and still be wrong in output, because small per-operation differences accumulate through a sensitive module — here a harmonic oscillator feeding an inverse STFT.

2. Validate by end-to-end perceptual quality on real input, never by runtime-versus-runtime tensor parity on a dummy input. A zeros-input parity check would have passed. The bug was only findable by listening.

The second point is the transferable one. Parity tests on synthetic input are reassuring and cheap, which is exactly why they get trusted past the point where they are informative.

Running it on Candle

The model graph is hand-ported and vendored in-tree rather than taken as a crate dependency, adapted from an existing Candle backend under a permissive licence with per-file attribution. Only the model code is vendored — the upstream project’s own phonemiser and its associated heavy dependencies are not pulled in.

Reading the checkpoint is the fiddly part. The published weights are a bare PyTorch pickle inside a ZIP, which Candle’s dict-only tensor loader cannot unwrap. The pickle metadata has to be parsed directly and the archive read with a ZIP crate. If you port this, budget time here specifically; it is the least interesting and most time-consuming piece.

CPU-only, on purpose

The engine pins to CPU explicitly. Candle 0.10’s Metal backend has no layer-norm kernel, and this model uses LayerNorm about thirty-one times, so it errors mid-synthesis with no runtime fallback. Rather than fight that, the port stays on CPU and takes its speed from the platform BLAS feature.

If you target a different accelerator, the calculation may differ — but verify the kernel exists on your backend and version before assuming a GPU helps. An absent kernel is a hard failure partway through, not a slow path.

Phonemes without espeak

The phonemiser runs with default features disabled, which drops its GPL fallback — that flag is the entire reason to use it, and it produces IPA directly.

Tokenisation is character-level against the config vocabulary, with unknown characters silently skipped and padding at both ends.

The out-of-vocabulary fallback is what makes espeak-free viable. Without it, unknown words get spelled letter by letter, which is audibly wrong. Detect that case and substitute a dictionary pronunciation, resolving in order: a direct dictionary hit; then a morphological suffix strip (-ing, -ed, -s, -ies, -er, -ly, -est); then a compound split where both halves are at least three characters, so masterclass becomes master plus class. Then map the dictionary’s stress digits onto IPA stress marks.

Apply user pronunciation overrides before phonemisation, so any word can be corrected without touching the fallback chain.

The call path

1. Normalise text — pronunciation rules first, then numerals, currency and ordinals into words, because the phonemiser emits nothing for digits.
2. Chunk into sentence spans. Spans beyond roughly 240 characters split further at clause then space boundaries, never mid-word. This is what gives first-audio after the first sentence rather than after the whole reply.
3. Per span: phonemise, tokenise, forward pass, producing f32 samples at 24 kHz. Zero any non-finite samples defensively.
4. Trim and rejoin. The model emits roughly 300–450 ms of trailing silence per sentence. Left in, these stack into audible ~800 ms stumbles in multi-sentence replies. Trim per span, then append one uniform short pause.
5. Encode and stream as chunk events, playing gaplessly from the first sentence.

Making CPU inference fast

Two independent wins, both worth replicating:

Platform BLAS routed matrix multiplications through the system library — roughly 2x on the decoder for a build-flag change.

Hot-path fixes in the vendored code, found by profiling, gave about 2.5x on top of that:

• LSTM. The naive port recomputed constant weight transposes every timestep and did the input projection step by step. Hoisting the transposes and batching the input projection into a single matmul took the predictor from ~1834 ms to ~61 ms, and the text encoder from ~243 ms to ~9 ms.
• STFT and inverse STFT. A hand-rolled O(N²) transform recomputed sine and cosine twiddle factors per element. Precomputing the tables once turned the hot loops into lookups.

End to end: ~6488 ms to ~2196 ms for five seconds of audio, from about 0.78x realtime to about 2.31x. The remainder is dominated by the decoder’s im2col convolutions, a known CPU limitation of the framework.

If you port a StyleTTS2-family graph, profile the LSTM and the STFT first. That is where a naive translation leaves the biggest and easiest wins, and both are mechanical fixes rather than redesigns.

Checklist

[ ] Enable the platform BLAS feature for your OS
[ ] Vendor the model code with per-file attribution
[ ] Write the checkpoint loader — expect this to be the fiddly bit
[ ] Force CPU unless you have confirmed your backend has a LayerNorm kernel
[ ] Phonemise without a GPL fallback; add a dictionary OOV shim with
    morphology and compound splitting
[ ] Chunk by sentence, trim trailing silence, insert a uniform pause, stream
[ ] Profile and fix the LSTM and STFT hot paths early
[ ] Validate by ear on real sentences, not by tensor parity on a dummy

That last line is the one that matters. It is how the metallic artifact was found, and it is how you will know a port is faithful rather than merely plausible.