mlboydaisuke/Mimi-LiteRT

Mimi (Kyutai 2024) — LiteRT on-device (hybrid GPU/CPU)

On-device LiteRT conversion of kyutai/mimi, the Kyutai/Moshi streaming neural audio codec (24 kHz, 12.5 Hz frame rate). The heavy SEANet convolutional halves run on the CompiledModel GPU (ML Drift / LITERT_CL); the two 8-layer Transformers and the split RVQ run on CPU. Device-verified on a Pixel 8a (Tensor G3): full round-trip at RTF ≈ 0.35 (faster than real-time), reconstruction at the codec's quality floor.

Files

File	Size	Placement	Input → Output
mimi_enc_conv_fp16.tflite	24 MB	GPU	audio [1,1,L] → feat [1,512,Se]
mimi_enc_tx_fp16.tflite	50 MB	CPU	feat [1,Se,512] → emb [1,512,Tc]
mimi_dec_tx_fp16.tflite	48 MB	CPU	emb [1,512,Tc] → conv_in [1,512,seq]
mimi_deconly_fp16.tflite	28 MB	GPU	conv_in [1,512,seq] → audio [1,1,L]
mimi_rvq.bin	69 MB	CPU	codes ↔ emb (32 codebooks, float32 LE)

Graphs are fixed-length (built per duration). The example set is for a 2 s clip (Se=50, Tc=25, seq=50).

Pipeline

audio →[GPU enc_conv]→ feat →[CPU enc_tx]→ emb →[CPU RVQ.encode]→ codes
      →[CPU RVQ.decode]→ emb →[CPU dec_tx]→ conv_in →[GPU deconly]→ audio

Why hybrid

Every op in all four graphs is GPU-clean (re-authored), and the convs are fp16-exact on Mali (decoder-only fed the exact transformer output = 48 dB SNR). But the decoder transformer's residual stream reaches |x|=27, where the Mali GPU delegate's internal fp16 compute loses precision — full-GPU decode drops to ~12 dB on real speech. The transformer behaves identically standalone and fused on device, so this is fp16 precision, not a fusion collapse. The transformers are tiny (8 layers × 512, seq ~50), so CPU is trivial and exact; the heavy SEANet convs stay on GPU. The split RVQ (1 semantic + 31 acoustic, Euclidean argmin + int64 indices) runs on CPU.

Re-authoring (litert-torch, parity ~1.0)

tanh-GELU · baked RoPE cos/sin + rotate_half · baked causal additive bias · MimiLayerScale→Linear · grouped ZeroStuffConvT1d (depthwise upsample, no TRANSPOSE_CONV) · baked constant conv pad · nn.ELU→relu(x)−relu(1−exp(min(x,0))) · replicate-pad→SLICE+CONCAT.

Usage

Sample app + conversion scripts: LiteRT-Models / mimi. Push the files to the app's filesDir with mimi/scripts/install_to_device.sh.

This conversion uses model-op rewrites + a GPU/CPU placement split, so it is hosted in a personal namespace (not a patch-free "clean" sample). License follows upstream Mimi (CC-BY-4.0).

---

Want a different model on-device? Open a request — free, open weights only; the export and its measured numbers get published publicly.