inferencelayer 0.2.3

//! VITS TTS (`facebook/mms-tts-eng` — MMS English) — text → waveform, end-to-end. Pure Rust, CPU.
//!
//! The engine's second synthesis model and its first normalizing-flow architecture. Char-level
//! input (no phonemizer — MMS English is trained on raw lowercased characters, blank-interspersed).
//! Pipeline, each stage gated against a `transformers` oracle (`tests/fixtures/export_vits.py`):
//!   1. tokenizer: keep-if-in-vocab (else lowercase), strip unknown, intersperse blank id 0
//!   2. text encoder: post-LN transformer with windowed RELATIVE attention (learned per-layer
//!      `emb_rel_k`/`emb_rel_v`, window ±4) and conv-based FFN (kernel 3, asymmetric pad);
//!      a 1×1-conv projection then yields the prior means/log-variances per input token
//!   3. stochastic duration predictor, run in REVERSE: zero latents (noise scale 0) through an
//!      inverse ElementwiseAffine + 3 inverse ConvFlows (rational-quadratic spline, DDS-conv
//!      conditioning) → log-durations → ceil(exp(·)) frames per token
//!   4. the prior means expand by repetition, then the residual-coupling flow runs in reverse
//!      (4 layers, each: half → 1×1 conv → gated WaveNet → mean; subtract; flip channels)
//!   5. HiFi-GAN decoder (VITS flavour: 192-ch input, biasless conv_post, no mel normalization)
//!
//! Determinism: HF VITS samples noise for the duration flow and the prior; the oracle export sets
//! both `noise_scale`s to 0 and this port hardcodes the same — synthesis is deterministic.

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

use anyhow::{Context, Result};
use rayon::prelude::*;

use crate::weights::LazySt;

/// `y[m,n] = x[m,k]·Wᵀ + b`, W `[n,k]` row-major (HF Linear). Parallel over rows.
fn linear(x: &[f32], w: &[f32], b: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
    let mut out = vec![0f32; m * n];
    out.par_chunks_mut(n).enumerate().for_each(|(i, orow)| {
        let xr = &x[i * k..][..k];
        for o in 0..n {
            let wr = &w[o * k..][..k];
            let mut acc = b[o];
            for c in 0..k {
                acc += xr[c] * wr[c];
            }
            orow[o] = acc;
        }
    });
    out
}

/// Row-wise LayerNorm over `d`, `[m, d]` (biased variance, eps 1e-5).
fn layer_norm(x: &[f32], d: usize, w: &[f32], b: &[f32]) -> Vec<f32> {
    let mut out = vec![0f32; x.len()];
    out.par_chunks_mut(d)
        .zip(x.par_chunks(d))
        .for_each(|(orow, row)| {
            let mean = row.iter().sum::<f32>() / d as f32;
            let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / d as f32;
            let inv = 1.0 / (var + 1e-5).sqrt();
            for i in 0..d {
                orow[i] = (row[i] - mean) * inv * w[i] + b[i];
            }
        });
    out
}

struct Linear {
    w: Vec<f32>,
    b: Vec<f32>,
    n: usize,
    k: usize,
}
impl Linear {
    fn load(st: &LazySt, prefix: &str, n: usize, k: usize) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let b = st.tensor_f32(&format!("{prefix}.bias"))?;
        anyhow::ensure!(w.len() == n * k, "{prefix}.weight {} != {n}x{k}", w.len());
        Ok(Self { w, b, n, k })
    }
    fn forward(&self, x: &[f32], m: usize) -> Vec<f32> {
        linear(x, &self.w, &self.b, m, self.n, self.k)
    }
}

struct Norm {
    w: Vec<f32>,
    b: Vec<f32>,
}
impl Norm {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        Ok(Self {
            w: st.tensor_f32(&format!("{prefix}.weight"))?,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
        })
    }
    fn forward(&self, x: &[f32], d: usize) -> Vec<f32> {
        layer_norm(x, d, &self.w, &self.b)
    }
}

/// Conv1d over channel-major `[in, t]` → `[out, t_out]`; optional bias, dilation, groups ∈ {1, in}.
struct Conv1d {
    w: Vec<f32>, // [out, in/groups, k]
    b: Vec<f32>, // empty = no bias
    out_ch: usize,
    in_ch: usize,
    k: usize,
    dilation: usize,
    pad_l: usize,
    depthwise: bool,
}
impl Conv1d {
    fn load(
        st: &LazySt,
        prefix: &str,
        out_ch: usize,
        in_ch: usize,
        k: usize,
        dilation: usize,
        depthwise: bool,
    ) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let per_in = if depthwise { 1 } else { in_ch };
        anyhow::ensure!(
            w.len() == out_ch * per_in * k,
            "{prefix}.weight {} != {out_ch}x{per_in}x{k}",
            w.len()
        );
        let b = st.tensor_f32(&format!("{prefix}.bias")).unwrap_or_default();
        let pad = (k * dilation - dilation) / 2;
        Ok(Self {
            w,
            b,
            out_ch,
            in_ch,
            k,
            dilation,
            pad_l: pad,
            depthwise,
        })
    }
    fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let mut y = vec![0f32; self.out_ch * t];
        y.par_chunks_mut(t).enumerate().for_each(|(o, orow)| {
            let bias = self.b.get(o).copied().unwrap_or(0.0);
            let (ic_range, w_base) = if self.depthwise {
                (o..o + 1, o * self.k) // groups == channels: out o reads in o only
            } else {
                (0..self.in_ch, o * self.in_ch * self.k)
            };
            for (ti, oval) in orow.iter_mut().enumerate() {
                let mut acc = bias;
                for (wi, ic) in ic_range.clone().enumerate() {
                    let xrow = &x[ic * t..][..t];
                    let wrow = &self.w[w_base + wi * self.k..][..self.k];
                    for (kk, wv) in wrow.iter().enumerate() {
                        let src = ti + kk * self.dilation;
                        if src >= self.pad_l && src - self.pad_l < t {
                            acc += xrow[src - self.pad_l] * wv;
                        }
                    }
                }
                *oval = acc;
            }
        });
        y
    }
}

/// LayerNorm across CHANNELS at each time step (input channel-major `[ch, t]`), eps 1e-5.
fn ln_channels(x: &[f32], ch: usize, t: usize, w: &[f32], b: &[f32]) -> Vec<f32> {
    let mut out = vec![0f32; ch * t];
    for ti in 0..t {
        let mut mean = 0f32;
        for c in 0..ch {
            mean += x[c * t + ti];
        }
        mean /= ch as f32;
        let mut var = 0f32;
        for c in 0..ch {
            let d = x[c * t + ti] - mean;
            var += d * d;
        }
        var /= ch as f32;
        let inv = 1.0 / (var + 1e-5).sqrt();
        for c in 0..ch {
            out[c * t + ti] = (x[c * t + ti] - mean) * inv * w[c] + b[c];
        }
    }
    out
}

/// Exact (erf) GELU — torch's default `F.gelu`.
fn gelu(v: f32) -> f32 {
    0.5 * v * (1.0 + libm::erff(v * std::f32::consts::FRAC_1_SQRT_2))
}

fn softplus(x: f32) -> f32 {
    // torch F.softplus, beta=1, threshold=20
    if x > 20.0 { x } else { (1.0 + x.exp()).ln() }
}

/// Dilated depth-separable conv stack (`VitsDilatedDepthSeparableConv`): per layer, depthwise
/// dilated conv → LN(channels) → gelu → pointwise conv → LN(channels) → gelu → residual.
struct DdsConv {
    dilated: Vec<Conv1d>,
    pointwise: Vec<Conv1d>,
    norms1: Vec<Norm>,
    norms2: Vec<Norm>,
}
impl DdsConv {
    fn load(st: &LazySt, prefix: &str, ch: usize, k: usize, layers: usize) -> Result<Self> {
        let mut dilated = Vec::new();
        let mut pointwise = Vec::new();
        let mut norms1 = Vec::new();
        let mut norms2 = Vec::new();
        for i in 0..layers {
            dilated.push(Conv1d::load(
                st,
                &format!("{prefix}.convs_dilated.{i}"),
                ch,
                ch,
                k,
                k.pow(i as u32),
                true,
            )?);
            pointwise.push(Conv1d::load(
                st,
                &format!("{prefix}.convs_pointwise.{i}"),
                ch,
                ch,
                1,
                1,
                false,
            )?);
            norms1.push(Norm::load(st, &format!("{prefix}.norms_1.{i}"))?);
            norms2.push(Norm::load(st, &format!("{prefix}.norms_2.{i}"))?);
        }
        Ok(Self {
            dilated,
            pointwise,
            norms1,
            norms2,
        })
    }
    /// `x` `[ch, t]`; `cond` (same shape) is added at entry, as HF does.
    fn forward(&self, x: &[f32], ch: usize, t: usize, cond: Option<&[f32]>) -> Vec<f32> {
        let mut x = x.to_vec();
        if let Some(c) = cond {
            for (xi, ci) in x.iter_mut().zip(c) {
                *xi += ci;
            }
        }
        for i in 0..self.dilated.len() {
            let h = self.dilated[i].forward(&x, t);
            let mut h = ln_channels(&h, ch, t, &self.norms1[i].w, &self.norms1[i].b);
            for v in h.iter_mut() {
                *v = gelu(*v);
            }
            let h = self.pointwise[i].forward(&h, t);
            let mut h = ln_channels(&h, ch, t, &self.norms2[i].w, &self.norms2[i].b);
            for v in h.iter_mut() {
                *v = gelu(*v);
            }
            for (xi, hi) in x.iter_mut().zip(h) {
                *xi += hi;
            }
        }
        x
    }
}

/// Inverse of the monotonic piecewise rational-quadratic spline (`tail_bound` identity outside).
/// `wh`/`hh` are the 10 unnormalized widths/heights (already ÷√filter), `dh` the 9 inner
/// unnormalized derivatives; boundary derivatives are pinned so that softplus(·)+1e-3 = 1.
fn spline_inverse(v: f32, wh: &[f32], hh: &[f32], dh: &[f32], tail: f32) -> f32 {
    let bins = wh.len();
    if !(-tail..=tail).contains(&v) {
        return v;
    }
    let (min_w, min_h, min_d) = (1e-3f32, 1e-3f32, 1e-3f32);
    let softmax = |u: &[f32]| {
        let m = u.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let e: Vec<f32> = u.iter().map(|x| (x - m).exp()).collect();
        let s: f32 = e.iter().sum();
        e.iter().map(|x| x / s).collect::<Vec<f32>>()
    };
    // cumulative bin edges in [-tail, tail]
    let edges = |uh: &[f32], min_b: f32| {
        let sm = softmax(uh);
        let mut cum = vec![0f32; bins + 1];
        let mut acc = 0f32;
        for i in 0..bins {
            acc += min_b + (1.0 - min_b * bins as f32) * sm[i];
            cum[i + 1] = 2.0 * tail * acc - tail;
        }
        cum[0] = -tail;
        cum[bins] = tail;
        cum
    };
    let cw = edges(wh, min_w);
    let mut ch = edges(hh, min_h);
    // derivatives: d[0] = d[bins] = softplus(constant)+min_d = 1 exactly (HF pins the raw value)
    let constant = ((1.0f32 - min_d).exp() - 1.0).ln();
    let mut d = vec![0f32; bins + 1];
    d[0] = min_d + softplus(constant);
    d[bins] = d[0];
    for i in 0..bins - 1 {
        d[i + 1] = min_d + softplus(dh[i]);
    }
    // reverse mode: bins located by cumheights (last edge nudged +1e-6)
    ch[bins] += 1e-6;
    let mut bin = 0usize;
    for (i, e) in ch.iter().enumerate() {
        if v >= *e {
            bin = i;
        }
    }
    let bin = bin.min(bins - 1);
    let (cw_b, w_b) = (cw[bin], cw[bin + 1] - cw[bin]);
    let (ch_b, h_b) = (ch[bin], ch[bin + 1] - ch[bin]);
    let delta = h_b / w_b;
    let (d_b, d_b1) = (d[bin], d[bin + 1]);
    let i1 = d_b + d_b1 - 2.0 * delta;
    let i2 = v - ch_b;
    let i3 = i2 * i1;
    let a = h_b * (delta - d_b) + i3;
    let b = h_b * d_b - i3;
    let c = -delta * i2;
    let disc = b * b - 4.0 * a * c;
    let root = (2.0 * c) / (-b - disc.sqrt());
    root * w_b + cw_b
}

/// One `VitsConvFlow` (spline flow over 2 channels, DDS-conv conditioning), reverse only.
struct ConvFlow {
    conv_pre: Conv1d, // 1 → filter
    dds: DdsConv,
    conv_proj: Conv1d, // filter → 3·bins−1
    bins: usize,
    tail: f32,
    filter: usize,
}
impl ConvFlow {
    /// `first`/`second` are the two 1-channel halves `[t]`; `cond` is `[filter, t]`.
    fn reverse(&self, first: &[f32], second: &mut [f32], cond: &[f32], t: usize) {
        let h = self.conv_pre.forward(first, t);
        let h = self.dds.forward(&h, self.filter, t, Some(cond));
        let h = self.conv_proj.forward(&h, t);
        let scale = 1.0 / (self.filter as f32).sqrt();
        let nb = self.bins;
        for ti in 0..t {
            let wh: Vec<f32> = (0..nb).map(|r| h[r * t + ti] * scale).collect();
            let hh: Vec<f32> = (0..nb).map(|r| h[(nb + r) * t + ti] * scale).collect();
            let dh: Vec<f32> = (0..nb - 1).map(|r| h[(2 * nb + r) * t + ti]).collect();
            second[ti] = spline_inverse(second[ti], &wh, &hh, &dh, self.tail);
        }
    }
}

/// The stochastic duration predictor, REVERSE (inference) path only, noise scale 0.
struct DurationPredictor {
    conv_pre: Conv1d,
    dds: DdsConv,
    conv_proj: Conv1d,
    ea_translate: Vec<f32>, // [2]
    ea_log_scale: Vec<f32>, // [2]
    flows: Vec<ConvFlow>,
}
impl DurationPredictor {
    /// encoder hidden, channel-major `[hidden, t]` → log-durations `[t]`.
    fn log_durations(&self, hidden_cm: &[f32], ch: usize, t: usize) -> Vec<f32> {
        let g = self.conv_pre.forward(hidden_cm, t);
        let g = self.dds.forward(&g, ch, t, None);
        let g = self.conv_proj.forward(&g, t);
        // latents start at zero (noise_scale_duration = 0). Flow order: reversed ConvFlows minus
        // the "useless vflow" (HF drops index -2), then the ElementwiseAffine; channel-flip
        // BEFORE each flow. With 2 channels the flip is a swap.
        let mut lat0 = vec![0f32; t];
        let mut lat1 = vec![0f32; t];
        for flow in self.flows.iter().rev() {
            std::mem::swap(&mut lat0, &mut lat1);
            flow.reverse(&lat0, &mut lat1, &g, t);
        }
        std::mem::swap(&mut lat0, &mut lat1);
        for c in 0..2 {
            let l = if c == 0 { &mut lat0 } else { &mut lat1 };
            let (tr, ls) = (self.ea_translate[c], self.ea_log_scale[c]);
            for v in l.iter_mut() {
                *v = (*v - tr) * (-ls).exp();
            }
        }
        lat0
    }
}

/// One WaveNet-conditioned residual coupling layer (`mean_only`), reverse path.
struct CouplingLayer {
    conv_pre: Conv1d, // half → hidden
    in_layers: Vec<Conv1d>,
    res_skip: Vec<Conv1d>,
    conv_post: Conv1d, // hidden → half
}
impl CouplingLayer {
    /// `x` `[flow, t]` split as first/second halves; second ← second − m(first).
    fn reverse(&self, x: &mut [f32], half: usize, hidden: usize, t: usize) {
        let first = &x[..half * t];
        let mut inputs = self.conv_pre.forward(first, t);
        let mut out = vec![0f32; hidden * t];
        let n = self.in_layers.len();
        for (l, in_layer) in self.in_layers.iter().enumerate() {
            let y = in_layer.forward(&inputs, t);
            // fused_add_tanh_sigmoid_multiply with zero conditioning
            let mut acts = vec![0f32; hidden * t];
            for i in 0..hidden * t {
                acts[i] = y[i].tanh() * (1.0 / (1.0 + (-y[hidden * t + i]).exp()));
            }
            let rs = self.res_skip[l].forward(&acts, t);
            if l < n - 1 {
                for i in 0..hidden * t {
                    inputs[i] += rs[i];
                    out[i] += rs[hidden * t + i];
                }
            } else {
                for i in 0..hidden * t {
                    out[i] += rs[i];
                }
            }
        }
        let m = self.conv_post.forward(&out, t);
        for i in 0..half * t {
            x[half * t + i] -= m[i];
        }
    }
}

/// VITS-flavoured HiFi-GAN decoder: latents `[flow, t]` → waveform. No input normalization,
/// biasless `conv_post`.
struct VitsDecoder {
    conv_pre: Conv1d,
    upsampler: Vec<ConvT1d>,
    resblocks: Vec<ResBlock>,
    conv_post: Conv1d,
    num_kernels: usize,
    slope: f32,
}
impl VitsDecoder {
    fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let mut x = self.conv_pre.forward(x, t);
        let mut t = t;
        for (i, up) in self.upsampler.iter().enumerate() {
            for e in x.iter_mut() {
                if *e < 0.0 {
                    *e *= self.slope;
                }
            }
            let (xu, tu) = up.forward(&x, t);
            t = tu;
            let ch = up.out_ch;
            let mut acc = self.resblocks[i * self.num_kernels].forward(&xu, ch, t, self.slope);
            for j in 1..self.num_kernels {
                let r = self.resblocks[i * self.num_kernels + j].forward(&xu, ch, t, self.slope);
                for (a, b) in acc.iter_mut().zip(r) {
                    *a += b;
                }
            }
            let inv = 1.0 / self.num_kernels as f32;
            for a in acc.iter_mut() {
                *a *= inv;
            }
            x = acc;
        }
        // HF calls leaky_relu WITHOUT a slope here — torch default 0.01, not the config's 0.1
        for e in x.iter_mut() {
            if *e < 0.0 {
                *e *= 0.01;
            }
        }
        let y = self.conv_post.forward(&x, t);
        y.iter().map(|v| v.tanh()).collect()
    }
}

/// ConvTranspose1d, weights `[in, out, k]` (torch layout).
struct ConvT1d {
    w: Vec<f32>,
    b: Vec<f32>,
    out_ch: usize,
    in_ch: usize,
    k: usize,
    stride: usize,
    pad: usize,
}
impl ConvT1d {
    fn load(
        st: &LazySt,
        prefix: &str,
        in_ch: usize,
        out_ch: usize,
        k: usize,
        stride: usize,
    ) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        anyhow::ensure!(
            w.len() == in_ch * out_ch * k,
            "{prefix}.weight {} != {in_ch}x{out_ch}x{k}",
            w.len()
        );
        Ok(Self {
            w,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
            out_ch,
            in_ch,
            k,
            stride,
            pad: (k - stride) / 2,
        })
    }
    fn forward(&self, x: &[f32], t: usize) -> (Vec<f32>, usize) {
        let t_out = (t - 1) * self.stride + self.k - 2 * self.pad;
        let mut y = vec![0f32; self.out_ch * t_out];
        y.par_chunks_mut(t_out).enumerate().for_each(|(o, orow)| {
            for (ti, oval) in orow.iter_mut().enumerate() {
                let mut acc = self.b[o];
                for kk in 0..self.k {
                    let up = ti + self.pad;
                    if up < kk || !(up - kk).is_multiple_of(self.stride) {
                        continue;
                    }
                    let src = (up - kk) / self.stride;
                    if src >= t {
                        continue;
                    }
                    for ic in 0..self.in_ch {
                        acc += x[ic * t + src] * self.w[(ic * self.out_ch + o) * self.k + kk];
                    }
                }
                *oval = acc;
            }
        });
        (y, t_out)
    }
}

/// One HiFi-GAN MRF residual block (leaky-relu pre-activation, dilated + plain conv pairs).
struct ResBlock {
    convs1: Vec<Conv1d>,
    convs2: Vec<Conv1d>,
}
impl ResBlock {
    fn forward(&self, x: &[f32], ch: usize, t: usize, slope: f32) -> Vec<f32> {
        let lrelu = |v: &mut Vec<f32>| {
            for e in v.iter_mut() {
                if *e < 0.0 {
                    *e *= slope;
                }
            }
        };
        let mut x = x.to_vec();
        for (c1, c2) in self.convs1.iter().zip(&self.convs2) {
            let mut h = x.clone();
            lrelu(&mut h);
            let mut h = c1.forward(&h, t);
            lrelu(&mut h);
            let h = c2.forward(&h, t);
            for i in 0..ch * t {
                x[i] += h[i];
            }
        }
        x
    }
}

struct VitsEncBlock {
    q: Linear,
    k: Linear,
    v: Linear,
    out: Linear,
    rel_k: Vec<f32>, // [2·window+1, head_dim]
    rel_v: Vec<f32>,
    ln1: Norm,
    conv1: Conv1d, // FFN in (kernel 3, asymmetric pad)
    conv2: Conv1d,
    ln2: Norm,
}

pub struct Vits {
    vocab: HashMap<String, u32>,
    // text encoder
    embed: Vec<f32>, // [vocab, hidden]
    enc_blocks: Vec<VitsEncBlock>,
    project: Conv1d, // 1×1 → [2·flow_size]
    // duration predictor + flow + decoder
    dp: DurationPredictor,
    flow_layers: Vec<CouplingLayer>,
    decoder: VitsDecoder,
    speaking_rate: f32,
    hidden: usize,
    heads: usize,
    hd: usize,
    window: i64,
    flow_size: usize,
    sampling_rate: u32,
}

impl Vits {
    pub fn load(dir: &Path) -> Result<Self> {
        let v: serde_json::Value =
            serde_json::from_slice(&std::fs::read(dir.join("config.json"))?).context("config")?;
        let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        let st = LazySt::open(dir)?;
        let hidden = g("hidden_size");
        let heads = g("num_attention_heads");
        let ffn = g("ffn_dim");
        let ffk = g("ffn_kernel_size");
        let vocab = v
            .get("vocab")
            .and_then(|m| m.as_object())
            .context("config vocab")?
            .iter()
            .filter_map(|(k, val)| val.as_u64().map(|i| (k.clone(), i as u32)))
            .collect();
        let enc_blocks = (0..g("num_hidden_layers"))
            .map(|i| {
                let b = format!("text_encoder.encoder.layers.{i}");
                let mut conv1 = Conv1d::load(
                    &st,
                    &format!("{b}.feed_forward.conv_1"),
                    ffn,
                    hidden,
                    ffk,
                    1,
                    false,
                )?;
                let mut conv2 = Conv1d::load(
                    &st,
                    &format!("{b}.feed_forward.conv_2"),
                    hidden,
                    ffn,
                    ffk,
                    1,
                    false,
                )?;
                // HF pads left (k-1)/2, right k/2 — identical for odd kernels (k=3 here)
                conv1.pad_l = (ffk - 1) / 2;
                conv2.pad_l = (ffk - 1) / 2;
                Ok(VitsEncBlock {
                    q: Linear::load(&st, &format!("{b}.attention.q_proj"), hidden, hidden)?,
                    k: Linear::load(&st, &format!("{b}.attention.k_proj"), hidden, hidden)?,
                    v: Linear::load(&st, &format!("{b}.attention.v_proj"), hidden, hidden)?,
                    out: Linear::load(&st, &format!("{b}.attention.out_proj"), hidden, hidden)?,
                    rel_k: st.tensor_f32(&format!("{b}.attention.emb_rel_k"))?,
                    rel_v: st.tensor_f32(&format!("{b}.attention.emb_rel_v"))?,
                    ln1: Norm::load(&st, &format!("{b}.layer_norm"))?,
                    conv1,
                    conv2,
                    ln2: Norm::load(&st, &format!("{b}.final_layer_norm"))?,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        // stochastic duration predictor (reverse-path modules only)
        let (dpk, dds_layers) = (
            g("duration_predictor_kernel_size"),
            g("depth_separable_num_layers"),
        );
        let bins = g("duration_predictor_flow_bins");
        let tail = v
            .get("duration_predictor_tail_bound")
            .and_then(|x| x.as_f64())
            .unwrap_or(5.0) as f32;
        let dp_flows = (2..=g("duration_predictor_num_flows"))
            .map(|i| {
                // ModuleList: [EA, CF0..CF3]; the reverse path uses CF3,CF2,CF1 (CF0 = HF's
                // "useless vflow", never evaluated at inference) — load indices 2..=4
                let b = format!("duration_predictor.flows.{i}");
                Ok(ConvFlow {
                    conv_pre: Conv1d::load(&st, &format!("{b}.conv_pre"), hidden, 1, 1, 1, false)?,
                    dds: DdsConv::load(&st, &format!("{b}.conv_dds"), hidden, dpk, dds_layers)?,
                    conv_proj: Conv1d::load(
                        &st,
                        &format!("{b}.conv_proj"),
                        3 * bins - 1,
                        hidden,
                        1,
                        1,
                        false,
                    )?,
                    bins,
                    tail,
                    filter: hidden,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        let dp = DurationPredictor {
            conv_pre: Conv1d::load(
                &st,
                "duration_predictor.conv_pre",
                hidden,
                hidden,
                1,
                1,
                false,
            )?,
            dds: DdsConv::load(&st, "duration_predictor.conv_dds", hidden, dpk, dds_layers)?,
            conv_proj: Conv1d::load(
                &st,
                "duration_predictor.conv_proj",
                hidden,
                hidden,
                1,
                1,
                false,
            )?,
            ea_translate: st.tensor_f32("duration_predictor.flows.0.translate")?,
            ea_log_scale: st.tensor_f32("duration_predictor.flows.0.log_scale")?,
            flows: dp_flows,
        };
        // residual coupling flow
        let flow_size = g("flow_size");
        let half = flow_size / 2;
        let (wk, wrate, wlayers) = (
            g("wavenet_kernel_size"),
            g("wavenet_dilation_rate"),
            g("prior_encoder_num_wavenet_layers"),
        );
        let flow_layers = (0..g("prior_encoder_num_flows"))
            .map(|i| {
                let b = format!("flow.flows.{i}");
                Ok(CouplingLayer {
                    conv_pre: Conv1d::load(
                        &st,
                        &format!("{b}.conv_pre"),
                        hidden,
                        half,
                        1,
                        1,
                        false,
                    )?,
                    in_layers: (0..wlayers)
                        .map(|l| {
                            Conv1d::load(
                                &st,
                                &format!("{b}.wavenet.in_layers.{l}"),
                                2 * hidden,
                                hidden,
                                wk,
                                wrate.pow(l as u32),
                                false,
                            )
                        })
                        .collect::<Result<Vec<_>>>()?,
                    res_skip: (0..wlayers)
                        .map(|l| {
                            let out = if l < wlayers - 1 { 2 * hidden } else { hidden };
                            Conv1d::load(
                                &st,
                                &format!("{b}.wavenet.res_skip_layers.{l}"),
                                out,
                                hidden,
                                1,
                                1,
                                false,
                            )
                        })
                        .collect::<Result<Vec<_>>>()?,
                    conv_post: Conv1d::load(
                        &st,
                        &format!("{b}.conv_post"),
                        half,
                        hidden,
                        1,
                        1,
                        false,
                    )?,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        // HiFi-GAN decoder
        let arr = |k: &str| -> Vec<usize> {
            v.get(k)
                .and_then(|x| x.as_array())
                .map(|a| {
                    a.iter()
                        .filter_map(|u| u.as_u64().map(|u| u as usize))
                        .collect()
                })
                .unwrap_or_default()
        };
        let ch0 = g("upsample_initial_channel");
        let rates = arr("upsample_rates");
        let kernels = arr("upsample_kernel_sizes");
        let res_k = arr("resblock_kernel_sizes");
        let dils: Vec<Vec<usize>> = v
            .get("resblock_dilation_sizes")
            .and_then(|x| x.as_array())
            .map(|a| {
                a.iter()
                    .map(|row| {
                        row.as_array()
                            .map(|r| {
                                r.iter()
                                    .filter_map(|u| u.as_u64().map(|u| u as usize))
                                    .collect()
                            })
                            .unwrap_or_default()
                    })
                    .collect()
            })
            .unwrap_or_default();
        let upsampler = rates
            .iter()
            .zip(&kernels)
            .enumerate()
            .map(|(i, (&r, &k))| {
                ConvT1d::load(
                    &st,
                    &format!("decoder.upsampler.{i}"),
                    ch0 >> i,
                    ch0 >> (i + 1),
                    k,
                    r,
                )
            })
            .collect::<Result<Vec<_>>>()?;
        let mut resblocks = Vec::new();
        for i in 0..rates.len() {
            let ch = ch0 >> (i + 1);
            for (j, (&k, dil)) in res_k.iter().zip(&dils).enumerate() {
                let b = format!("decoder.resblocks.{}", i * res_k.len() + j);
                resblocks.push(ResBlock {
                    convs1: dil
                        .iter()
                        .enumerate()
                        .map(|(n, &d)| {
                            Conv1d::load(&st, &format!("{b}.convs1.{n}"), ch, ch, k, d, false)
                        })
                        .collect::<Result<Vec<_>>>()?,
                    convs2: dil
                        .iter()
                        .enumerate()
                        .map(|(n, _)| {
                            Conv1d::load(&st, &format!("{b}.convs2.{n}"), ch, ch, k, 1, false)
                        })
                        .collect::<Result<Vec<_>>>()?,
                });
            }
        }
        let decoder = VitsDecoder {
            conv_pre: Conv1d::load(&st, "decoder.conv_pre", ch0, flow_size, 7, 1, false)?,
            upsampler,
            resblocks,
            conv_post: Conv1d::load(&st, "decoder.conv_post", 1, ch0 >> rates.len(), 7, 1, false)?,
            num_kernels: res_k.len(),
            slope: v
                .get("leaky_relu_slope")
                .and_then(|x| x.as_f64())
                .unwrap_or(0.1) as f32,
        };
        Ok(Self {
            vocab,
            embed: st.tensor_f32("text_encoder.embed_tokens.weight")?,
            enc_blocks,
            project: Conv1d::load(
                &st,
                "text_encoder.project",
                2 * flow_size,
                hidden,
                1,
                1,
                false,
            )?,
            dp,
            flow_layers,
            decoder,
            speaking_rate: v
                .get("speaking_rate")
                .and_then(|x| x.as_f64())
                .unwrap_or(1.0) as f32,
            hidden,
            heads,
            hd: hidden / heads,
            window: g("window_size") as i64,
            flow_size,
            sampling_rate: g("sampling_rate") as u32,
        })
    }

    pub fn sampling_rate(&self) -> u32 {
        self.sampling_rate
    }

    /// Text → blank-interspersed ids, matching `VitsTokenizer` (normalize=True, add_blank=True):
    /// keep a char verbatim if it is in the vocab, otherwise lowercase it; drop anything still
    /// unknown; trim; intersperse token id 0 between/around all characters.
    pub fn tokenize(&self, text: &str) -> Vec<u32> {
        let mut kept: Vec<u32> = Vec::with_capacity(text.len());
        let mut chars: Vec<char> = Vec::with_capacity(text.len());
        for ch in text.chars() {
            let s = ch.to_string();
            if let Some(&id) = self.vocab.get(s.as_str()) {
                kept.push(id);
                chars.push(ch);
            } else {
                for lc in ch.to_lowercase() {
                    if let Some(&id) = self.vocab.get(lc.to_string().as_str()) {
                        kept.push(id);
                        chars.push(lc);
                    }
                }
            }
        }
        // strip(): trim leading/trailing whitespace survivors (' ' is in the vocab)
        let start = chars
            .iter()
            .position(|c| !c.is_whitespace())
            .unwrap_or(chars.len());
        let end = chars
            .iter()
            .rposition(|c| !c.is_whitespace())
            .map_or(0, |e| e + 1);
        let kept = &kept[start..end.max(start)];
        let mut ids = Vec::with_capacity(kept.len() * 2 + 1);
        ids.push(0);
        for &id in kept {
            ids.push(id);
            ids.push(0);
        }
        ids
    }

    /// Windowed relative self-attention (`emb_rel_k`/`emb_rel_v`, zero outside ±window):
    /// `score(i,j) = q_i·k_j·s + q_i·s·rel_k[j−i]`, `out_i = Σ p_ij·(v_j + rel_v[j−i])`.
    fn enc_attention(&self, blk: &VitsEncBlock, x: &[f32], t: usize) -> Vec<f32> {
        let (h, heads, hd, w) = (self.hidden, self.heads, self.hd, self.window);
        let q = blk.q.forward(x, t);
        let k = blk.k.forward(x, t);
        let v = blk.v.forward(x, t);
        let scale = 1.0 / (hd as f32).sqrt();
        let ctx: Vec<Vec<f32>> = (0..heads)
            .into_par_iter()
            .map(|head| {
                let mut oh = vec![0f32; t * hd];
                let mut srow = vec![0f32; t];
                for i in 0..t {
                    let qi = &q[i * h + head * hd..][..hd];
                    for j in 0..t {
                        let kj = &k[j * h + head * hd..][..hd];
                        let mut dot = 0f32;
                        for c in 0..hd {
                            dot += qi[c] * kj[c];
                        }
                        let rel = j as i64 - i as i64;
                        if rel.abs() <= w {
                            let re = &blk.rel_k[(rel + w) as usize * hd..][..hd];
                            for c in 0..hd {
                                dot += qi[c] * re[c];
                            }
                        }
                        srow[j] = dot * scale;
                    }
                    let max = srow.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                    let mut sum = 0.0;
                    for s in srow.iter_mut() {
                        *s = (*s - max).exp();
                        sum += *s;
                    }
                    let inv = 1.0 / sum;
                    let orow = &mut oh[i * hd..][..hd];
                    for j in 0..t {
                        let p = srow[j] * inv;
                        let vj = &v[j * h + head * hd..][..hd];
                        for c in 0..hd {
                            orow[c] += p * vj[c];
                        }
                        let rel = j as i64 - i as i64;
                        if rel.abs() <= w {
                            let re = &blk.rel_v[(rel + w) as usize * hd..][..hd];
                            for c in 0..hd {
                                orow[c] += p * re[c];
                            }
                        }
                    }
                }
                oh
            })
            .collect();
        let mut merged = vec![0f32; t * h];
        for (head, oh) in ctx.iter().enumerate() {
            for i in 0..t {
                merged[i * h + head * hd..][..hd].copy_from_slice(&oh[i * hd..][..hd]);
            }
        }
        blk.out.forward(&merged, t)
    }

    /// ids → (encoder hidden `[t, hidden]`, prior means `[t, flow]`, prior log-variances).
    pub fn encode_text(&self, ids: &[u32]) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
        let (h, t) = (self.hidden, ids.len());
        let scale = (h as f32).sqrt();
        let mut x = vec![0f32; t * h];
        for (i, &id) in ids.iter().enumerate() {
            let emb = &self.embed[id as usize * h..][..h];
            for c in 0..h {
                x[i * h + c] = emb[c] * scale;
            }
        }
        for blk in &self.enc_blocks {
            let a = self.enc_attention(blk, &x, t);
            for (xi, ai) in x.iter_mut().zip(a) {
                *xi += ai;
            }
            x = blk.ln1.forward(&x, h);
            // conv FFN operates channel-major
            let mut cm = vec![0f32; h * t];
            for i in 0..t {
                for c in 0..h {
                    cm[c * t + i] = x[i * h + c];
                }
            }
            let mut mid = blk.conv1.forward(&cm, t);
            for m in mid.iter_mut() {
                *m = m.max(0.0); // hidden_act relu
            }
            let f = blk.conv2.forward(&mid, t);
            for i in 0..t {
                for c in 0..h {
                    x[i * h + c] += f[c * t + i];
                }
            }
            x = blk.ln2.forward(&x, h);
        }
        // project (1×1 conv) → split means / log-variances, back to time-major
        let mut cm = vec![0f32; h * t];
        for i in 0..t {
            for c in 0..h {
                cm[c * t + i] = x[i * h + c];
            }
        }
        let stats = self.project.forward(&cm, t);
        let fs = self.flow_size;
        let mut means = vec![0f32; t * fs];
        let mut logvar = vec![0f32; t * fs];
        for i in 0..t {
            for c in 0..fs {
                means[i * fs + c] = stats[c * t + i];
                logvar[i * fs + c] = stats[(fs + c) * t + i];
            }
        }
        (x, means, logvar)
    }

    /// Encoder hidden (time-major `[t, hidden]`) → (log-durations `[t]`, per-token frame counts).
    /// Deterministic: the duration flow runs in reverse from ZERO latents (noise scale 0).
    pub fn predict_durations(&self, hidden: &[f32]) -> (Vec<f32>, Vec<usize>) {
        let h = self.hidden;
        let t = hidden.len() / h;
        let mut cm = vec![0f32; h * t];
        for i in 0..t {
            for c in 0..h {
                cm[c * t + i] = hidden[i * h + c];
            }
        }
        let log_dur = self.dp.log_durations(&cm, h, t);
        let length_scale = 1.0 / self.speaking_rate;
        let durations = log_dur
            .iter()
            .map(|&ld| (ld.exp() * length_scale).ceil() as usize)
            .collect();
        (log_dur, durations)
    }

    /// Reverse residual-coupling flow: prior latents (time-major `[t_out, flow]`) → decoder
    /// latents, same shape. Channel-flip before every coupling layer, layers in reverse order.
    pub fn flow_reverse(&self, z_tm: &[f32]) -> Vec<f32> {
        let fs = self.flow_size;
        let t = z_tm.len() / fs;
        let mut z = vec![0f32; fs * t];
        for i in 0..t {
            for c in 0..fs {
                z[c * t + i] = z_tm[i * fs + c];
            }
        }
        for layer in self.flow_layers.iter().rev() {
            // torch.flip on the channel dim
            for c in 0..fs / 2 {
                for i in 0..t {
                    z.swap(c * t + i, (fs - 1 - c) * t + i);
                }
            }
            layer.reverse(&mut z, fs / 2, self.hidden, t);
        }
        let mut out = vec![0f32; t * fs];
        for i in 0..t {
            for c in 0..fs {
                out[i * fs + c] = z[c * t + i];
            }
        }
        out
    }

    /// HiFi-GAN decode: latents (time-major `[t_out, flow]`) → waveform samples.
    pub fn decode(&self, latents_tm: &[f32]) -> Vec<f32> {
        let fs = self.flow_size;
        let t = latents_tm.len() / fs;
        let mut cm = vec![0f32; fs * t];
        for i in 0..t {
            for c in 0..fs {
                cm[c * t + i] = latents_tm[i * fs + c];
            }
        }
        self.decoder.forward(&cm, t)
    }

    /// End-to-end: text → waveform. Deterministic (both VITS noise scales are 0 by design).
    pub fn synthesize(&self, text: &str) -> Vec<f32> {
        let ids = self.tokenize(text);
        let (hidden, means, _logvar) = self.encode_text(&ids);
        let (_, durations) = self.predict_durations(&hidden);
        let fs = self.flow_size;
        let t_out: usize = durations.iter().sum::<usize>().max(1);
        let mut z = Vec::with_capacity(t_out * fs);
        for (i, &d) in durations.iter().enumerate() {
            for _ in 0..d {
                z.extend_from_slice(&means[i * fs..][..fs]);
            }
        }
        let latents = self.flow_reverse(&z);
        self.decode(&latents)
    }
}