README.md

---
license: cc-by-4.0
base_model: kyutai/mimi
tags:
- litert
- on-device
- neural-codec
- audio
library_name: litert
---

# Mimi (Kyutai 2024) — LiteRT on-device

On-device [LiteRT](https://ai.google.dev/edge/litert) conversion of [**kyutai/mimi**](https://huggingface.co/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** delegate (`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.

![Mimi — original vs 32-codebook RVQ round-trip on-device (LiteRT)](samples/sample.png)

## Files

| File | Size | Delegate | In → Out |
|------|------|----------|----------|
| `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
```

## Minimal usage

**Android (Kotlin, LiteRT CompiledModel)**

```kotlin
fun load(name: String, acc: Accelerator) =                       // models staged in filesDir
    CompiledModel.create(File(filesDir, name).absolutePath, CompiledModel.Options(acc), null)

val encConv = load("mimi_enc_conv_fp16.tflite", Accelerator.GPU) // audio[1,1,48000] -> feat[1,512,50]
val encTx   = load("mimi_enc_tx_fp16.tflite",  Accelerator.CPU)  // feat.T[1,50,512] -> emb[1,512,25]
val decTx   = load("mimi_dec_tx_fp16.tflite",  Accelerator.CPU)  // emb[1,512,25] -> convIn[1,512,50]
val deconly = load("mimi_deconly_fp16.tflite", Accelerator.GPU)  // convIn -> audio[1,1,48000]

val inB = encConv.createInputBuffers(); val outB = encConv.createOutputBuffers()
inB[0].writeFloat(audio)     // 48000 floats = 2 s @ 24 kHz, [-1,1]
encConv.run(inB, outB)
val feat = outB[0].readFloat()  // transpose (1,512,50)->(1,50,512), feed encTx, then host RVQ.
// Full chain incl. the split-RVQ host code: compiled_model_api/audio_codec in litert-samples.
```

**Python (desktop verification, full round-trip incl. RVQ codes)**

```python
import numpy as np, soundfile as sf
from ai_edge_litert.interpreter import Interpreter

def run(path, x):
    it = Interpreter(model_path=path); it.allocate_tensors()
    it.set_tensor(it.get_input_details()[0]["index"], x.astype(np.float32)); it.invoke()
    return it.get_tensor(it.get_output_details()[0]["index"])

wav, _ = sf.read("in.wav", dtype="float32")               # 24 kHz mono
x = np.zeros((1, 1, 48000), np.float32); n = min(len(wav), 48000); x[0, 0, :n] = wav[:n]

# 1) encode: SEANet convs (GPU graph) -> transformer + downsample -> emb [1,512,25]
feat = run("mimi_enc_conv_fp16.tflite", x)                # [1,512,50]
emb = run("mimi_enc_tx_fp16.tflite", feat.transpose(0, 2, 1))[0]   # [512,25]

# 2) split RVQ (host). mimi_rvq.bin = sem_Win[256,512], aco_Win[256,512], sem_Wout[512,256],
#    aco_Wout[512,256], sem_CB[2048,256], 31x aco_CB[2048,256] — float32 LE, contiguous.
D, S, H = 256, 2048, 512
w, o = np.fromfile("mimi_rvq.bin", "<f4"), 0
def take(*sh):
    global o; n = int(np.prod(sh)); a = w[o:o + n].reshape(sh); o += n; return a
sem_Win, aco_Win, sem_Wout, aco_Wout = take(D, H), take(D, H), take(H, D), take(H, D)
sem_CB, aco_CB = take(S, D), [take(S, D) for _ in range(31)]

def nearest(r_T, CB):                                     # Euclidean argmin over the codebook
    return ((CB * CB).sum(1)[None] - 2.0 * (r_T @ CB.T)).argmin(1)

codes = np.zeros((32, emb.shape[1]), np.int64)            # 1 semantic + 31 acoustic @ 12.5 Hz
codes[0] = nearest((sem_Win @ emb).T, sem_CB)
res = aco_Win @ emb                                       # acoustic residual loop
for i in range(31):
    codes[1 + i] = nearest(res.T, aco_CB[i]); res -= aco_CB[i][codes[1 + i]].T

# 3) decode: codes -> emb -> transformer + upsample -> SEANet deconv (GPU graph)
q = sem_Wout @ sem_CB[codes[0]].T + aco_Wout @ sum(aco_CB[i][codes[1 + i]] for i in range(31)).T
conv_in = run("mimi_dec_tx_fp16.tflite", q[None])         # [1,512,50]
sf.write("roundtrip.wav", run("mimi_deconly_fp16.tflite", conv_in)[0, 0], 24000)
```

## 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 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 artifact. 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 (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.

## Sample app

A complete Android sample app + the conversion/RVQ-export scripts are in the official LiteRT samples
repository under **`compiled_model_api/audio_codec`** (google-ai-edge/litert-samples). Push these files
to the app's `filesDir` with that sample's `install_to_device.sh`.

License follows upstream Mimi (CC-BY-4.0).