Pocket-TTS-LiteRT

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):

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.

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 graph table above
# (64-frame blocks with 32-frame overlap, then one 256-frame SEANet window).

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 two snippets above are the reference for wiring the graphs. The graph build script (build_pockettts.py: litert-torch export of every graph, the fp16 pass, and a tflite-vs-eager parity check per stage) and the Android sample app live in LiteRT-Models/pockettts.

Bundled voices

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

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 LiteRT-Models is MIT.
• Voice prompt states: Kyutai, per-voice licenses above (CC BY 4.0 attribution: VCTK, alba-mackenna).
• Built with litert-torch and LiteRT.