Struct AudioTokenizer 

pub struct AudioTokenizer {
    pub cfg: TokenizerConfig,
    /* private fields */
}

The MiMo audio tokenizer encoder and its residual VQ.

Fields

cfg: TokenizerConfig

Implementations

impl AudioTokenizer

pub fn features( &self, mel: &[f32], m: usize, pool: Option<&Pool>, ) -> Result<Vec<f32>, String>

Pre-RVQ features of ONE segment of m mel frames (mel [m][n_mels]): rows [⌈⌈m/2⌉/2⌉][d] after down_sample_norm.

pub fn quantize( &self, feats: &[f32], rows: usize, bf16_books: bool, pool: Option<&Pool>, ) -> Vec<u32>

Residual VQ in f32: per level, the nearest codeword by ‖E‖² − 2 r·E (the row’s ‖r‖² does not move the argmin), then r −= E[idx]. bf16_books selects the bf16-rounded codebooks a bf16 serving load holds. Returns [rows][levels].

Examples found in repository

examples/mimo_audio_dump.rs (line 253)

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let cmd = args.get(1).map(String::as_str).unwrap_or("");
    let out = PathBuf::from(arg(&args, "--out").expect("--out"));
    match cmd {
        "decode" => {
            let wav = std::fs::read(arg(&args, "--wav").expect("--wav")).unwrap();
            let w = mimo_audio::decode_wav(&wav).unwrap();
            let flat: Vec<f32> = w.channels.concat();
            save_f32(&out, &[w.channels.len(), w.frames()], &flat);
            println!(
                "rate {} channels {} frames {}",
                w.sample_rate,
                w.channels.len(),
                w.frames()
            );
        }
        "frontend" => {
            std::fs::create_dir_all(&out).unwrap();
            let wav = std::fs::read(arg(&args, "--wav").expect("--wav")).unwrap();
            let t0 = Instant::now();
            let w = mimo_audio::decode_wav(&wav).unwrap();
            save_f32(
                &out.join("dec.npy"),
                &[w.channels.len(), w.frames()],
                &w.channels.concat(),
            );
            let chans: Vec<Vec<f32>> = w
                .channels
                .iter()
                .map(|c| mimo_audio::resample_sinc(c, w.sample_rate, mimo_audio::SAMPLE_RATE))
                .collect();
            save_f32(
                &out.join("chan24k.npy"),
                &[chans.len(), chans[0].len()],
                &chans.concat(),
            );
            let mono = mimo_audio::wav_to_mono_24k(&w).unwrap();
            save_f32(&out.join("wave24k.npy"), &[mono.len()], &mono);
            let pool = cortiq_engine::pool::Pool::from_env();
            let (mel, m) = mimo_audio::log_mel(&mono, pool.as_deref()).unwrap();
            save_f32(&out.join("mel.npy"), &[m, mimo_audio::N_MELS], &mel);
            println!(
                "rate {} channels {} frames {} -> {} samples, {m} mel frames, K {} ({:.3}s)",
                w.sample_rate,
                w.channels.len(),
                w.frames(),
                mono.len(),
                mimo_audio::audio_token_count(m, 4),
                t0.elapsed().as_secs_f64()
            );
        }
        "tower" => {
            std::fs::create_dir_all(&out).unwrap();
            let src = PathBuf::from(arg(&args, "--src").expect("--src"));
            let t0 = Instant::now();
            let audio = if src.extension().is_some_and(|e| e == "cmf") {
                let model = Arc::new(cortiq_core::CmfModel::open(&src).expect("open cmf"));
                MimoAudio::from_model(&model).expect("load towers")
            } else {
                MimoAudio::from_hf_dir(&src).expect("load towers")
            };
            let t_load = t0.elapsed().as_secs_f64();
            let (mel, m) = if let Some(mp) = arg(&args, "--mel") {
                let (shape, v) = load_f32(Path::new(&mp));
                assert_eq!(shape[1], mimo_audio::N_MELS);
                (v, shape[0])
            } else {
                let wav = std::fs::read(arg(&args, "--wav").expect("--wav or --mel")).unwrap();
                audio.wav_to_mel(&wav).unwrap()
            };
            let t1 = Instant::now();
            let feats = audio.features(&mel, m).unwrap();
            let t_feats = t1.elapsed().as_secs_f64();
            let d = audio.tokenizer.cfg.d_model;
            let rows = feats.len() / d;
            save_f32(&out.join("feats.npy"), &[rows, d], &feats);
            let t2 = Instant::now();
            let exact = audio.tokenizer.quantize(&feats, rows, false, audio.pool());
            let t_rvq = t2.elapsed().as_secs_f64();
            let rounded = audio.tokenizer.quantize(&feats, rows, true, audio.pool());
            let levels = exact.len() / rows;
            save_i32(&out.join("codes_exact.npy"), &[rows, levels], &exact);
            save_i32(&out.join("codes_bf16books.npy"), &[rows, levels], &rounded);
            let own = mimo_audio::AudioCodes {
                frames: rows,
                levels,
                codes: if audio.bf16_codebooks {
                    rounded.clone()
                } else {
                    exact.clone()
                },
            };
            let t3 = Instant::now();
            let emb_own = audio.embed_codes(&own).unwrap();
            let t_enc = t3.elapsed().as_secs_f64();
            save_f32(
                &out.join("embeds_own.npy"),
                &[emb_own.n_tokens, emb_own.dim],
                &emb_own.rows,
            );
            let fixed = match arg(&args, "--codes") {
                Some(cp) => {
                    let (shape, v) = load_codes(Path::new(&cp));
                    mimo_audio::AudioCodes {
                        frames: shape[0],
                        levels: shape[1],
                        codes: v,
                    }
                }
                None => own.clone(),
            };
            let emb = audio.embed_codes(&fixed).unwrap();
            save_f32(&out.join("embeds.npy"), &[emb.n_tokens, emb.dim], &emb.rows);
            let k = mimo_audio::audio_token_count(m, audio.encoder.cfg.group);
            let meta = serde_json::json!({
                "src": src.display().to_string(),
                "mel_frames": m,
                "segments": mimo_audio::segment_lengths(m),
                "codes": rows,
                "placeholder_count_K": k,
                "embed_rows_own": emb_own.n_tokens,
                "embed_rows_fixed": emb.n_tokens,
                "load_s": t_load,
                "features_s": t_feats,
                "rvq_s": t_rvq,
                "encoder_s": t_enc,
                "bf16_codebooks_default": audio.bf16_codebooks,
                "threads": cortiq_engine::pool::Pool::effective_threads(),
            });
            std::fs::write(
                out.join("tower.json"),
                serde_json::to_string_pretty(&meta).unwrap(),
            )
            .unwrap();
            println!("{meta}");
            assert_eq!(emb_own.n_tokens, k, "placeholder count != encoder rows");
        }
        "calib" => {
            let src = PathBuf::from(arg(&args, "--src").expect("--src"));
            let model = Arc::new(cortiq_core::CmfModel::open(&src).expect("open cmf"));
            let audio = MimoAudio::from_model(&model).expect("load towers");
            let dir = PathBuf::from(arg(&args, "--wav-dir").expect("--wav-dir"));
            let mut wavs: Vec<PathBuf> = std::fs::read_dir(&dir)
                .unwrap()
                .filter_map(|e| e.ok().map(|e| e.path()))
                .filter(|p| p.extension().is_some_and(|e| e == "wav"))
                .collect();
            wavs.sort();
            let t0 = Instant::now();
            cortiq_engine::gptq_capture::begin(true);
            let mut frames = 0usize;
            for w in &wavs {
                let emb = audio.embed_wav(&std::fs::read(w).unwrap()).unwrap();
                frames += emb.n_tokens;
                eprintln!(
                    "  {} -> {} rows ({:.0}s)",
                    w.display(),
                    emb.n_tokens,
                    t0.elapsed().as_secs_f64()
                );
            }
            let hess = cortiq_engine::gptq_capture::end();
            save_hessians(&out, &hess);
            println!(
                "{} clips, {frames} LLM rows, {} linears -> {} ({:.0}s)",
                wavs.len(),
                hess.len(),
                out.display(),
                t0.elapsed().as_secs_f64()
            );
        }
        _ => {
            eprintln!(
                "usage: mimo_audio_dump (decode|frontend|tower|calib) --out ... (see the source header)"
            );
            std::process::exit(2);
        }
    }
}