--- 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 | | | 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) ```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 { 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).