richardr1126/Kokoro-82M-CoreAI

Kokoro-82M — Core AI

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.