Pocket-TTS-LiteRT

Pocket TTS (Kyutai, ~100M params) converted to LiteRT CompiledModel graphs for the phone GPU. Stateless graphs + host-side orchestration reproduce the reference pocket_tts pipeline: the 100M language model, the flow head and the SEANet vocoder run on the mobile GPU; the small 2-layer Mimi decoder transformer runs on CPU (placement notes below).

Real output from the phone (Pixel 8a, nothing cloud, nothing post-processed):

voice	sample
alba
marius

How it is put together

Pocket TTS is a flow-matching LM over continuous 32-dim Mimi latents: per 12.5 Hz frame a 6-layer/1024-wide causal transformer conditions a 6-block AdaLN MLP that turns one Gaussian draw into the next latent (LSD, 1 step); a 20M tiny Mimi (×16 ConvTranspose upsample + 2-layer transformer + SEANet) decodes latents to 24 kHz audio.

graph	I/O	role
pt_flowlm_fused	emb[1,1,1024] + cos/sin + mask[1,16,1,513] + packed KV [1,96,512,64] + noise[1,32] → [1,12321] = eos ∣ latent ∣ new-k ∣ new-v	one full AR frame (step + flow head) in one invocation with one readback — the variant the Android sample runs; on Mali the per-frame cost is dispatch/sync-bound, and fusing removes one invocation + three readbacks per frame
pt_flowlm_step	emb[1,1,1024] + cos/sin + mask[1,16,1,513] + packed KV [1,96,512,64] → cond, eos, new k/v	one AR step; KV cache lives on the host (split reference variant)
pt_flow_head	cond[1,1024] + noise[1,32] → latent[1,32]	LSD time embeddings (s=0, t=1) baked into the cond bias (split reference variant)
pt_mimi_dec_tx	lat[1,65,32] → feat[1,512,1024]	Mimi decoder transformer in 64-frame blocks (32-frame overlap: the 2-layer sliding-window attention has a 498-position stacked receptive field) — runs on CPU in the shipped app (see Placement)
pt_mimi_deconly	feat[1,512,4096] → audio[1,1,491520]	SEANet decoder, one-shot 256-frame window (causal ⇒ exact per frame)

Host side (a few hundred lines of Kotlin/Python, no FFT anywhere): sentencepiece unigram tokenizer, fp16 token-embedding lookup, 32→1024 input projection, RoPE cos/sin per step, KV cache + additive mask bookkeeping, Gaussian noise (std √0.3), EOS threshold −4.

The pt_voice_*.bin files are Kyutai's published per-voice prompt states repacked for the packed-KV layout (fp16). Text is chunked at ≤50 tokens along sentence boundaries, exactly like the reference implementation.

Placement and measured numbers

Every graph compiles fully on the GPU (LITERT_CL, zero CPU-fallback nodes) on both devices tried. The shipped placement still runs pt_mimi_dec_tx on CPU: on a Pixel 8a's Mali the GPU output of that one graph is audibly degraded (alba voicing HNR 0.9 dB on GPU vs 2.8 dB on CPU — CPU matches the fp32 desktop reference exactly), and requesting FP32 GPU precision does not recover it. This mirrors what the Mimi zoo module documents for its own decoder transformer, so the same split ships here: heavy compute (LM, flow head, SEANet) on GPU, the 2-layer decoder transformer on CPU. It is 7 small calls per utterance — on the Pixel the whole pipeline goes 1.03× → 1.01× real-time.

LiteRT 2.1.6, fp16 graphs, decode after generation (no streaming), app process warm:

• Pixel 8a (Tensor G3, Mali-G715), shipped placement: 8.8 s of speech in 8.75 s — ~1.0× real-time (alba); per-step cost is KV upload + dispatch overhead, not arithmetic.
• Samsung SM-S942Q (Snapdragon SM8850, Adreno), measured with the all-GPU placement and the split step+head graphs: 8.2 s in 1.63 s — 5.0× real-time; 13.0 s in 3.05 s — 4.3× (3 chunks). The decoder-transformer-on-CPU delta measured ~2% on the Pixel.

Numbers move with device, thermals and text length; treat them as one measured point, not a benchmark.

Conversion fidelity

Everything is a numerically-equivalent re-authoring except one op: erf-GELU is replaced by a fitted odd tanh-polynomial (max |gelu err| 7.1e-5, ~15× closer than the classic tanh-GELU). Measured against the eager reference (fp32 host):

• flow-LM step: cond corr 1.000000 (max|d| 3.7e-6 per step)
• flow head: max|d| 2.4e-7 (bit-exact-grade; the +noise integration is in-graph)
• Mimi dec_tx blocks vs full-sequence eager: corr 1.000000 (max|d| 5.2e-4)
• dec_tx + SEANet vs eager decode_from_latent: corr 1.000000 (max|d| 2.0e-4)
• full free-running pipeline vs eager, same noise: audio corr 0.997, EOS step identical

RoPE's interleaved pairs are de-interleaved by baking a row permutation into the QKV projection (bit-exact — q and k share the permutation, so q·k is unchanged). The KV-step FULLY_CONNECTED shapes need LiteRT ≥ 2.1.5 on Mali; this build uses 2.1.6.

Minimal usage — Python (CompiledModel)

# pip install ai-edge-litert sentencepiece numpy huggingface_hub
import numpy as np, sentencepiece as spm
from ai_edge_litert.compiled_model import CompiledModel
from huggingface_hub import hf_hub_download as dl

R = "mlboydaisuke/Pocket-TTS-LiteRT"
lm  = CompiledModel.from_file(dl(R, "pt_flowlm_step_fp16.tflite"))
hd  = CompiledModel.from_file(dl(R, "pt_flow_head_fp16.tflite"))
sp  = spm.SentencePieceProcessor(dl("kyutai/pocket-tts-without-voice-cloning",
                                    "languages/english/tokenizer.model"))
emb = np.fromfile(dl(R, "pt_embed_f16.bin"), np.float16).reshape(4001, 1024)
inw = np.fromfile(dl(R, "pt_input_linear_f32.bin"), np.float32).reshape(1024, 32)
bos = np.fromfile(dl(R, "pt_bos_input_f32.bin"), np.float32)
raw = open(dl(R, "voices/pt_voice_alba.bin"), "rb").read()
T = np.frombuffer(raw, np.int32, 1)[0]                      # voice prompt length
kv = np.frombuffer(raw, np.float16, offset=4).astype(np.float32).reshape(2, 96, T, 64)
pk = np.zeros((1, 96, 512, 64), np.float32); pk[0, :, :T] = kv[0]
pv = np.zeros((1, 96, 512, 64), np.float32); pv[0, :, :T] = kv[1]
pos, freqs = int(T), 10000.0 ** (-np.arange(32) / 32.0)
lin, lout = lm.create_input_buffers(0), lm.create_output_buffers(0)
hin, hout = hd.create_input_buffers(0), hd.create_output_buffers(0)

def step(x):
    global pos
    ang = pos * freqs
    mask = np.full((16, 513), -1e4, np.float32); mask[:, :pos] = 0; mask[:, 512] = 0
    for b, a in zip(lin, [x, np.tile(np.cos(ang), 2), np.tile(np.sin(ang), 2),
                          mask, pk, pv]):
        b.write(np.ascontiguousarray(a, np.float32).ravel())
    lm.run_by_index(0, lin, lout)
    cond, eos = lout[0].read(1024, np.float32), lout[1].read(1, np.float32)[0]
    pk[0, :, pos] = lout[2].read(96 * 64, np.float32).reshape(96, 64)
    pv[0, :, pos] = lout[3].read(96 * 64, np.float32).reshape(96, 64)
    pos += 1
    return cond, eos

for t in sp.encode("Hello from LiteRT!"):                    # text prompt
    step(emb[t].astype(np.float32))
lat, x, eos_at = [], bos, None
for g in range(200):                                         # AR loop
    cond, eos = step(x)
    if eos > -4 and eos_at is None: eos_at = g
    if eos_at is not None and g >= eos_at + 3: break         # frames_after_eos
    hin[0].write(cond); hin[1].write(np.random.randn(32).astype(np.float32) * 0.3**0.5)
    hd.run_by_index(0, hin, hout)
    lat.append(hout[0].read(32, np.float32)); x = lat[-1] @ inw.T
# decode `lat` with pt_mimi_dec_tx + pt_mimi_deconly — block layout in the model card,
# full reference: scripts/build_pockettts.py (stage_pipeline) in the repo linked below.

Minimal usage — Kotlin (Android)

val lm = CompiledModel.create(File(dir, "pt_flowlm_step_fp16.tflite").absolutePath,
                              CompiledModel.Options(Accelerator.GPU), null)
val lmIn = lm.createInputBuffers(); val lmOut = lm.createOutputBuffers()
// pk/pv: FloatArray(96*512*64) seeded from pt_voice_alba.bin (fp16), pos = voice length
fun step(embRow: FloatArray): Pair<FloatArray, Float> {
    for (j in 0 until 32) {                       // RoPE at absolute position `pos`
        val a = pos * Math.pow(10000.0, -j / 32.0)
        cos[j] = cos(a).toFloat(); cos[j + 32] = cos[j]
        sin[j] = sin(a).toFloat(); sin[j + 32] = sin[j]
    }
    lmIn[0].writeFloat(embRow); lmIn[1].writeFloat(cos); lmIn[2].writeFloat(sin)
    lmIn[3].writeFloat(mask); lmIn[4].writeFloat(pk); lmIn[5].writeFloat(pv)
    lm.run(lmIn, lmOut)
    val cond = lmOut[0].readFloat(); val eos = lmOut[1].readFloat()[0]
    val nk = lmOut[2].readFloat(); val nv = lmOut[3].readFloat()
    for (g in 0 until 96) {                       // append this step's K/V at `pos`
        System.arraycopy(nk, g * 64, pk, g * 512 * 64 + pos * 64, 64)
        System.arraycopy(nv, g * 64, pv, g * 512 * 64 + pos * 64, 64)
    }
    for (h in 0 until 16) mask[h * 513 + pos] = 0f
    pos++
    return cond to eos
}
// per frame: step() -> flow head (cond + gaussian(32)*sqrt(0.3)) -> latent ->
// input_linear(latent) feeds the next step; stop on eos > -4. Decode latents with
// pt_mimi_dec_tx (64-frame blocks, keep right 512 positions) + pt_mimi_deconly.

The complete, runnable Android app (voice picker, chunking, WAV export) lives in the LiteRT-Models zoo, pockettts/ — including scripts/build_pockettts.py, which rebuilds every graph from the released checkpoint and re-checks all the parity numbers above.

Bundled voices

Only permissively-licensed voices from kyutai/tts-voices and kyutai/pocket-tts-without-voice-cloning are repacked here:

voice	source	license
alba	alba-mackenna	CC BY 4.0
marius, javert	Kyutai voice donations	CC0
charles, mary, eve	VCTK corpus	CC BY 4.0

The Expresso- and EARS-derived voices (e.g. cosette, jean) are CC BY-NC and are not included.

Voice cloning

These graphs cover the released preset-voice path. Voice cloning needs the Mimi encoder, whose weights ship only in the gated kyutai/pocket-tts repo (they are zeroed in the ungated one) — accept Kyutai's terms there if you want to build that path; the encoder recipe is the same SEANet re-authoring used for the decoder.

License & credits

• Model weights: converted from kyutai/pocket-tts (Kyutai), CC BY 4.0 — this repo carries the same license. Code in the linked zoo is MIT.
• Voice prompt states: Kyutai, per-voice licenses above (CC BY 4.0 attribution: VCTK, alba-mackenna).
• Built with litert-torch and LiteRT.