Mimi (Kyutai 2024) — LiteRT on-device

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 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.

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)

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)

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

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
TFLite benchmark_model (TfLiteGpuDelegateV2) — mimi_dec_tx_fp16.tflite	GPU (OpenCL)	689 / 689	21.7 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) — mimi_deconly_fp16.tflite	GPU (OpenCL)	220 / 220	114.3 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) — mimi_enc_conv_fp16.tflite	GPU (OpenCL)	189 / 189	93.4 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) — mimi_enc_tx_fp16.tflite	GPU (OpenCL)	681 / 681	28.0 ms
TFLite benchmark_model — mimi_dec_tx_fp16.tflite	CPU (XNNPACK, 4 threads)	—	XNNPACK declined the graph
TFLite benchmark_model — mimi_deconly_fp16.tflite	CPU (XNNPACK, 4 threads)	—	936.2 ms
TFLite benchmark_model — mimi_enc_conv_fp16.tflite	CPU (XNNPACK, 4 threads)	—	XNNPACK declined the graph
TFLite benchmark_model — mimi_enc_tx_fp16.tflite	CPU (XNNPACK, 4 threads)	—	53.1 ms

Any on-device figure recorded when this model shipped came from a different runtime. It was taken through LiteRT's own CompiledModel accelerator (logcat reports it as LITERT_CL), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

XNNPACK declines these fp16 graphs — it reports failed to delegate DEPTHWISE_CONV_2D and then fails to allocate tensors — so there is no usable CPU number. Disabling XNNPACK falls back to reference kernels, which measured about 20× slower than the GPU on models of this size and would not represent CPU inference anyone would ship.