continuous-mimi-onnx

Continuous Mimi Encoder

This is a modified ONNX export of the Mimi audio codec encoder.

It bypasses the Residual Vector Quantization (RVQ) to directly output continuous feature representations instead of discrete tokens. This preserves fine-grained acoustic details, making it ideal for downstream tasks like real-time dialogue analysis or speech emotion recognition.

Figure: Architecture showing the RVQ bypass for continuous latents. (Adapted from the Moshi technical report)

License and Attribution

This model is a derivative work based on the "Mimi" model. * Original Developer: Kyutai * License: CC BY 4.0 * Modifications: Bypassed the RVQ module to output continuous embeddings and exported to ONNX format.

Speed and accuracy (reference)

Numbers below are illustrative (desktop CPU, FP32 ONNX vs PyTorch kyutai/mimi, streaming chunk emit=1280 @ 16 kHz, inference.py --compare). Your machine, thread counts, and ORT / PyTorch builds will shift them.

Speed — RTF (processing time ÷ real-time audio length per chunk)
In repeated 1-second windows, Torch often lands around RTF ≈ 0.22–0.32 and FP32 ONNX around RTF ≈ 0.14–0.18 — i.e. ONNX is commonly ~1.5–2× less wall time than Torch for the same window. Both RTF values stay below 1 in those traces, so either path can run faster than real time for this configuration.

Accuracy — Torch vs FP32 ONNX embeddings
With identical windows fed to both backends over a long stream (400 steps = 32 s, well past the 10 s point where the transformer's KV ring buffer wraps), the worst per-step max_abs is ≤ 1.3e−6 for both the 1-frame and 2-frame exports, with cosine similarity ≈ 1.0 throughout. The published FP32 ONNX pair is the intended match to Torch; INT8 trades more numerical error for CPU throughput.

Changelog

2026-10 — fix for drift after 10 s of streaming. Earlier exports diverged from Torch once the encoder transformer's KV cache (250 steps = 10 s) wrapped around: the attention mask assumed the ring buffer was in chronological order, so after the wrap a query could attend to a key one step in its future, and the buffer was one slot too small for a 2-step call. Differences grew from ~1e−7 to ~1e−3 … 5e−2 (cosine down to ~0.97) and stayed there. The mask is now built from the absolute position stored in each slot, and the cache holds sliding_window + steps_per_call − 1 slots (max_past_len 251 for the 1-frame files, 253 for _2f). The _2f files are also re-exported from the original kyutai/mimi weights (the previous _2f upload used fine-tuned weights by mistake). No code change is needed: runtimes read max_past_len from the .json. The previous files remain available in this repository's history.

Getting Started

Requirements

Use Python 3.10+ and install the following (versions should match your CUDA/CPU setup for PyTorch).

Recommended: transformers==5.5.3 (streaming Mimi + cache behaviour used by the ONNX export and inference.py are validated against this release).

Example install for CPU ONNX + Torch:

pip install "torch" "transformers==5.5.3" einops numpy huggingface_hub onnxruntime

For GPU ONNX, replace onnxruntime with onnxruntime-gpu and install a matching torch build. For live microphone streaming, add pyaudio when your platform supports it.

The Torch encoder path downloads kyutai/mimi weights via Hugging Face on first use. The ONNX path, if you do not pass local paths, downloads continuous_mimi_fp32.onnx / .json (or INT8 counterparts) from this Hugging Face repository using huggingface_hub.

inference.py

inference.py provides a small CLI around the streaming Mimi encoder: PyTorch (kyutai/mimi) or ONNX, optional 16 kHz microphone streaming, and an optional Torch–ONNX comparison mode. Run it with python inference.py ... after installing the packages above.

Quick checks

# One synthetic streaming step; PyTorch Mimi encoder (HF weights)
python inference.py --backend torch

# Same, but ONNX runtime (FP32 bundle from this Hub, cached locally after first run)
python inference.py --backend onnx --onnx-precision fp32

# Use ONNX files you already downloaded next to the script
python inference.py --backend onnx --onnx-model ./continuous_mimi_fp32.onnx --onnx-meta ./continuous_mimi_fp32.json

# INT8 ONNX is intended for CPU in this script (do not pair INT8 with CUDA here)
python inference.py --backend onnx --onnx-precision int8 --device cpu

Microphone (16 kHz, mono float32)

Requires pyaudio. Prints RTF (processing time ÷ real-time audio duration) and per-chunk timing on a fixed wall-clock interval (default 1 s).

python inference.py --mic --backend onnx
python inference.py --mic --backend onnx --mic-device-index 1 --max-steps 500 --mic-report-interval 1.0

Torch vs ONNX comparison

Initializes both backends and feeds identical streaming windows; reports numeric differences and per-backend timing / RTF.

# Synthetic multi-step stream (default 10 steps; override with --compare-steps)
python inference.py --compare --compare-steps 20

# Same comparison using live microphone chunks
python inference.py --compare --mic

Example (--compare on CPU, 400 synthetic steps = 32 s)

Each line reports embedding agreement (max_abs, mean_abs, cos) for one streaming step, plus inference time and RTF for Torch and ONNX separately. Step 125 is where the transformer's KV cache first wraps around.

$ python inference.py --compare --compare-steps 400
Froze EncoderMimi!
Froze EncoderMimiOnnx!
Compare mode: 400 streaming steps, emit=1280, overlap=320, device=cpu
step    0  max_abs=1.713634e-07  mean_abs=4.131152e-08  cos=1.00000000  |  torch=82.48ms RTF=1.031  onnx=13.33ms RTF=0.167
step    1  max_abs=2.235174e-07  mean_abs=5.715674e-08  cos=1.00000012  |  torch=16.68ms RTF=0.208  onnx=12.29ms RTF=0.154
...
step  124  max_abs=2.682209e-07  mean_abs=5.849643e-08  cos=1.00000000  |  torch=15.41ms RTF=0.193  onnx=10.16ms RTF=0.127
step  125  max_abs=2.533197e-07  mean_abs=5.629795e-08  cos=1.00000000  |  torch=14.61ms RTF=0.183  onnx=8.92ms RTF=0.112
step  126  max_abs=3.278255e-07  mean_abs=6.289383e-08  cos=1.00000012  |  torch=14.26ms RTF=0.178  onnx=9.13ms RTF=0.114
...
step  399  max_abs=2.533197e-07  mean_abs=6.198024e-08  cos=0.99999994  |  torch=15.76ms RTF=0.197  onnx=10.84ms RTF=0.136
Done. worst max_abs_diff=4.470348e-07 at step 141  |  avg torch=15.21ms RTF=0.190  avg onnx=9.75ms RTF=0.122

--compare --mic prints the same statistics aggregated per wall-clock second ([1s] lines) from live microphone audio.

When --compare is set, --backend is ignored (both paths run). With --compare and without --mic, --compare-steps sets how many synthetic chunks to run.

Other useful flags