rlx-mimi 0.2.11

Kyutai Mimi neural audio codec (12.5 Hz, 24 kHz) for RLX
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// Decode/encode throughput for Mimi across every compiled rlx backend.
//
//   RLX_MIMI_DIR=/path/to/mimi cargo run -p rlx-mimi --example bench --release \
//     --features cpu,metal,mlx,gpu -- --dur 4 --iters 5
//
// Weights: a kyutai/mimi dir with `model.safetensors` + `config.json`. Real codes
// come from encoding synthetic PCM.

use anyhow::Result;
use rlx_core::AudioCodec;
use rlx_core::codec_bench as cb;
use rlx_mimi::MimiCodec;
use rlx_runtime::Device;

const CRATE: &str = "rlx-mimi";

// `candidates` is seeded empty then extended via cfg-gated push()es (below),
// which can't be folded into the vec![] literal across feature combos.
#[allow(clippy::vec_init_then_push)]
fn main() -> Result<()> {
    let (dur, iters) = cb::parse_dur_iters();
    let dir = std::env::var_os("RLX_MIMI_DIR")
        .map(std::path::PathBuf::from)
        .unwrap_or_else(|| std::path::PathBuf::from(".cache/mimi"));
    if !dir.join("model.safetensors").is_file() {
        eprintln!(
            "skip {CRATE}: missing {}/model.safetensors (set RLX_MIMI_DIR)",
            dir.display()
        );
        return Ok(());
    }

    #[allow(unused_mut)] // push()es below are cfg-gated on backend features
    let mut candidates = vec![];
    #[cfg(feature = "metal")]
    candidates.push(Device::Metal);
    #[cfg(feature = "mlx")]
    candidates.push(Device::Mlx);
    #[cfg(feature = "gpu")]
    candidates.push(Device::Gpu);

    for dev in cb::available(&candidates) {
        if let Err(e) = bench_device(&dir, dev, dur, iters) {
            cb::report_fail(CRATE, dev, "all", &e.to_string());
        }
    }
    Ok(())
}

fn bench_device(dir: &std::path::Path, dev: Device, dur: f64, iters: usize) -> Result<()> {
    let codec = MimiCodec::open_on(dir, dev)?;
    let info = codec.info();
    let n_samples = (dur * info.sample_rate as f64) as usize;
    let pcm = cb::synth_pcm(n_samples, 0.3, 0x717010);

    let codes = match cb::time_once_ms(|| codec.encode_pcm(&pcm, None)) {
        Ok((codes, compile_ms)) => {
            let t = codes.num_frames();
            let audio_s = (t * info.hop_length) as f64 / info.sample_rate as f64;
            let run_ms = cb::time_median_ms(0, iters, || codec.encode_pcm(&pcm, None))?;
            cb::report(CRATE, dev, "encode", audio_s, t, compile_ms, run_ms);
            codes
        }
        Err(e) => {
            cb::report_fail(CRATE, dev, "encode", &e.to_string());
            return Ok(());
        }
    };

    let t = codes.num_frames();
    let audio_s = (t * info.hop_length) as f64 / info.sample_rate as f64;
    match cb::time_once_ms(|| codec.decode_codes(&codes)) {
        Ok((_, compile_ms)) => {
            let run_ms = cb::time_median_ms(0, iters, || codec.decode_codes(&codes))?;
            cb::report(CRATE, dev, "decode", audio_s, t, compile_ms, run_ms);
        }
        Err(e) => cb::report_fail(CRATE, dev, "decode", &e.to_string()),
    }
    Ok(())
}