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