README.md

---
license: cc-by-4.0
base_model: kyutai/pocket-tts
language:
- en
pipeline_tag: text-to-speech
tags:
- litert
- tflite
- android
- on-device
- tts
- kyutai
---

# Pocket-TTS-LiteRT

[Pocket TTS](https://huggingface.co/kyutai/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).

![hero](assets/hero.png)

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

| voice | sample |
|---|---|
| alba | <audio controls src="https://huggingface.co/mlboydaisuke/Pocket-TTS-LiteRT/resolve/main/samples/alba_device.wav"></audio> |
| marius | <audio controls src="https://huggingface.co/mlboydaisuke/Pocket-TTS-LiteRT/resolve/main/samples/marius_device.wav"></audio> |

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

```python
# 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)

```kotlin
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/`](https://github.com/john-rocky/LiteRT-Models) — 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](https://huggingface.co/kyutai/tts-voices)
and [kyutai/pocket-tts-without-voice-cloning](https://huggingface.co/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](https://huggingface.co/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](https://huggingface.co/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](https://datashare.ed.ac.uk/handle/10283/3443), alba-mackenna).
* Built with [litert-torch](https://github.com/google-ai-edge/litert-torch) and
  [LiteRT](https://github.com/google-ai-edge/LiteRT).