a2f-mlx 0.2.0

//! Round-robin A/B of model variants within one process (robust to drifting
//! GPU/CPU contention from other apps). Each spec is `precision:compile:batch`,
//! e.g. `f32:0:1 f16:1:8`.
//!   ab BUNDLE WAV REPS spec1 spec2 ...
use std::time::Instant;
use a2f_core::{read_wav_16k_mono, BUFFER_LEN};
use a2f_core::{Emotion, OUT_COUNT};
use a2f_mlx::{A2FMlx, A2FMlxCompiled, Options};

enum Model {
    Plain(A2FMlx),
    Compiled(A2FMlxCompiled),
}

impl Model {
    fn run(&self, w: &[&[f32]], e: &Emotion) -> anyhow::Result<Vec<[f32; OUT_COUNT]>> {
        match self { Model::Plain(m) => m.coefficients_batch(w, e), Model::Compiled(c) => c.coefficients_batch(w, e) }
    }
}

fn main() -> anyhow::Result<()> {
    let a: Vec<String> = std::env::args().collect();
    let audio = read_wav_16k_mono(a[2].as_ref())?;
    let reps: usize = a[3].parse()?;
    let mut variants = Vec::new();
    for spec in &a[4..] {
        let f: Vec<&str> = spec.split(':').collect();
        let opts = Options { precision: f[0].parse()? };
        let plain = A2FMlx::load_with(a[1].as_ref(), None, opts)?;
        let emo = plain.default_emotion().clone();
        let model = if f[1] == "1" { Model::Compiled(plain.compiled()) } else { Model::Plain(plain) };
        let batch: usize = f[2].parse()?;
        variants.push((spec.clone(), emo, batch, model));
    }
    let wins: Vec<Vec<f32>> = (0..8).map(|k| audio[20000 + k * 533..20000 + k * 533 + BUFFER_LEN].to_vec()).collect();
    let refs: Vec<&[f32]> = wins.iter().map(|w| w.as_slice()).collect();
    let mut t = vec![Vec::new(); variants.len()];
    for r in 0..reps + 3 {
        for (i, (_, emo, batch, model)) in variants.iter().enumerate() {
            let s = Instant::now();
            model.run(&refs[..*batch], emo)?;
            if r >= 3 { t[i].push(s.elapsed().as_secs_f64() * 1e3); }
        }
    }
    for ((spec, _, batch, _), v) in variants.iter().zip(t.iter_mut()) {
        v.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("{spec:<16} call: min {:6.2}  p25 {:6.2}  median {:6.2} ms | per frame: min {:5.2}  median {:5.2} ms",
            v[0], v[v.len() / 4], v[v.len() / 2], v[0] / *batch as f64, v[v.len() / 2] / *batch as f64);
    }
    Ok(())
}