Kokoro-82M (StyleTTS2 + iSTFTNet)
exported to Apple's Core AI .aimodel format for iOS 27 and macOS 26.
Built for OpenReader iOS, an
on-device document reader. Published because the export differs from the
community recipe in one way that matters, and because a working Core AI TTS
export was hard enough to arrive at that it is worth not repeating.
What is here
Three fixed-shape bundles with two host steps between them. Kokoro's acoustic
graph has exactly one data-dependent length — the duration-to-alignment
expansion L = sum(pred_dur) — so the model is cut where that length appears
rather than being made dynamic.
text --(misaki G2P, host)--> phoneme ids
1. kokoro_predictor.aimodel
in input_ids[1,128] int32, ref_s[1,256], attn_mask[1,128]
out duration[1,128], d[1,128,640], t_en[1,512,128]
host pred_dur = round(duration).clamp(min=1)
alignment one-hot aln[1,128,512]; frame_mask[1,512]
2. kokoro_prosody.aimodel
in d, t_en, aln, ref_s, frame_mask
out asr[1,512,512], F0[1,512], N[1,512]
host f0_upsamp -> SineGen -> STFT => har[1,22,512] (fp32; see below)
3. kokoro_vocoder.aimodel
in asr, F0, N, har, ref_s, frame_mask
out audio[1, 512*600] (24 kHz)
| bundle |
size |
kokoro_predictor.aimodel |
80 MB |
kokoro_prosody.aimodel |
36 MB |
kokoro_vocoder.aimodel |
204 MB |
Token bucket 128, frame bucket 512 — one pass renders 12.8 s of audio.
Fixed buckets are not a shortcut. nn.LSTM specialises its sequence length
under torch.export; dynamic shapes fail.
What differs from the community recipe
The bidirectional LSTMs are not unrolled.
A fused bidirectional LSTM is wrong on right-padded input: the reverse pass
starts inside the padding and accumulates state across it before reaching a
real token. The usual fix is to unroll to 128 masked cells, which is correct
and enormous — the unrolled predictor carries a 128-step loop in its graph, and
specialisation cost tracks the program rather than the weights.
This export replaces the unroll with an exact rewrite. Gather the sequence
through an index map that reverses the first n positions and leaves padding
where it is, run a unidirectional LSTM, then gather back through the same
map — the map is an involution, so the same indices undo it. Two lstm ops
instead of 256 cells' worth of arithmetic.
Measured on a physical iPhone 15 Pro (A17 Pro), iOS 27.0, from an empty
specialisation cache:
| compute unit |
unrolled |
this export |
| CPU |
96.6 s |
58.7 s |
| GPU |
69.8 s |
35.0 s |
| Neural Engine |
107.1 s |
49.6 s |
The vocoder is the control: it holds no LSTM, it is by far the largest bundle,
and it moved 1.37 s → 1.31 s. The predictor, which carried the unroll, dropped
51%.
A second, smaller difference: an nn.LSTM must not be constructed inside the
traced region. __init__ runs reset_parameters(), which puts aten.uniform
in the exported graph and coreai-torch refuses to lower it. The export calls
torch._VF.lstm directly with the weights that already exist.
Accuracy
Verified against upstream PyTorch Kokoro (KModel.forward_with_tokens) on the
same kokoro-v1_0.pth weights:
| text |
voice |
length |
magspec corr |
waveform corr |
| "Hello, world. This is a test of on device speech generation." |
af_heart |
4.50 s |
0.9994 |
0.9820 |
| "Chapter One" |
af_heart |
1.57 s |
0.9994 |
0.9917 |
| "The Metamorphosis" |
af_heart |
1.82 s |
0.9994 |
0.9896 |
| "It was a bright cold day in April, and the clocks were striking thirteen." |
bm_george |
5.25 s |
0.9993 |
0.9824 |
Run it yourself:
uv run export_kokoro.py --out-dir . --verify --voice af_heart --text "..."
Checked on macOS with cpu_only, so this validates the exported graph. The
Neural Engine computes in fp16 and has not been compared numerically.
Performance
One pass renders 12.8 s of audio, so anything faster than that beats realtime.
Same device, per-model, excluding the host steps:
| stage |
CPU |
GPU |
ANE |
| predictor |
52.2 ms |
125.2 ms |
115.7 ms |
| prosody |
21.2 ms |
102.6 ms |
91.7 ms |
| vocoder |
1225.5 ms |
2349.6 ms |
1139.4 ms |
| pipeline |
1299 ms |
2577 ms |
1347 ms |
| vs realtime |
9.9× |
5.0× |
9.5× |
The GPU is the best unit for loading and the worst for running. Since a unit is
chosen once, at specialisation, the useful assignment picks per stage:
| stage |
unit |
cold |
inference |
| predictor |
CPU |
19.1 s |
52.2 ms |
| prosody |
ANE |
15.6 s |
91.7 ms |
| vocoder |
ANE |
1.5 s |
1139.4 ms |
| 36.1 s |
1283 ms — 10.0× realtime |
Second and subsequent loads are 1–4 ms; specialisation is cached per device and
OS. Declining the GPU costs 1.1 s of first-launch time and doubles throughput.
SpecializationOptions takes a preferredComputeUnitKind per model, and these
are three separate bundles, so the split above is expressible directly.
Notes for anyone using these
- The host steps are yours to write. Alignment between stages 1 and 2 is a
one-hot expansion; between 2 and 3 it is
f0_upsamp → SineGen → STFT. The
SineGen phase accumulator needs fp32.
- A short segment costs a full bucket. A two-word heading fills the same
512 frames as a long sentence and takes the same ~1.3 s. Batch short
segments if you care.
- Kokoro prepends a sentinel. Every segment carries roughly 0.45 s of
leading sentinel audio. It is 9% of a five-second sentence and 29% of a
1.5-second heading, and it is not silence — the vocoder's upsampling
smears the first phoneme backwards into it, so trimming
durations[0]
blindly will clip the first word. This is upstream behaviour, not an export
artefact; it reproduces in PyTorch, which is why the correlations above stay
at 0.9994.
Provenance
- Weights: hexgrad/Kokoro-82M
(
kokoro-v1_0.pth), Apache-2.0.
- Export recipe derived from
coreai-model-zoo by
John Rocky, modified as described above.
export_kokoro.py is included.
- Toolchain:
coreai-core 1.0.0b2, coreai-torch 0.4.2, torch 2.9.0,
Python 3.11.
- G2P at inference time is misaki; it is
not part of these bundles.
Reproduce with:
NATIVE_BILSTM=1 uv run export_kokoro.py --out-dir .
Without NATIVE_BILSTM=1 the script produces the unrolled variant.
