a2f-mlx 0.2.0

//! a2f-mlx — Audio2Face-3D v2.3.1 (James) as a hand-written MLX graph (Rust).
//!
//! Mirrors the validated Python port (~/dev/a2f-mlx/a2f_mlx.py), which itself
//! mirrors soniqo/speech-swift's Audio2Face3DMLXRuntime.swift. Weights come
//! from `tools/export_mlx_weights.py` — extracted from NVIDIA's official
//! network.onnx (keep-NVIDIA provenance; no third-party repack).
//!
//! Graph:
//!   audio branch : wav2vec2-style conv frontend + 4-layer transformer
//!   freq branch  : 25 hann windows → 32-lag autocorrelation → 5× conv2d
//!   join         : 4× conv2d stride(2,1) + emotion broadcast → dual MLP heads
//!   out          : 169 = 140 skin | 10 tongue | 15 jaw | 4 eye

use std::cell::RefCell;
use std::collections::HashMap;
use std::path::Path;

use anyhow::{bail, Context, Result};
use mlx_rs::{
    error::Exception,
    ops::{self, indexing::IndexOp},
    Array, Dtype,
};
use safetensors::SafeTensors;

use a2f_core::{
    hann640, Emotion, EXPLICIT_EMO, IMPLICIT_EMO, N_EMO, OUT_COUNT, SKIN, TONGUE,
};

pub const HEADS: usize = 12;
const SQRT2: f32 = std::f32::consts::SQRT_2;
const EPS: f32 = 1e-5;

/// Tensor-name families ("variants") present in the official ONNX. For this
/// bundle the live one is `"2230"` — decided by ONNX parity, not by name.
pub const VARIANT_2230: Variant = Variant {
    emo_proj: "2230",
    fe_scale: "2235",
    fe_bias: "2236",
    feat_proj: "2237",
    pos_conv: "2240",
    audio_map: "2289",
    layers: [
        [2241, 2244, 2247, 2250, 2251, 2252],
        [2253, 2256, 2259, 2262, 2263, 2264],
        [2265, 2268, 2271, 2274, 2275, 2276],
        [2277, 2280, 2283, 2286, 2287, 2288],
    ],
};

#[derive(Clone, Copy)]
pub struct Variant {
    pub emo_proj: &'static str,
    pub fe_scale: &'static str,
    pub fe_bias: &'static str,
    pub feat_proj: &'static str,
    pub pos_conv: &'static str,
    pub audio_map: &'static str,
    /// [q, k, v, o, ff1, ff2] per transformer layer (MatMul initializer names).
    pub layers: [[u32; 6]; 4],
}

/// Weight storage precision. `Q8`/`Q4` quantize the transformer + projection
/// linears (affine, group 64) and keep everything else f32.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum Precision {
    #[default]
    F32,
    F16,
    BF16,
    Q8,
    Q4,
}

impl std::str::FromStr for Precision {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self> {
        Ok(match s.to_ascii_lowercase().as_str() {
            "f32" => Self::F32,
            "f16" => Self::F16,
            "bf16" => Self::BF16,
            "q8" | "int8" => Self::Q8,
            "q4" | "int4" => Self::Q4,
            _ => bail!("precision `{s}`: want f32|f16|bf16|q8|q4"),
        })
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Options {
    pub precision: Precision,
}

const Q_GROUP: i32 = 64;


/// A stored weight: dense, or an `mlx::quantize`d [out,in] matrix.
#[derive(Clone)]
enum Weight {
    Dense(Array),
    Quant { w: Array, scales: Array, biases: Array, bits: i32 },
}

fn p(layer: usize, rest: &str) -> String {
    format!("wave2vec_model.encoder.transformer.layers.{layer}.{rest}")
}

/// Load safetensors into MLX arrays with layouts baked for the graph:
/// conv1d torch [O,I,K] → [O,K,I]; conv2d torch [O,I,kH,kW] → [O,kH,kW,I];
/// torch-style matmul [out,in] → [in,out]. Named as-is otherwise.
fn load_weights(bytes: &[u8], v: &Variant) -> Result<HashMap<String, Array>> {
    let st = SafeTensors::deserialize(bytes)?;
    let mut m = HashMap::new();
    let get = |m: &mut HashMap<String, Array>, name: &str, data: &[u8], shape: &[usize]| {
        let f: Vec<f32> = data
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();
        let sh: Vec<i32> = shape.iter().map(|&d| d as i32).collect();
        m.insert(name.to_string(), Array::from_slice(&f, &sh));
    };
    for (name, view) in st.iter() {
        match view.dtype() {
            safetensors::Dtype::F32 => {
                get(&mut m, name, view.data(), view.shape())
            }
            other => bail!("tensor `{name}` has dtype {other:?}, need F32 — re-export"),
        }
    }
    // bake layouts (transpose once at load, not per call)
    let conv1d = [
        "wave2vec_model.feature_extractor.conv_layers.0.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.1.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.2.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.3.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.4.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.5.conv.weight",
        "wave2vec_model.feature_extractor.conv_layers.6.conv.weight",
        v.pos_conv,
    ];
    for name in conv1d {
        let a = m.remove(name).unwrap();
        m.insert(name.into(), a.transpose_axes(&[0, 2, 1])?);
    }
    let conv2d = ["freq1.weight", "freq2.weight", "freq3.weight", "freq4.weight", "freq5.weight",
        "time1.weight", "time2.weight", "time3.weight", "time4.weight"];
    for name in conv2d {
        let a = m.remove(name).unwrap();
        m.insert(name.into(), a.transpose_axes(&[0, 2, 3, 1])?);
    }
    let outin = ["time5.weight", "out_mapping.weight", "out_mapping1.weight",
        "out_mapping_tongue.weight", "out_mapping1_tongue.weight"];
    for name in outin {
        let a = m.remove(name).unwrap();
        m.insert(name.into(), a.transpose()?);
    }
    // transformer MatMul initializers are [in, out] already — keep.
    Ok(m)
}

/// Linears eligible for `Q8`/`Q4`: the 24 transformer matmuls plus the
/// feature projection and audio feature map (~29M of 39M params). All have an
/// input dim divisible by the quantization group.
fn quantizable(v: &Variant) -> Vec<String> {
    let mut names: Vec<String> = v.layers.iter().flatten().map(|id| id.to_string()).collect();
    names.push(v.feat_proj.into());
    names.push(v.audio_map.into());
    names
}

fn prepare(raw: HashMap<String, Array>, v: &Variant, precision: Precision) -> Result<HashMap<String, Weight>> {
    let quant: Vec<String> = quantizable(v);
    let mut out = HashMap::with_capacity(raw.len());
    for (name, a) in raw {
        let w = match precision {
            Precision::F32 => Weight::Dense(a),
            Precision::F16 => Weight::Dense(a.as_dtype(Dtype::Float16)?),
            Precision::BF16 => Weight::Dense(a.as_dtype(Dtype::Bfloat16)?),
            Precision::Q8 | Precision::Q4 if quant.contains(&name) => {
                let bits = if precision == Precision::Q8 { 8 } else { 4 };
                // quantize along `in`: store [out, in], matmul with transpose=true
                let (w, scales, biases) = ops::quantize(a.transpose()?, Q_GROUP, bits)?;
                Weight::Quant { w, scales, biases, bits }
            }
            Precision::Q8 | Precision::Q4 => Weight::Dense(a),
        };
        // Materialize per tensor (cast/quantize/transpose → contiguous) so the
        // first window doesn't pay for it and the f32 source is freed as we go.
        let w = match w {
            Weight::Dense(a) => {
                let a = a.contiguous()?;
                a.eval()?;
                Weight::Dense(a)
            }
            Weight::Quant { w, scales, biases, bits } => {
                mlx_rs::transforms::eval([&w, &scales, &biases])?;
                Weight::Quant { w, scales, biases, bits }
            }
        };
        out.insert(name, w);
    }
    Ok(out)
}

/// The weights + graph. Cloning is cheap (MLX arrays are refcounted), which
/// lets a compiled closure own its own copy.
#[derive(Clone)]
struct Net {
    w: HashMap<String, Weight>,
    v: Variant,
    dtype: Dtype,
    hann: Array,
    combine: Vec<f32>, // normalized combine weights
}

type CompiledFn = Box<dyn FnMut(&[Array]) -> std::result::Result<Vec<Array>, Exception>>;

pub struct A2FMlx {
    net: Net,
    opts: Options,
    input_strength: f32,
    default_emotion: Emotion,
}

/// [`A2FMlx`] with its forward pass traced by `mlx::compile` (fused
/// elementwise kernels, one trace per batch size). mlx-rs's compiled closure
/// holds `Rc`s, so this wrapper is `!Send` — build it on the thread that runs
/// inference. Plain [`A2FMlx`] stays `Send`.
pub struct A2FMlxCompiled {
    model: A2FMlx,
    f: RefCell<CompiledFn>,
}

impl A2FMlxCompiled {
    pub fn model(&self) -> &A2FMlx {
        &self.model
    }

    pub fn coefficients(&self, window: &[f32], emotion: &Emotion) -> Result<Vec<f32>> {
        Ok(self.coefficients_batch(&[window], emotion)?.remove(0).to_vec())
    }

    pub fn coefficients_batch(&self, windows: &[&[f32]], emotion: &Emotion) -> Result<Vec<[f32; OUT_COUNT]>> {
        self.model.run_batch(windows, emotion, |audio, emo| {
            Ok((self.f.borrow_mut())(&[audio, emo])?.remove(0))
        })
    }
}

impl A2FMlx {
    pub fn load(bundle_dir: &Path, input_strength: Option<f32>) -> Result<Self> {
        Self::load_with(bundle_dir, input_strength, Options::default())
    }

    pub fn load_with(bundle_dir: &Path, input_strength: Option<f32>, opts: Options) -> Result<Self> {
        Self::load_with_variant(bundle_dir, VARIANT_2230, input_strength, opts)
    }

    pub fn load_with_variant(
        bundle_dir: &Path,
        variant: Variant,
        input_strength: Option<f32>,
        opts: Options,
    ) -> Result<Self> {
        let bytes = std::fs::read(bundle_dir.join("audio2face3d.safetensors"))
            .with_context(|| format!("read {}", bundle_dir.join("audio2face3d.safetensors").display()))?;
        let raw = load_weights(&bytes, &variant)?;

        let mc: serde_json::Value = serde_json::from_str(
            &std::fs::read_to_string(bundle_dir.join("model_config.json"))
                .unwrap_or_else(|_| "{}".into()),
        )?;
        let input_strength = input_strength.unwrap_or(
            mc.get("config")
                .and_then(|c| c.get("input_strength"))
                .and_then(|v| v.as_f64())
                .unwrap_or(1.3) as f32,
        );
        let default_emotion = Emotion::from_raw(
            &std::fs::read(bundle_dir.join("default_emotion.f32"))
                .context("default_emotion.f32 missing from bundle")?,
        )?;

        // combine_weights / Σ (python: cw = cw / cw.sum())
        let cw_arr = raw.get("combine_weights").context("combine_weights")?;
        let cw: Vec<f32> = cw_arr.as_slice::<f32>().to_vec();
        let sum: f32 = cw.iter().sum();
        let combine = cw.into_iter().map(|c| c / sum).collect();

        let dtype = match opts.precision {
            Precision::F16 => Dtype::Float16,
            Precision::BF16 => Dtype::Bfloat16,
            _ => Dtype::Float32,
        };
        let w = prepare(raw, &variant, opts.precision)?;
        let net = Net {
            w,
            v: variant,
            dtype,
            hann: Array::from_slice(&hann640(), &[1, 1, 1, 640]).as_dtype(dtype)?,
            combine,
        };

        Ok(Self { net, opts, input_strength, default_emotion })
    }

    /// Wrap in a `mlx::compile`d forward pass (see [`A2FMlxCompiled`]).
    pub fn compiled(self) -> A2FMlxCompiled {
        let n = self.net.clone();
        let f = move |args: &[Array]| -> Vec<Array> {
            vec![n.forward(&args[0], &args[1]).expect("a2f forward")]
        };
        let f: CompiledFn = Box::new(mlx_rs::transforms::compile::compile(f, false));
        A2FMlxCompiled { model: self, f: RefCell::new(f) }
    }

    pub fn default_emotion(&self) -> &Emotion {
        &self.default_emotion
    }

    pub fn options(&self) -> Options {
        self.opts
    }

    /// window: BUFFER_LEN f32 samples; emotion: N_EMO f32. Returns 169 coefs.
    pub fn coefficients(&self, window: &[f32], emotion: &Emotion) -> Result<Vec<f32>> {
        Ok(self.coefficients_batch(&[window], emotion)?.remove(0).to_vec())
    }

    /// Several windows (e.g. consecutive 30 fps frames) in one forward pass.
    /// Each window is BUFFER_LEN samples; returns one 169-vector per window.
    pub fn coefficients_batch(&self, windows: &[&[f32]], emotion: &Emotion) -> Result<Vec<[f32; OUT_COUNT]>> {
        self.run_batch(windows, emotion, |audio, emo| self.net.forward(&audio, &emo))
    }

    fn run_batch(
        &self,
        windows: &[&[f32]],
        emotion: &Emotion,
        forward: impl FnOnce(Array, Array) -> Result<Array>,
    ) -> Result<Vec<[f32; OUT_COUNT]>> {
        use a2f_core::BUFFER_LEN;
        let b = windows.len();
        anyhow::ensure!(b > 0, "empty batch");
        let mut scaled = Vec::with_capacity(b * BUFFER_LEN);
        for w in windows {
            assert_eq!(w.len(), BUFFER_LEN);
            scaled.extend(w.iter().map(|s| s * self.input_strength));
        }
        let mut emo = Vec::with_capacity(b * N_EMO);
        for _ in 0..b {
            emo.extend_from_slice(&emotion.0);
        }
        let audio = Array::from_slice(&scaled, &[b as i32, 1, BUFFER_LEN as i32]);
        let emo = Array::from_slice(&emo, &[b as i32, 1, N_EMO as i32]);
        let out = forward(audio, emo)?;
        let flat = out.as_slice::<f32>();
        Ok(flat
            .chunks_exact(OUT_COUNT)
            .map(|c| c.try_into().expect("OUT_COUNT"))
            .collect())
    }

    /// audio [B,1,T]; emotion [B,1,26] → [B,1,169] f32
    pub fn forward(&self, audio: &Array, emotion: &Array) -> Result<Array> {
        self.net.forward(audio, emotion)
    }

    pub fn input_strength(&self) -> f32 {
        self.input_strength
    }

    pub fn emotion_kind(&self) -> (usize, usize) {
        (IMPLICIT_EMO, EXPLICIT_EMO)
    }

    pub fn out_count(&self) -> usize {
        OUT_COUNT
    }

    /// Bench/profiling hook: median ms of each graph stage on `audio` [B,1,T].
    #[doc(hidden)]
    pub fn stage_times(&self, audio: &Array, emotion: &Array, reps: usize) -> Result<Vec<(&'static str, f64)>> {
        let audio = audio.as_dtype(self.net.dtype)?;
        let time = |f: &dyn Fn() -> Result<Array>| -> Result<f64> {
            f()?.eval()?;
            let mut v: Vec<f64> = (0..reps)
                .map(|_| {
                    let t = std::time::Instant::now();
                    f().and_then(|a| Ok(a.eval()?)).map(|_| t.elapsed().as_secs_f64() * 1e3)
                })
                .collect::<Result<_>>()?;
            v.sort_by(|a, b| a.partial_cmp(b).unwrap());
            Ok(v[0]) // min: contention only adds time
        };
        Ok(vec![
            ("freq_branch", time(&|| self.net.freq_branch(&audio))?),
            ("conv_frontend", time(&|| self.net.conv_frontend(&audio))?),
            ("+feat_proj", time(&|| self.net.feat_proj(&self.net.conv_frontend(&audio)?))?),
            ("+pos_embed", time(&|| self.net.pos_embed(&self.net.feat_proj(&self.net.conv_frontend(&audio)?)?))?),
            ("audio_branch", time(&|| self.net.audio_branch(&audio))?),
            ("forward", time(&|| self.net.forward(&audio, emotion))?),
        ])
    }

    /// Parity helper: run one fixture window/emotion pair, return coefficients.
    pub fn run_fixture(&self, window: &[f32], emotion: &[f32]) -> Result<Vec<f32>> {
        self.coefficients(window, &Emotion(emotion.to_vec()))
    }
}

impl Net {
    fn t(&self, name: &str) -> Result<Array> {
        match self.w.get(name) {
            Some(Weight::Dense(a)) => Ok(a.clone()),
            Some(Weight::Quant { .. }) => bail!("tensor `{name}` is quantized; use linear()"),
            None => bail!("missing tensor `{name}`"),
        }
    }

    fn bias1(&self, name: &str) -> Result<Array> {
        self.t(name)
    }

    /// x @ W for a stored [in,out] weight (dense) or [out,in] quantized one.
    fn linear(&self, x: &Array, name: &str) -> Result<Array> {
        match self.w.get(name) {
            Some(Weight::Dense(a)) => Ok(x.matmul(a)?),
            Some(Weight::Quant { w, scales, biases, bits }) => {
                Ok(ops::quantized_matmul(x, w, scales, biases, true, Q_GROUP, *bits)?)
            }
            None => bail!("missing tensor `{name}`"),
        }
    }

    fn w3(&self, x: &Array, id: u32) -> Result<Array> {
        self.linear(x, &id.to_string())
    }

    fn scalar(&self, v: f32) -> Array {
        Array::from_f32(v).as_dtype(self.dtype).expect("scalar cast")
    }

    /// x: NCL [B,C,T]; weight pre-baked [O,K,I/g]; returns NCL [B,O,T].
    fn conv1d_ncl(&self, x: &Array, weight: &Array, bias: Option<&Array>,
                  stride: i32, padding: i32, groups: i32) -> Result<Array> {
        let y = ops::conv1d(x.transpose_axes(&[0, 2, 1])?, weight, stride, padding, 1, groups)?;
        let y = y.transpose_axes(&[0, 2, 1])?; // back to NCL
        match bias {
            Some(b) => {
                let b = b.reshape(&[1, -1, 1])?;
                Ok(y + b)
            }
            None => Ok(y),
        }
    }

    /// x: NCHW; weight pre-baked [O,kH,kW,I]; returns NCHW.
    fn conv2d_nchw(&self, x: &Array, weight: &Array, bias: &Array,
                   stride: (i32, i32), padding: (i32, i32)) -> Result<Array> {
        let y = ops::conv2d(x.transpose_axes(&[0, 2, 3, 1])?, weight, stride, padding, (1, 1), 1)?;
        let y = y.transpose_axes(&[0, 3, 1, 2])? + bias.reshape(&[1, -1, 1, 1])?;
        Ok(y)
    }

    fn gelu(&self, x: &Array) -> Result<Array> {
        let t = x / self.scalar(SQRT2);
        let t = ops::erf(&t)?;
        Ok(x * (t + self.scalar(1.0)) * self.scalar(0.5))
    }

    fn layer_norm(&self, x: &Array, wname: &str, bname: &str) -> Result<Array> {
        Ok(mlx_rs::fast::layer_norm(x, &self.t(wname)?, &self.t(bname)?, EPS)?)
    }

    /// wav2vec2 conv feature extractor. Input [B,1,T]; returns NCL [B,512,T'].
    fn conv_frontend(&self, x: &Array) -> Result<Array> {
        let fe = |i: u32| {
            format!("wave2vec_model.feature_extractor.conv_layers.{i}.conv.weight")
        };
        let mut x = self.conv1d_ncl(x, &self.t(&fe(0))?, None, 5, 0, 1)?;
        // instance norm with learned scale/bias, then gelu
        let mean = x.mean_axis(-1, Some(true))?;
        let centered = &x - &mean;
        let var = (&centered * &centered).mean_axis(-1, Some(true))?;
        x = &centered / (&var + self.scalar(EPS)).sqrt()?
            * self.t(self.v.fe_scale)?.reshape(&[1, -1, 1])?
            + self.t(self.v.fe_bias)?.reshape(&[1, -1, 1])?;
        x = self.gelu(&x)?;
        for i in 1..7 {
            x = self.conv1d_ncl(&x, &self.t(&fe(i))?, None, 2, 0, 1)?;
            x = self.gelu(&x)?;
        }
        Ok(x)
    }

    /// wav2vec2-style audio branch. Input [B,1,T]; returns [B,T',C] combined.
    fn audio_branch(&self, x: &Array) -> Result<Array> {
        let x = self.conv_frontend(x)?;
        let x = self.feat_proj(&x)?;
        let x = self.pos_embed(&x)?;
        self.transformer(&x)
    }

    /// NCL conv features → [B,T,768] projected.
    fn feat_proj(&self, x: &Array) -> Result<Array> {
        // feature projection: NLC + LN + linear
        let mut x = x.transpose_axes(&[0, 2, 1])?; // [B,T,C]
        x = self.layer_norm(
            &x,
            "wave2vec_model.encoder.feature_projection.layer_norm.weight",
            "wave2vec_model.encoder.feature_projection.layer_norm.bias",
        )?;
        x = self.linear(&x, self.v.feat_proj)?
            + self.bias1("wave2vec_model.encoder.feature_projection.projection.bias")?;
        Ok(x)
    }

    /// positional conv embedding + encoder layer norm.
    fn pos_embed(&self, x: &Array) -> Result<Array> {
        let mut x = x.clone();
        // positional conv (groups=16, pad 64, truncate) + gelu, residual
        let t_len = x.shape()[1];
        let pos = self.conv1d_ncl(
            &x.transpose_axes(&[0, 2, 1])?,
            &self.t(self.v.pos_conv)?,
            Some(&self.bias1("wave2vec_model.encoder.transformer.pos_conv_embed.conv.bias")?),
            1, 64, 16,
        )?;
        let pos = pos.index((.., .., ..(t_len as i32)));
        x = x + self.gelu(&pos)?.transpose_axes(&[0, 2, 1])?;

        x = self.layer_norm(
            &x,
            "wave2vec_model.encoder.transformer.layer_norm.weight",
            "wave2vec_model.encoder.transformer.layer_norm.bias",
        )?;
        Ok(x)
    }

    /// 4 post-LN transformer layers + learned weighted combine of their outputs.
    fn transformer(&self, x: &Array) -> Result<Array> {
        let mut x = x.clone();
        let shape = x.shape();
        let (b, t_len, hidden) = (shape[0], shape[1], shape[2]);
        let head_dim = hidden / HEADS as i32;
        let mut outputs = Vec::with_capacity(4);
        for (li, ids) in self.v.layers.iter().enumerate() {
            let pre = p(li, "attention");
            let q = (self.w3(&x, ids[0])? + self.bias1(&format!("{pre}.q_proj.bias"))?)
                .reshape(&[b, t_len, HEADS as i32, head_dim])?
                .transpose_axes(&[0, 2, 1, 3])?;
            let k = (self.w3(&x, ids[1])? + self.bias1(&format!("{pre}.k_proj.bias"))?)
                .reshape(&[b, t_len, HEADS as i32, head_dim])?;
            let v = (self.w3(&x, ids[2])? + self.bias1(&format!("{pre}.v_proj.bias"))?)
                .reshape(&[b, t_len, HEADS as i32, head_dim])?
                .transpose_axes(&[0, 2, 1, 3])?;
            let k = k.transpose_axes(&[0, 2, 1, 3])?;
            let ctx = mlx_rs::fast::scaled_dot_product_attention(
                &q, &k, &v, 1.0 / (head_dim as f32).sqrt(), None, None)?;
            let ctx = ctx
                .transpose_axes(&[0, 2, 1, 3])?
                .reshape(&[b, t_len, hidden])?;
            x = x + self.w3(&ctx, ids[3])?
                + self.bias1(&format!("{pre}.out_proj.bias"))?;
            x = self.layer_norm(&x, &p(li, "layer_norm.weight"), &p(li, "layer_norm.bias"))?;
            let ff = self.w3(&x, ids[4])?
                + self.bias1(&p(li, "feed_forward.intermediate_dense.bias"))?;
            let ff = self.gelu(&ff)?;
            x = x + self.w3(&ff, ids[5])?
                + self.bias1(&p(li, "feed_forward.output_dense.bias"))?;
            x = self.layer_norm(
                &x,
                &p(li, "final_layer_norm.weight"),
                &p(li, "final_layer_norm.bias"),
            )?;
            outputs.push(x.clone());
        }
        // weighted combine (weights normalized at load)
        let mut acc = &outputs[0] * self.scalar(self.combine[0]);
        for (o, c) in outputs.iter().zip(self.combine.iter()).skip(1) {
            acc = acc + o * self.scalar(*c);
        }
        Ok(acc)
    }

    /// 25 hann-windowed 640-sample frames (hop 320) → 32-lag autocorrelation,
    /// as two strided views + one batched matmul instead of ~100 small ops.
    fn autocorr(&self, scaled: &Array) -> Result<Array> {
        let b = scaled.shape()[0];
        let t = scaled.shape()[2] as i64;
        let frames = ops::as_strided(scaled, &[b, 1, 25, 640][..], &[t, t, 320, 1][..], 0)?;
        let mean = frames.mean_axis(-1, Some(true))?;
        let centered = (&frames - mean) * &self.hann; // [B,1,25,640]
        let padded = ops::pad(&centered, &[(0, 0), (0, 0), (0, 0), (0, 32)][..], None, None)?;
        let lagged = ops::as_strided(&padded, &[b, 1, 25, 32, 640][..],
                                     &[25 * 672, 25 * 672, 672, 1, 1][..], 0)?;
        let corr = lagged.matmul(&centered.reshape(&[b, 1, 25, 640, 1])?)?; // [B,1,25,32,1]
        Ok(corr.reshape(&[b, 1, 25, 32])? * self.scalar(1.0 / 640.0))
    }

    /// Autocorrelation frequency branch. Input [B,1,T] (raw audio); [B,c,h,w] NCHW.
    fn freq_branch(&self, x: &Array) -> Result<Array> {
        let scaled = x * self.t("in_scale_autocorr")?;
        let mut out = self.autocorr(&scaled)?;
        let spec = [
            ("freq1", (1, 2), (0, 1)),
            ("freq2", (1, 2), (0, 1)),
            ("freq3", (1, 2), (0, 1)),
            ("freq4", (1, 2), (0, 1)),
            ("freq5", (1, 1), (0, 0)),
        ];
        for (name, stride, padding) in spec {
            let y = self.conv2d_nchw(
                &out,
                &self.t(&format!("{name}.weight"))?,
                &self.t(&format!("{name}.bias"))?,
                stride,
                padding,
            )?;
            out = ops::maximum(&y, self.scalar(0.0))? * self.scalar(SQRT2);
        }
        Ok(out)
    }

    /// audio [B,1,T]; emotion [B,1,26] (f32) → [B,1,169] f32
    fn forward(&self, audio: &Array, emotion: &Array) -> Result<Array> {
        let audio = audio.as_dtype(self.dtype)?;
        let emotion = emotion.as_dtype(self.dtype)?;
        let implicit = emotion.index((.., .., ..(IMPLICIT_EMO as i32)));
        let explicit = emotion.index((.., .., (IMPLICIT_EMO as i32)..));
        let projected = explicit.matmul(&self.t(self.v.emo_proj)?)?
            + self.bias1("emo_linear1.bias")?;

        let freq = self.freq_branch(&audio)?;
        let aud = self.audio_branch(&audio)?;

        let head = ops::maximum(
            &(self.linear(&aud, self.v.audio_map)?
                + self.bias1("audio_feature_map.bias")?),
            self.scalar(0.0),
        )?;
        let hs = head.shape();
        // [B,1,C,T] → [B,1,C,T,1]
        let aud5 = head
            .reshape(&[hs[0], 1, hs[1], hs[2], 1])?
            .transpose_axes(&[0, 1, 3, 2, 4])?; // [B,1,T,C,1]
        let fs = freq.shape();
        let freq5 = freq.reshape(&[fs[0], 1, fs[1], fs[2], fs[3]])?;

        let joined = ops::concatenate(&[aud5, freq5], 2)?;
        let js = joined.shape();
        let mut x = joined.reshape(&[js[0] * js[1], js[2], js[3], js[4]])?;

        for name in ["time1", "time2", "time3", "time4"] {
            let y = self.conv2d_nchw(
                &x,
                &self.t(&format!("{name}.weight"))?,
                &self.t(&format!("{name}.bias"))?,
                (2, 1),
                (1, 0),
            )?;
            x = ops::maximum(&y, self.scalar(0.0))? * self.scalar(SQRT2);
            let xs = x.shape();
            let (b, _, h, w) = (xs[0], xs[1], xs[2], xs[3]);
            let imp = ops::broadcast_to(
                &implicit.reshape(&[b, IMPLICIT_EMO as i32, 1, 1])?,
                &[b, IMPLICIT_EMO as i32, h, w],
            )?;
            let prj = ops::broadcast_to(
                &projected.reshape(&[b, 8, 1, 1])?,
                &[b, 8, h, w],
            )?;
            x = ops::concatenate(&[x, imp, prj], 1)?;
        }

        let xs = x.shape();
        let flat = x.reshape(&[xs[0], xs[1] * xs[2] * xs[3]])?;
        let time5 = (flat.matmul(&self.t("time5.weight")?)?
            + self.bias1("time5.bias")?)
            .reshape(&[xs[0], 1, 256])?;
        let time5 = ops::maximum(&time5, self.scalar(0.0))? * self.scalar(SQRT2);
        let head_in = ops::concatenate(&[time5, implicit, projected], 2)?;
        let his = head_in.shape();
        let head_flat = head_in.reshape(&[his[0] * his[1], his[2]])?;

        let face = (head_flat.matmul(&self.t("out_mapping.weight")?)?
            + self.bias1("out_mapping.bias")?)
            .reshape(&[his[0] * his[1], a2f_core::FACE_COUNT as i32])?;
        let face = (face.matmul(&self.t("out_mapping1.weight")?)?
            + self.bias1("out_mapping1.bias")?)
            .reshape(&[his[0], his[1], a2f_core::FACE_COUNT as i32])?;
        let tongue = (head_flat.matmul(&self.t("out_mapping_tongue.weight")?)?
            + self.bias1("out_mapping_tongue.bias")?)
            .reshape(&[his[0] * his[1], TONGUE as i32])?;
        let tongue = (tongue.matmul(&self.t("out_mapping1_tongue.weight")?)?
            + self.bias1("out_mapping1_tongue.bias")?)
            .reshape(&[his[0], his[1], TONGUE as i32])?;

        let skin = face.index((.., .., ..(SKIN as i32)));
        let jaw_eye = face.index((.., .., (SKIN as i32)..));
        Ok(ops::concatenate(&[skin, tongue, jaw_eye], 2)?.as_dtype(Dtype::Float32)?)
    }
}

#[cfg(test)]
mod tests {
    fn assert_send<T: Send>() {}

    #[test]
    fn model_is_send() {
        assert_send::<super::A2FMlx>();
    }
}