inferencelayer 0.2.4

//! Kokoro-82M (`hexgrad/Kokoro-82M`, StyleTTS2/ISTFTNet lineage) — phonemes + voice style →
//! 24 kHz waveform. Pure Rust, CPU. The engine's third TTS architecture and its most composite:
//!
//!   PLBERT (ALBERT, one shared layer ×12) → duration encoder (3× BiLSTM + AdaLayerNorm, style-
//!   conditioned) → durations (sigmoid-sum, round, ≥1) → frame expansion → shared BiLSTM →
//!   F0/N chains (AdaIN residual blocks, 2× upsample) → text encoder (CNN + BiLSTM) →
//!   ISTFTNet decoder: AdaIN-Snake resblocks + harmonic source (SineGen → 20-point STFT) →
//!   spec/phase head → inverse STFT.
//!
//! Gated per stage against the `kokoro` reference package (`tests/fixtures/export_kokoro.py`).
//! Like KModel, this is language-blind: input is an IPA phoneme string (G2P lives upstream) plus
//! a 256-d voice style row (`ref_s`, selected from the voice pack by phoneme count).
//!
//! Determinism: the reference SineGen draws random initial harmonic phases and additive noise;
//! the oracle zeroes both and this port generates none — synthesis is deterministic by design.
//!
//! Numerics: every stage is bit-exact vs the reference (cos 1.0) EXCEPT the harmonic source
//! ([`Kokoro::sine_source`]). ISTFTNet takes sin() of cumulative phases up to ~1.2e5 rad, where
//! f32's ULP (≈0.016) exceeds sin's stability, so torch's fused-kernel f32 phase rounding is not
//! portably reproducible. This port computes that phase in f64 (the mathematically-intended value);
//! it diverges from torch's f32-noisy source only in the phase of harmonics ABOVE the fundamental
//! — perceptually inaudible. The rest of the generator is bit-exact (the parity gate proves it by
//! feeding the reference's own source STFT through it: cos 1.0).

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

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

use crate::weights::LazySt;

const STYLE_DIM: usize = 128;
const SR: f32 = 24000.0;

/// `y[m,n] = x[m,k]·Wᵀ + b` (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
}

fn lrelu(v: &mut [f32], slope: f32) {
    for e in v.iter_mut() {
        if *e < 0.0 {
            *e *= slope;
        }
    }
}

/// tanh-approximation GELU (`gelu_new` — ALBERT's default).
fn gelu_new(v: f32) -> f32 {
    0.5 * v * (1.0 + (0.797_884_6 * (v + 0.044715 * v * v * v)).tanh())
}

fn sigmoid(v: f32) -> f32 {
    1.0 / (1.0 + (-v).exp())
}

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"))?;
        anyhow::ensure!(w.len() == n * k, "{prefix}.weight {} != {n}x{k}", w.len());
        Ok(Self {
            w,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
            n,
            k,
        })
    }
    fn forward(&self, x: &[f32], m: usize) -> Vec<f32> {
        linear(x, &self.w, &self.b, m, self.n, self.k)
    }
}

/// Conv1d over channel-major `[in, t]`: stride, dilation, groups ∈ {1, in}, optional bias.
struct Conv1d {
    w: Vec<f32>,
    b: Vec<f32>, // empty = no bias
    out_ch: usize,
    in_ch: usize,
    k: usize,
    stride: usize,
    dilation: usize,
    pad: usize,
    depthwise: bool,
}
impl Conv1d {
    #[allow(clippy::too_many_arguments)]
    fn load(
        st: &LazySt,
        prefix: &str,
        out_ch: usize,
        in_ch: usize,
        k: usize,
        stride: 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()
        );
        Ok(Self {
            w,
            b: st.tensor_f32(&format!("{prefix}.bias")).unwrap_or_default(),
            out_ch,
            in_ch,
            k,
            stride,
            dilation,
            pad: (k * dilation - dilation) / 2,
            depthwise,
        })
    }
    fn out_len(&self, t: usize) -> usize {
        (t + 2 * self.pad - self.dilation * (self.k - 1) - 1) / self.stride + 1
    }
    fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let to = self.out_len(t);
        let mut y = vec![0f32; self.out_ch * to];
        y.par_chunks_mut(to).enumerate().for_each(|(o, orow)| {
            let bias = self.b.get(o).copied().unwrap_or(0.0);
            let (ic0, nic, wbase) = if self.depthwise {
                (o, 1, o * self.k)
            } else {
                (0, self.in_ch, o * self.in_ch * self.k)
            };
            for (ti, oval) in orow.iter_mut().enumerate() {
                let base = ti * self.stride;
                let mut acc = bias;
                for wi in 0..nic {
                    let xrow = &x[(ic0 + wi) * t..][..t];
                    let wrow = &self.w[wbase + wi * self.k..][..self.k];
                    for (kk, wv) in wrow.iter().enumerate() {
                        let src = base + kk * self.dilation;
                        if src >= self.pad && src - self.pad < t {
                            acc += xrow[src - self.pad] * wv;
                        }
                    }
                }
                *oval = acc;
            }
        });
        y
    }
}

/// ConvTranspose1d, weights `[in, out/groups, k]`; groups ∈ {1, in} (depthwise), output_padding.
struct ConvT1d {
    w: Vec<f32>,
    b: Vec<f32>,
    out_ch: usize,
    in_ch: usize,
    k: usize,
    stride: usize,
    pad: usize,
    out_pad: usize,
    depthwise: bool,
}
impl ConvT1d {
    #[allow(clippy::too_many_arguments)]
    fn load(
        st: &LazySt,
        prefix: &str,
        in_ch: usize,
        out_ch: usize,
        k: usize,
        stride: usize,
        pad: usize,
        out_pad: usize,
        depthwise: bool,
    ) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let per = if depthwise { 1 } else { out_ch };
        anyhow::ensure!(
            w.len() == in_ch * per * k,
            "{prefix}.weight {} != {in_ch}x{per}x{k}",
            w.len()
        );
        Ok(Self {
            w,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
            out_ch,
            in_ch,
            k,
            stride,
            pad,
            out_pad,
            depthwise,
        })
    }
    fn forward(&self, x: &[f32], t: usize) -> (Vec<f32>, usize) {
        let to = (t - 1) * self.stride + self.k - 2 * self.pad + self.out_pad;
        let mut y = vec![0f32; self.out_ch * to];
        y.par_chunks_mut(to).enumerate().for_each(|(o, orow)| {
            for (ti, oval) in orow.iter_mut().enumerate() {
                let mut acc = self.b[o];
                let up = ti + self.pad;
                for kk in 0..self.k {
                    if up < kk || !(up - kk).is_multiple_of(self.stride) {
                        continue;
                    }
                    let src = (up - kk) / self.stride;
                    if src >= t {
                        continue;
                    }
                    if self.depthwise {
                        acc += x[o * t + src] * self.w[o * self.k + kk];
                    } else {
                        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, to)
    }
}

/// Single-layer bidirectional LSTM (torch gate order i,f,g,o), time-major in/out `[t, dims]`.
struct BiLstm {
    w_ih: Vec<f32>, // [4H, in]
    w_hh: Vec<f32>, // [4H, H]
    b: Vec<f32>,    // [4H] (b_ih + b_hh, pre-summed)
    w_ih_r: Vec<f32>,
    w_hh_r: Vec<f32>,
    b_r: Vec<f32>,
    input: usize,
    hidden: usize,
}
impl BiLstm {
    fn load(st: &LazySt, prefix: &str, input: usize, hidden: usize) -> Result<Self> {
        let sum =
            |a: Vec<f32>, b: Vec<f32>| a.iter().zip(&b).map(|(x, y)| x + y).collect::<Vec<f32>>();
        let w_ih = st.tensor_f32(&format!("{prefix}.weight_ih_l0"))?;
        anyhow::ensure!(w_ih.len() == 4 * hidden * input, "{prefix} lstm dims");
        Ok(Self {
            w_ih,
            w_hh: st.tensor_f32(&format!("{prefix}.weight_hh_l0"))?,
            b: sum(
                st.tensor_f32(&format!("{prefix}.bias_ih_l0"))?,
                st.tensor_f32(&format!("{prefix}.bias_hh_l0"))?,
            ),
            w_ih_r: st.tensor_f32(&format!("{prefix}.weight_ih_l0_reverse"))?,
            w_hh_r: st.tensor_f32(&format!("{prefix}.weight_hh_l0_reverse"))?,
            b_r: sum(
                st.tensor_f32(&format!("{prefix}.bias_ih_l0_reverse"))?,
                st.tensor_f32(&format!("{prefix}.bias_hh_l0_reverse"))?,
            ),
            input,
            hidden,
        })
    }
    fn run_dir(&self, x: &[f32], t: usize, reverse: bool) -> Vec<f32> {
        let (h, kin) = (self.hidden, self.input);
        let (w_ih, w_hh, bias) = if reverse {
            (&self.w_ih_r, &self.w_hh_r, &self.b_r)
        } else {
            (&self.w_ih, &self.w_hh, &self.b)
        };
        let mut out = vec![0f32; t * h];
        let mut hs = vec![0f32; h];
        let mut cs = vec![0f32; h];
        let mut gates = vec![0f32; 4 * h];
        for step in 0..t {
            let ti = if reverse { t - 1 - step } else { step };
            let xr = &x[ti * kin..][..kin];
            gates.par_chunks_mut(1).enumerate().for_each(|(g, gv)| {
                let mut acc = bias[g];
                let wr = &w_ih[g * kin..][..kin];
                for c in 0..kin {
                    acc += xr[c] * wr[c];
                }
                let hr = &w_hh[g * h..][..h];
                for c in 0..h {
                    acc += hs[c] * hr[c];
                }
                gv[0] = acc;
            });
            for c in 0..h {
                let i = sigmoid(gates[c]);
                let f = sigmoid(gates[h + c]);
                let g = gates[2 * h + c].tanh();
                let o = sigmoid(gates[3 * h + c]);
                cs[c] = f * cs[c] + i * g;
                hs[c] = o * cs[c].tanh();
            }
            out[ti * h..][..h].copy_from_slice(&hs);
        }
        out
    }
    /// `[t, input]` → `[t, 2·hidden]` (forward ‖ backward).
    fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let (fwd, bwd) = rayon::join(|| self.run_dir(x, t, false), || self.run_dir(x, t, true));
        let h = self.hidden;
        let mut out = vec![0f32; t * 2 * h];
        for ti in 0..t {
            out[ti * 2 * h..][..h].copy_from_slice(&fwd[ti * h..][..h]);
            out[ti * 2 * h + h..][..h].copy_from_slice(&bwd[ti * h..][..h]);
        }
        out
    }
}

/// AdaIN1d: InstanceNorm1d(affine) over time per channel, then `(1+γ)·x̂ + β` from a style FC.
struct AdaIn {
    norm_w: Vec<f32>,
    norm_b: Vec<f32>,
    fc: Linear, // style → [2·ch]
    ch: usize,
}
impl AdaIn {
    fn load(st: &LazySt, prefix: &str, ch: usize) -> Result<Self> {
        Ok(Self {
            norm_w: st.tensor_f32(&format!("{prefix}.norm.weight"))?,
            norm_b: st.tensor_f32(&format!("{prefix}.norm.bias"))?,
            fc: Linear::load(st, &format!("{prefix}.fc"), 2 * ch, STYLE_DIM)?,
            ch,
        })
    }
    fn forward(&self, x: &[f32], t: usize, s: &[f32]) -> Vec<f32> {
        let h = self.fc.forward(s, 1);
        let mut out = vec![0f32; x.len()];
        out.par_chunks_mut(t).enumerate().for_each(|(c, orow)| {
            let row = &x[c * t..][..t];
            let mean = row.iter().sum::<f32>() / t as f32;
            let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / t as f32;
            let inv = 1.0 / (var + 1e-5).sqrt();
            let (gamma, beta) = (h[c], h[self.ch + c]);
            for (o, v) in orow.iter_mut().zip(row) {
                let normed = (v - mean) * inv * self.norm_w[c] + self.norm_b[c];
                *o = (1.0 + gamma) * normed + beta;
            }
        });
        out
    }
}

/// `AdainResBlk1d`: AdaIN → lrelu(0.2) → (optional ×2 depthwise-ConvT pool) → conv → AdaIN →
/// lrelu → conv; shortcut nearest-×2 (+1×1 if dims change); sum ÷ √2.
struct AdainResBlk {
    norm1: AdaIn,
    norm2: AdaIn,
    conv1: Conv1d,
    conv2: Conv1d,
    conv1x1: Option<Conv1d>,
    pool: Option<ConvT1d>,
    dim_in: usize,
}
impl AdainResBlk {
    fn load(
        st: &LazySt,
        prefix: &str,
        dim_in: usize,
        dim_out: usize,
        upsample: bool,
    ) -> Result<Self> {
        Ok(Self {
            norm1: AdaIn::load(st, &format!("{prefix}.norm1"), dim_in)?,
            norm2: AdaIn::load(st, &format!("{prefix}.norm2"), dim_out)?,
            conv1: Conv1d::load(
                st,
                &format!("{prefix}.conv1"),
                dim_out,
                dim_in,
                3,
                1,
                1,
                false,
            )?,
            conv2: Conv1d::load(
                st,
                &format!("{prefix}.conv2"),
                dim_out,
                dim_out,
                3,
                1,
                1,
                false,
            )?,
            conv1x1: (dim_in != dim_out)
                .then(|| {
                    Conv1d::load(
                        st,
                        &format!("{prefix}.conv1x1"),
                        dim_out,
                        dim_in,
                        1,
                        1,
                        1,
                        false,
                    )
                })
                .transpose()?,
            pool: upsample
                .then(|| {
                    ConvT1d::load(
                        st,
                        &format!("{prefix}.pool"),
                        dim_in,
                        dim_in,
                        3,
                        2,
                        1,
                        1,
                        true,
                    )
                })
                .transpose()?,
            dim_in,
        })
    }
    fn forward(&self, x: &[f32], t: usize, s: &[f32]) -> (Vec<f32>, usize) {
        // residual branch
        let mut r = self.norm1.forward(x, t, s);
        lrelu(&mut r, 0.2);
        let (r, t_out) = match &self.pool {
            Some(p) => p.forward(&r, t),
            None => (r, t),
        };
        let r = self.conv1.forward(&r, t_out);
        let mut r = self.norm2.forward(&r, t_out, s);
        lrelu(&mut r, 0.2);
        let r = self.conv2.forward(&r, t_out);
        // shortcut branch
        let short = if self.pool.is_some() {
            let mut up = vec![0f32; self.dim_in * t_out];
            for c in 0..self.dim_in {
                for ti in 0..t_out {
                    up[c * t_out + ti] = x[c * t + ti / 2]; // nearest ×2
                }
            }
            up
        } else {
            x.to_vec()
        };
        let short = match &self.conv1x1 {
            Some(c) => c.forward(&short, t_out),
            None => short,
        };
        let inv_sqrt2 = 1.0 / 2f32.sqrt();
        let out = r
            .iter()
            .zip(&short)
            .map(|(a, b)| (a + b) * inv_sqrt2)
            .collect();
        (out, t_out)
    }
}

/// `AdaINResBlock1` (the ISTFTNet MRF block): 3× [AdaIN → Snake → dilated conv → AdaIN → Snake →
/// conv] with residual adds. Snake: `x + sin²(αx)/α`.
struct SnakeResBlock {
    convs1: Vec<Conv1d>,
    convs2: Vec<Conv1d>,
    adain1: Vec<AdaIn>,
    adain2: Vec<AdaIn>,
    alpha1: Vec<Vec<f32>>, // [ch] each
    alpha2: Vec<Vec<f32>>,
}
impl SnakeResBlock {
    fn load(st: &LazySt, prefix: &str, ch: usize, k: usize, dilations: &[usize]) -> Result<Self> {
        let mut b = Self {
            convs1: Vec::new(),
            convs2: Vec::new(),
            adain1: Vec::new(),
            adain2: Vec::new(),
            alpha1: Vec::new(),
            alpha2: Vec::new(),
        };
        for (i, &d) in dilations.iter().enumerate() {
            b.convs1.push(Conv1d::load(
                st,
                &format!("{prefix}.convs1.{i}"),
                ch,
                ch,
                k,
                1,
                d,
                false,
            )?);
            b.convs2.push(Conv1d::load(
                st,
                &format!("{prefix}.convs2.{i}"),
                ch,
                ch,
                k,
                1,
                1,
                false,
            )?);
            b.adain1
                .push(AdaIn::load(st, &format!("{prefix}.adain1.{i}"), ch)?);
            b.adain2
                .push(AdaIn::load(st, &format!("{prefix}.adain2.{i}"), ch)?);
            b.alpha1
                .push(st.tensor_f32(&format!("{prefix}.alpha1.{i}"))?);
            b.alpha2
                .push(st.tensor_f32(&format!("{prefix}.alpha2.{i}"))?);
        }
        Ok(b)
    }
    fn snake(x: &mut [f32], alpha: &[f32], t: usize) {
        for (c, a) in alpha.iter().enumerate() {
            for v in x[c * t..][..t].iter_mut() {
                let s = (a * *v).sin();
                *v += s * s / a;
            }
        }
    }
    fn forward(&self, x: &[f32], t: usize, s: &[f32]) -> Vec<f32> {
        let mut x = x.to_vec();
        for i in 0..self.convs1.len() {
            let mut xt = self.adain1[i].forward(&x, t, s);
            Self::snake(&mut xt, &self.alpha1[i], t);
            let xt = self.convs1[i].forward(&xt, t);
            let mut xt = self.adain2[i].forward(&xt, t, s);
            Self::snake(&mut xt, &self.alpha2[i], t);
            let xt = self.convs2[i].forward(&xt, t);
            for (a, b) in x.iter_mut().zip(xt) {
                *a += b;
            }
        }
        x
    }
}

struct AlbertLayer {
    q: Linear,
    k: Linear,
    v: Linear,
    dense: Linear,
    attn_ln_w: Vec<f32>,
    attn_ln_b: Vec<f32>,
    ffn: Linear,
    ffn_out: Linear,
    full_ln_w: Vec<f32>,
    full_ln_b: Vec<f32>,
}

/// The style-conditioned duration encoder: alternating BiLSTM / AdaLayerNorm, style concatenated
/// back after every stage.
struct DurationEncoder {
    lstms: Vec<BiLstm>,  // input d_model+style → d_model
    ada_fc: Vec<Linear>, // style → 2·d_model (AdaLayerNorm has NO learned norm params)
}

pub struct Kokoro {
    vocab: HashMap<String, u32>,
    // PLBERT
    word_emb: Vec<f32>, // [n_token, emb]
    pos_emb: Vec<f32>,
    type_emb: Vec<f32>,
    emb_ln_w: Vec<f32>,
    emb_ln_b: Vec<f32>,
    emb_map: Linear, // emb → bert hidden
    albert: AlbertLayer,
    bert_layers: usize,
    bert_hidden: usize,
    bert_heads: usize,
    emb_dim: usize,
    bert_encoder: Linear, // bert hidden → hidden_dim
    // prosody predictor
    dur_enc: DurationEncoder,
    lstm: BiLstm, // hidden+style → hidden
    duration_proj: Linear,
    shared: BiLstm,
    f0_blocks: Vec<AdainResBlk>,
    n_blocks: Vec<AdainResBlk>,
    f0_proj: Conv1d,
    n_proj: Conv1d,
    // text encoder
    te_emb: Vec<f32>,                          // [n_token, hidden]
    te_cnn: Vec<(Conv1d, Vec<f32>, Vec<f32>)>, // conv, ln gamma, ln beta (LN over channels)
    te_lstm: BiLstm,
    // decoder
    f0_conv: Conv1d,
    n_conv: Conv1d,
    encode: AdainResBlk,
    asr_res: Conv1d,
    decode: Vec<AdainResBlk>,
    // generator
    ups: Vec<ConvT1d>,
    noise_convs: Vec<Conv1d>,
    noise_res: Vec<SnakeResBlock>,
    resblocks: Vec<SnakeResBlock>,
    conv_post: Conv1d,
    m_source_linear: Linear,  // 9 harmonics → 1
    unvoiced_phase: Vec<f32>, // reference STFT phase column for constant frames (see export)
    n_fft: usize,
    hop: usize,
    num_kernels: usize,
    upsample_total: usize, // ∏ upsample_rates · hop
    hidden: usize,
}

impl Kokoro {
    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 vocab: HashMap<String, u32> = 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 p = v.get("plbert_effective").context("plbert_effective")?;
        let pg = |k: &str| p.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        let (emb_dim, bh) = (pg("embedding_size"), pg("hidden_size"));
        let hidden = g("hidden_dim");
        let al = "bert.encoder.albert_layer_groups.0.albert_layers.0";
        let albert = AlbertLayer {
            q: Linear::load(&st, &format!("{al}.attention.query"), bh, bh)?,
            k: Linear::load(&st, &format!("{al}.attention.key"), bh, bh)?,
            v: Linear::load(&st, &format!("{al}.attention.value"), bh, bh)?,
            dense: Linear::load(&st, &format!("{al}.attention.dense"), bh, bh)?,
            attn_ln_w: st.tensor_f32(&format!("{al}.attention.LayerNorm.weight"))?,
            attn_ln_b: st.tensor_f32(&format!("{al}.attention.LayerNorm.bias"))?,
            ffn: Linear::load(&st, &format!("{al}.ffn"), pg("intermediate_size"), bh)?,
            ffn_out: Linear::load(
                &st,
                &format!("{al}.ffn_output"),
                bh,
                pg("intermediate_size"),
            )?,
            full_ln_w: st.tensor_f32(&format!("{al}.full_layer_layer_norm.weight"))?,
            full_ln_b: st.tensor_f32(&format!("{al}.full_layer_layer_norm.bias"))?,
        };
        let n_layer = g("n_layer");
        let dur_enc = DurationEncoder {
            lstms: (0..n_layer)
                .map(|i| {
                    BiLstm::load(
                        &st,
                        &format!("predictor.text_encoder.lstms.{}", 2 * i),
                        hidden + STYLE_DIM,
                        hidden / 2,
                    )
                })
                .collect::<Result<Vec<_>>>()?,
            ada_fc: (0..n_layer)
                .map(|i| {
                    Linear::load(
                        &st,
                        &format!("predictor.text_encoder.lstms.{}.fc", 2 * i + 1),
                        2 * hidden,
                        STYLE_DIM,
                    )
                })
                .collect::<Result<Vec<_>>>()?,
        };
        let half = hidden / 2;
        let load_chain = |name: &str| -> Result<Vec<AdainResBlk>> {
            Ok(vec![
                AdainResBlk::load(&st, &format!("predictor.{name}.0"), hidden, hidden, false)?,
                AdainResBlk::load(&st, &format!("predictor.{name}.1"), hidden, half, true)?,
                AdainResBlk::load(&st, &format!("predictor.{name}.2"), half, half, false)?,
            ])
        };
        let te_k = g("text_encoder_kernel_size");
        let te_cnn = (0..n_layer)
            .map(|i| {
                Ok((
                    Conv1d::load(
                        &st,
                        &format!("text_encoder.cnn.{i}.0"),
                        hidden,
                        hidden,
                        te_k,
                        1,
                        1,
                        false,
                    )?,
                    st.tensor_f32(&format!("text_encoder.cnn.{i}.1.gamma"))?,
                    st.tensor_f32(&format!("text_encoder.cnn.{i}.1.beta"))?,
                ))
            })
            .collect::<Result<Vec<_>>>()?;
        let ist = v.get("istftnet").context("istftnet cfg")?;
        let ig = |k: &str| ist.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        let iarr = |k: &str| -> Vec<usize> {
            ist.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 rates = iarr("upsample_rates");
        let up_ks = iarr("upsample_kernel_sizes");
        let res_ks = iarr("resblock_kernel_sizes");
        let dils: Vec<Vec<usize>> = ist
            .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 ch0 = ig("upsample_initial_channel");
        let n_fft = ig("gen_istft_n_fft");
        let hop = ig("gen_istft_hop_size");
        let mut ups = Vec::new();
        let mut noise_convs = Vec::new();
        let mut noise_res = Vec::new();
        let mut resblocks = Vec::new();
        for (i, (&u, &k)) in rates.iter().zip(&up_ks).enumerate() {
            let c_cur = ch0 >> (i + 1);
            ups.push(ConvT1d::load(
                &st,
                &format!("decoder.generator.ups.{i}"),
                ch0 >> i,
                c_cur,
                k,
                u,
                (k - u) / 2,
                0,
                false,
            )?);
            for (j, (&rk, dil)) in res_ks.iter().zip(&dils).enumerate() {
                resblocks.push(SnakeResBlock::load(
                    &st,
                    &format!("decoder.generator.resblocks.{}", i * res_ks.len() + j),
                    c_cur,
                    rk,
                    dil,
                )?);
            }
            if i + 1 < rates.len() {
                let stride_f0: usize = rates[i + 1..].iter().product();
                let mut nc = Conv1d::load(
                    &st,
                    &format!("decoder.generator.noise_convs.{i}"),
                    c_cur,
                    n_fft + 2,
                    stride_f0 * 2,
                    stride_f0,
                    1,
                    false,
                )?;
                nc.pad = stride_f0.div_ceil(2);
                noise_convs.push(nc);
                noise_res.push(SnakeResBlock::load(
                    &st,
                    &format!("decoder.generator.noise_res.{i}"),
                    c_cur,
                    7,
                    &[1, 3, 5],
                )?);
            } else {
                noise_convs.push(Conv1d::load(
                    &st,
                    &format!("decoder.generator.noise_convs.{i}"),
                    c_cur,
                    n_fft + 2,
                    1,
                    1,
                    1,
                    false,
                )?);
                noise_res.push(SnakeResBlock::load(
                    &st,
                    &format!("decoder.generator.noise_res.{i}"),
                    c_cur,
                    11,
                    &[1, 3, 5],
                )?);
            }
        }
        let last_ch = ch0 >> rates.len();
        Ok(Self {
            vocab,
            word_emb: st.tensor_f32("bert.embeddings.word_embeddings.weight")?,
            pos_emb: st.tensor_f32("bert.embeddings.position_embeddings.weight")?,
            type_emb: st.tensor_f32("bert.embeddings.token_type_embeddings.weight")?,
            emb_ln_w: st.tensor_f32("bert.embeddings.LayerNorm.weight")?,
            emb_ln_b: st.tensor_f32("bert.embeddings.LayerNorm.bias")?,
            emb_map: Linear::load(&st, "bert.encoder.embedding_hidden_mapping_in", bh, emb_dim)?,
            albert,
            bert_layers: pg("num_hidden_layers"),
            bert_hidden: bh,
            bert_heads: pg("num_attention_heads"),
            emb_dim,
            bert_encoder: Linear::load(&st, "bert_encoder", hidden, bh)?,
            dur_enc,
            lstm: BiLstm::load(&st, "predictor.lstm", hidden + STYLE_DIM, hidden / 2)?,
            duration_proj: Linear::load(
                &st,
                "predictor.duration_proj.linear_layer",
                g("max_dur"),
                hidden,
            )?,
            shared: BiLstm::load(&st, "predictor.shared", hidden + STYLE_DIM, hidden / 2)?,
            f0_blocks: load_chain("F0")?,
            n_blocks: load_chain("N")?,
            f0_proj: Conv1d::load(&st, "predictor.F0_proj", 1, half, 1, 1, 1, false)?,
            n_proj: Conv1d::load(&st, "predictor.N_proj", 1, half, 1, 1, 1, false)?,
            te_emb: st.tensor_f32("text_encoder.embedding.weight")?,
            te_cnn,
            te_lstm: BiLstm::load(&st, "text_encoder.lstm", hidden, hidden / 2)?,
            f0_conv: Conv1d::load(&st, "decoder.F0_conv", 1, 1, 3, 2, 1, false)?,
            n_conv: Conv1d::load(&st, "decoder.N_conv", 1, 1, 3, 2, 1, false)?,
            encode: AdainResBlk::load(&st, "decoder.encode", hidden + 2, 1024, false)?,
            asr_res: Conv1d::load(&st, "decoder.asr_res.0", 64, hidden, 1, 1, 1, false)?,
            decode: vec![
                AdainResBlk::load(&st, "decoder.decode.0", 1024 + 2 + 64, 1024, false)?,
                AdainResBlk::load(&st, "decoder.decode.1", 1024 + 2 + 64, 1024, false)?,
                AdainResBlk::load(&st, "decoder.decode.2", 1024 + 2 + 64, 1024, false)?,
                AdainResBlk::load(&st, "decoder.decode.3", 1024 + 2 + 64, hidden, true)?,
            ],
            ups,
            noise_convs,
            noise_res,
            resblocks,
            conv_post: Conv1d::load(
                &st,
                "decoder.generator.conv_post",
                n_fft + 2,
                last_ch,
                7,
                1,
                1,
                false,
            )?,
            m_source_linear: Linear::load(&st, "decoder.generator.m_source.l_linear", 1, 9)?,
            unvoiced_phase: st.tensor_f32("kokoro_unvoiced_phase").unwrap_or_default(),
            n_fft,
            hop,
            num_kernels: res_ks.len(),
            upsample_total: rates.iter().product::<usize>() * hop,
            hidden,
        })
    }

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

    /// Phoneme string → `[0, ids…, 0]` (unknown phonemes dropped, exactly like KModel).
    pub fn phonemes_to_ids(&self, phonemes: &str) -> Vec<u32> {
        let mut ids = vec![0u32];
        ids.extend(
            phonemes
                .chars()
                .filter_map(|c| self.vocab.get(c.to_string().as_str()).copied()),
        );
        ids.push(0);
        ids
    }

    fn bert_ln(x: &mut [f32], d: usize, w: &[f32], b: &[f32], eps: f32) {
        for row in x.chunks_mut(d) {
            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 + eps).sqrt();
            for (i, v) in row.iter_mut().enumerate() {
                *v = (*v - mean) * inv * w[i] + b[i];
            }
        }
    }

    /// PLBERT: token ids → last hidden `[t, bert_hidden]` (one shared ALBERT layer, applied N×).
    pub fn bert_forward(&self, ids: &[u32]) -> Vec<f32> {
        let (e, bh, t) = (self.emb_dim, self.bert_hidden, ids.len());
        let mut x = vec![0f32; t * e];
        for (i, &id) in ids.iter().enumerate() {
            for c in 0..e {
                x[i * e + c] =
                    self.word_emb[id as usize * e + c] + self.pos_emb[i * e + c] + self.type_emb[c];
            }
        }
        Self::bert_ln(&mut x, e, &self.emb_ln_w, &self.emb_ln_b, 1e-12);
        let mut x = self.emb_map.forward(&x, t);
        let (heads, hd) = (self.bert_heads, bh / self.bert_heads);
        let scale = 1.0 / (hd as f32).sqrt();
        for _ in 0..self.bert_layers {
            let l = &self.albert;
            let q = l.q.forward(&x, t);
            let k = l.k.forward(&x, t);
            let v = l.v.forward(&x, t);
            let mut ctx = vec![0f32; t * bh];
            ctx.par_chunks_mut(bh).enumerate().for_each(|(i, orow)| {
                let mut srow = vec![0f32; t];
                for head in 0..heads {
                    let qi = &q[i * bh + head * hd..][..hd];
                    for j in 0..t {
                        let kj = &k[j * bh + head * hd..][..hd];
                        srow[j] = qi.iter().zip(kj).map(|(a, b)| a * b).sum::<f32>() * scale;
                    }
                    let max = srow.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                    let mut sum = 0.0;
                    for sv in srow.iter_mut() {
                        *sv = (*sv - max).exp();
                        sum += *sv;
                    }
                    let inv = 1.0 / sum;
                    let orow = &mut orow[head * hd..][..hd];
                    for j in 0..t {
                        let w = srow[j] * inv;
                        let vj = &v[j * bh + head * hd..][..hd];
                        for c in 0..hd {
                            orow[c] += w * vj[c];
                        }
                    }
                }
            });
            let proj = l.dense.forward(&ctx, t);
            let mut a = x
                .iter()
                .zip(proj)
                .map(|(xi, pi)| xi + pi)
                .collect::<Vec<f32>>();
            Self::bert_ln(&mut a, bh, &l.attn_ln_w, &l.attn_ln_b, 1e-12);
            let mut mid = l.ffn.forward(&a, t);
            for m in mid.iter_mut() {
                *m = gelu_new(*m);
            }
            let f = l.ffn_out.forward(&mid, t);
            let mut o = a
                .iter()
                .zip(f)
                .map(|(ai, fi)| ai + fi)
                .collect::<Vec<f32>>();
            Self::bert_ln(&mut o, bh, &l.full_ln_w, &l.full_ln_b, 1e-12);
            x = o;
        }
        x
    }

    /// The duration encoder: bert-projected features `[t, hidden]` + style → `[t, hidden+style]`.
    fn duration_encode(&self, d_en: &[f32], t: usize, s: &[f32]) -> Vec<f32> {
        let h = self.hidden;
        // x: time-major [t, h+style]
        let cat = |x: &[f32], per: usize| -> Vec<f32> {
            let mut out = vec![0f32; t * (per + STYLE_DIM)];
            for i in 0..t {
                out[i * (per + STYLE_DIM)..][..per].copy_from_slice(&x[i * per..][..per]);
                out[i * (per + STYLE_DIM) + per..][..STYLE_DIM].copy_from_slice(s);
            }
            out
        };
        let mut x = cat(d_en, h);
        for (lstm, fc) in self.dur_enc.lstms.iter().zip(&self.dur_enc.ada_fc) {
            let y = lstm.forward(&x, t); // [t, h]
            // AdaLayerNorm: plain LN over channels (no learned norm params) + style affine
            let g = fc.forward(s, 1); // [2h]: gamma, beta
            let mut y = y;
            for row in y.chunks_mut(h) {
                let mean = row.iter().sum::<f32>() / h as f32;
                let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / h as f32;
                let inv = 1.0 / (var + 1e-5).sqrt();
                for (c, v) in row.iter_mut().enumerate() {
                    *v = (1.0 + g[c]) * ((*v - mean) * inv) + g[h + c];
                }
            }
            x = cat(&y, h);
        }
        x
    }

    /// Everything the prosody predictor yields. Tensors match the oracle layouts.
    #[allow(clippy::type_complexity)]
    pub fn predict_prosody(
        &self,
        bert_hidden: &[f32],
        ref_s: &[f32],
        speed: f32,
    ) -> (Vec<f32>, Vec<f32>, Vec<usize>, Vec<f32>, Vec<f32>, Vec<f32>) {
        let h = self.hidden;
        let t = bert_hidden.len() / self.bert_hidden;
        let s = &ref_s[STYLE_DIM..]; // prosody style = second half
        let d_en = self.bert_encoder.forward(bert_hidden, t); // [t, h]
        let d = self.duration_encode(&d_en, t, s); // [t, h+style]
        let x = self.lstm.forward(&d, t); // [t, h]
        let dur = self.duration_proj.forward(&x, t); // [t, max_dur]
        let md = self.duration_proj.n;
        let duration: Vec<f32> = (0..t)
            .map(|i| dur[i * md..][..md].iter().map(|&l| sigmoid(l)).sum::<f32>() / speed)
            .collect();
        let pred_dur: Vec<usize> = duration
            .iter()
            .map(|&v| (v.round() as i64).max(1) as usize)
            .collect();
        // en = dᵀ @ align: repeat token column i pred_dur[i] times → [h+style, frames]
        let frames: usize = pred_dur.iter().sum();
        let dc = h + STYLE_DIM;
        let mut en = vec![0f32; dc * frames];
        let mut fi = 0;
        for (i, &pd) in pred_dur.iter().enumerate() {
            for _ in 0..pd {
                for c in 0..dc {
                    en[c * frames + fi] = d[i * dc + c];
                }
                fi += 1;
            }
        }
        // F0Ntrain: shared LSTM over enᵀ, then the two AdaIN chains (2× upsample inside)
        let mut en_tm = vec![0f32; frames * dc];
        for i in 0..frames {
            for c in 0..dc {
                en_tm[i * dc + c] = en[c * frames + i];
            }
        }
        let xs = self.shared.forward(&en_tm, frames); // [frames, h]
        let mut x_cm = vec![0f32; h * frames];
        for i in 0..frames {
            for c in 0..h {
                x_cm[c * frames + i] = xs[i * h + c];
            }
        }
        let run_chain = |blocks: &[AdainResBlk], proj: &Conv1d| -> Vec<f32> {
            let mut cur = x_cm.clone();
            let mut ct = frames;
            for b in blocks {
                let (nx, nt) = b.forward(&cur, ct, s);
                cur = nx;
                ct = nt;
            }
            proj.forward(&cur, ct) // [1, 2·frames]
        };
        let (f0, n) = rayon::join(
            || run_chain(&self.f0_blocks, &self.f0_proj),
            || run_chain(&self.n_blocks, &self.n_proj),
        );
        (d, duration, pred_dur, en, f0, n)
    }

    /// Text encoder: ids → `[hidden, t]` channel-major (CNN stack + BiLSTM).
    pub fn text_encode(&self, ids: &[u32]) -> Vec<f32> {
        let (h, t) = (self.hidden, ids.len());
        let mut x = vec![0f32; h * t];
        for (i, &id) in ids.iter().enumerate() {
            for c in 0..h {
                x[c * t + i] = self.te_emb[id as usize * h + c];
            }
        }
        for (conv, gamma, beta) in &self.te_cnn {
            let y = conv.forward(&x, t);
            // LayerNorm over channels at each t, then LeakyReLU(0.2)
            let mut y2 = vec![0f32; h * t];
            for ti in 0..t {
                let mut mean = 0f32;
                for c in 0..h {
                    mean += y[c * t + ti];
                }
                mean /= h as f32;
                let mut var = 0f32;
                for c in 0..h {
                    let dv = y[c * t + ti] - mean;
                    var += dv * dv;
                }
                var /= h as f32;
                let inv = 1.0 / (var + 1e-5).sqrt();
                for c in 0..h {
                    let v = (y[c * t + ti] - mean) * inv * gamma[c] + beta[c];
                    y2[c * t + ti] = if v < 0.0 { v * 0.2 } else { v };
                }
            }
            x = y2;
        }
        let mut tm = vec![0f32; t * h];
        for i in 0..t {
            for c in 0..h {
                tm[i * h + c] = x[c * t + i];
            }
        }
        let y = self.te_lstm.forward(&tm, t); // [t, h]
        let mut out = vec![0f32; h * t];
        for i in 0..t {
            for c in 0..h {
                out[c * t + i] = y[i * h + c];
            }
        }
        out
    }

    /// Expand `[ch, t]` by repeating column i `dur[i]` times → `[ch, Σdur]`.
    pub fn expand(&self, x: &[f32], t: usize, dur: &[usize]) -> Vec<f32> {
        let ch = x.len() / t;
        let frames: usize = dur.iter().sum();
        let mut out = vec![0f32; ch * frames];
        let mut fi = 0;
        for (i, &pd) in dur.iter().enumerate() {
            for _ in 0..pd {
                for c in 0..ch {
                    out[c * frames + fi] = x[c * t + i];
                }
                fi += 1;
            }
        }
        out
    }

    /// Decoder front: asr `[hidden, frames]`, F0/N `[2·frames]`, decoder style → the generator
    /// input `[hidden, 2·frames]`.
    pub fn decoder_pre(
        &self,
        asr: &[f32],
        f0_curve: &[f32],
        n_curve: &[f32],
        s: &[f32],
    ) -> Vec<f32> {
        let frames = f0_curve.len() / 2;
        let f0 = self.f0_conv.forward(f0_curve, f0_curve.len()); // [1, frames]
        let n = self.n_conv.forward(n_curve, n_curve.len());
        let cat3 = |x: &[f32], xc: usize| -> Vec<f32> {
            let mut out = Vec::with_capacity((xc + 2) * frames);
            out.extend_from_slice(x);
            out.extend_from_slice(&f0);
            out.extend_from_slice(&n);
            out
        };
        let x = cat3(asr, self.hidden);
        let (mut x, _) = self.encode.forward(&x, frames, s);
        let asr_res = self.asr_res.forward(asr, frames); // [64, frames]
        let mut res = true;
        let mut ct = frames;
        for block in &self.decode {
            if res {
                let mut cat = Vec::with_capacity((1024 + 64 + 2) * ct);
                cat.extend_from_slice(&x);
                cat.extend_from_slice(&asr_res);
                cat.extend_from_slice(&f0);
                cat.extend_from_slice(&n);
                x = cat;
            }
            let (nx, nt) = block.forward(&x, ct, s);
            x = nx;
            if nt != ct {
                res = false;
                ct = nt;
            }
        }
        x // [hidden, 2·frames]
    }

    /// The raw merged sine-harmonic source at audio rate (deterministic: zero initial phases,
    /// no additive noise).
    #[doc(hidden)]
    pub fn sine_source(&self, f0_curve: &[f32]) -> Vec<f32> {
        let up = self.upsample_total; // e.g. 10·6·5 = 300
        let t_in = f0_curve.len();
        let l = t_in * up; // upsampled length (nearest)
        // 9 harmonics: rad_h[t] = (f0[t]·(h+1)/SR) mod 1, on the ×up grid (nearest upsample)
        let nh = 9usize;
        // The reference upsamples f0 (nearest), rad = (f0·h/SR) mod 1, downsamples rad ×1/up
        // (linear), cumsums, ×2π, ×up, then upsamples ×up (linear) and takes sin. The phase reaches
        // ~10⁴ rad, where f32 has < 3 fractional digits — so torch's own phase carries ~0.04 rad of
        // f32 rounding that its (accurate) sin faithfully reproduces, and no reimplementation of the
        // f32 accumulation order matches it bit-for-bit. We instead compute the phase in f64 and
        // reduce mod 2π before sin: the mathematically intended result. This diverges from torch's
        // f32-noisy source by ~0.04 in the raw harmonic (an additive secondary path in the decoder;
        // the main signal path stays bit-exact), which is perceptually inaudible.
        let mut sines = vec![0f32; l * nh]; // [l, nh] time-major
        let t_dn = l / up; // == t_in
        let two_pi = 2.0 * std::f64::consts::PI;
        for hi in 0..nh {
            // rad_dn[i]: rad sampled at block-center i (the ×1/up linear downsample lands both taps
            // inside constant block i, so it reduces to the block value).
            let mut phase_dn = vec![0f64; t_dn];
            let mut acc = 0f64;
            for (i, p) in phase_dn.iter_mut().enumerate() {
                let v = f0_curve[i.min(t_in - 1)] as f64 * (hi as f64 + 1.0) / SR as f64;
                acc += v - v.floor(); // cumsum of rad
                *p = acc * two_pi * up as f64;
            }
            for t in 0..l {
                let src = ((t as f64 + 0.5) / up as f64 - 0.5)
                    .max(0.0)
                    .min(t_dn as f64 - 1.0);
                let i0 = (src as usize).min(t_dn - 1);
                let i1 = (i0 + 1).min(t_dn - 1);
                let w = src - i0 as f64;
                let ph = phase_dn[i0] * (1.0 - w) + phase_dn[i1] * w;
                sines[t * nh + hi] = (ph.rem_euclid(two_pi)).sin() as f32;
            }
        }
        // sine_waves·amp·uv → merge linear → tanh
        let mut har = vec![0f32; l];
        har.par_iter_mut().enumerate().for_each(|(t, o)| {
            let uv = if f0_curve[(t / up).min(t_in - 1)] > 10.0 {
                1.0f32
            } else {
                0.0
            };
            let mut acc = self.m_source_linear.b[0];
            for hi in 0..nh {
                acc += sines[t * nh + hi] * 0.1 * uv * self.m_source_linear.w[hi];
            }
            *o = acc.tanh();
        });
        har
    }

    /// STFT (center=true, reflect pad, periodic hann) of the source → spec ‖ phase
    /// `[n_fft+2, l/hop + 1]`.
    #[doc(hidden)]
    pub fn stft_spec_phase(&self, har: &[f32]) -> (Vec<f32>, usize) {
        let l = har.len();
        let (n_fft, hop) = (self.n_fft, self.hop);
        let nb = n_fft / 2 + 1;
        let frames = l / hop + 1;
        let win: Vec<f32> = (0..n_fft)
            .map(|i| 0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / n_fft as f32).cos()))
            .collect();
        let padded = |i: i64| -> f32 {
            // reflect padding by n_fft/2 on both sides
            let li = l as i64;
            let j = if i < 0 {
                -i
            } else if i >= li {
                2 * li - 2 - i
            } else {
                i
            };
            har[j as usize]
        };
        let mut out = vec![0f32; (n_fft + 2) * frames];
        let (spec_rows, phase_rows) = out.split_at_mut(nb * frames);
        spec_rows
            .par_chunks_mut(frames)
            .zip(phase_rows.par_chunks_mut(frames))
            .enumerate()
            .for_each(|(b, (srow, prow))| {
                for f in 0..frames {
                    let start = f as i64 * hop as i64 - (n_fft / 2) as i64;
                    let (mut re, mut im) = (0f64, 0f64);
                    for (j, &wj) in win.iter().enumerate() {
                        let v = (padded(start + j as i64) * wj) as f64;
                        let ang = 2.0 * std::f64::consts::PI * b as f64 * j as f64 / n_fft as f64;
                        re += v * ang.cos();
                        im -= v * ang.sin();
                    }
                    srow[f] = ((re * re + im * im).sqrt()) as f32;
                    prow[f] = im.atan2(re) as f32;
                }
            });
        // Fully-constant (unvoiced) frames: X[b] = c·FFT(win)[b], so every AC bin has magnitude 0
        // in exact math and its angle is pure FFT rounding noise. torch produces one deterministic
        // column for all such frames; replay the reference's (see export_kokoro.py). Detect by an
        // all-samples-equal frame — the source is locally constant only in unvoiced regions (f0 ≤
        // threshold), all at the same value c = tanh(m_source bias), so one dumped column suffices.
        // (A frame-internal test avoids depending on c's exact ULP, which differs between torch's
        // tanh(l_linear(0)) and Rust's tanh(bias).)
        if !self.unvoiced_phase.is_empty() {
            for f in 0..frames {
                let start = f as i64 * hop as i64 - (n_fft / 2) as i64;
                let c0 = padded(start);
                if c0 != 0.0 && (1..n_fft).all(|j| padded(start + j as i64) == c0) {
                    for b in 0..nb {
                        phase_rows[b * frames + f] = self.unvoiced_phase[b];
                    }
                }
            }
        }
        (out, frames)
    }

    /// ISTFTNet generator: `x` `[hidden, 2·frames]`, decoder style, F0 curve → waveform.
    pub fn generator_forward(&self, x: &[f32], s: &[f32], f0_curve: &[f32]) -> Result<Vec<f32>> {
        Ok(self.generator_stages(x, s, f0_curve, None)?.4)
    }

    /// Generator with every intermediate exposed (debug/parity surface):
    /// (harmonic spec‖phase, per-upsample-stage outputs, spec, phase, audio). `har_override`
    /// injects a precomputed source STFT (to isolate the generator from the harmonic source).
    #[doc(hidden)]
    #[allow(clippy::type_complexity)]
    pub fn generator_stages(
        &self,
        x: &[f32],
        s: &[f32],
        f0_curve: &[f32],
        har_override: Option<&[f32]>,
    ) -> Result<(Vec<f32>, Vec<Vec<f32>>, Vec<f32>, Vec<f32>, Vec<f32>)> {
        let (har, har_frames) = match har_override {
            Some(h) => (h.to_vec(), h.len() / (self.n_fft + 2)),
            None => self.stft_spec_phase(&self.sine_source(f0_curve)),
        };
        let mut stages = Vec::new();
        let mut x = x.to_vec();
        let mut t = f0_curve.len();
        for i in 0..self.ups.len() {
            lrelu(&mut x, 0.1);
            let xs = self.noise_convs[i].forward(&har, har_frames);
            let xs_t = self.noise_convs[i].out_len(har_frames);
            let xs = self.noise_res[i].forward(&xs, xs_t, s);
            let (mut xu, mut tu) = self.ups[i].forward(&x, t);
            if i == self.ups.len() - 1 {
                // reflection pad (1, 0): prepend column 1
                let ch = self.ups[i].out_ch;
                let mut padded = vec![0f32; ch * (tu + 1)];
                for c in 0..ch {
                    padded[c * (tu + 1)] = xu[c * tu + 1];
                    padded[c * (tu + 1) + 1..][..tu].copy_from_slice(&xu[c * tu..][..tu]);
                }
                xu = padded;
                tu += 1;
            }
            anyhow::ensure!(xs_t == tu, "noise/up length mismatch {xs_t} vs {tu}");
            for (a, b) in xu.iter_mut().zip(&xs) {
                *a += b;
            }
            let ch = self.ups[i].out_ch;
            let mut acc = self.resblocks[i * self.num_kernels].forward(&xu, tu, s);
            for j in 1..self.num_kernels {
                let r = self.resblocks[i * self.num_kernels + j].forward(&xu, tu, s);
                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;
            }
            debug_assert_eq!(acc.len(), ch * tu);
            x = acc;
            t = tu;
            stages.push(x.clone());
        }
        lrelu(&mut x, 0.01); // torch default slope — HF/kokoro call without an argument
        let y = self.conv_post.forward(&x, t); // [n_fft+2, t]
        let nb = self.n_fft / 2 + 1;
        // spec = exp, phase = sin → iSTFT
        let spec: Vec<f32> = y[..nb * t].iter().map(|v| v.exp()).collect();
        let phase: Vec<f32> = y[nb * t..].iter().map(|v| v.sin()).collect();
        let audio = self.istft(&spec, &phase, t);
        Ok((har, stages, spec, phase, audio))
    }

    /// torch.istft (center=true): complex spec+phase `[n_fft/2+1, frames]` → time signal.
    fn istft(&self, spec: &[f32], phase: &[f32], frames: usize) -> Vec<f32> {
        let (n_fft, hop) = (self.n_fft, self.hop);
        let nb = n_fft / 2 + 1;
        let win: Vec<f64> = (0..n_fft)
            .map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / n_fft as f64).cos()))
            .collect();
        let full = (frames - 1) * hop + n_fft;
        let mut y = vec![0f64; full];
        let mut norm = vec![0f64; full];
        let inv_n = 1.0 / n_fft as f64;
        for f in 0..frames {
            // irfft of the frame
            for j in 0..n_fft {
                let mut acc = 0f64;
                for b in 0..nb {
                    let mag = spec[b * frames + f] as f64;
                    let ph = phase[b * frames + f] as f64;
                    let (re, im) = (mag * ph.cos(), mag * ph.sin());
                    let ang = 2.0 * std::f64::consts::PI * b as f64 * j as f64 / n_fft as f64;
                    let term = re * ang.cos() - im * ang.sin();
                    acc += if b == 0 || b == nb - 1 {
                        term
                    } else {
                        2.0 * term
                    };
                }
                let v = acc * inv_n * win[j];
                y[f * hop + j] += v;
                norm[f * hop + j] += win[j] * win[j];
            }
        }
        let start = n_fft / 2;
        let len = (frames - 1) * hop;
        (start..start + len)
            .map(|i| (y[i] / norm[i].max(1e-11)) as f32)
            .collect()
    }

    /// End-to-end: IPA phonemes + 256-d voice row (`ref_s`) → 24 kHz waveform. Deterministic.
    pub fn synthesize(&self, phonemes: &str, ref_s: &[f32], speed: f32) -> Result<Vec<f32>> {
        anyhow::ensure!(ref_s.len() == 2 * STYLE_DIM, "ref_s must be 256-d");
        let ids = self.phonemes_to_ids(phonemes);
        let bert_hidden = self.bert_forward(&ids);
        let (_, _, pred_dur, en, f0, n) = self.predict_prosody(&bert_hidden, ref_s, speed);
        let _ = en;
        let t_en = self.text_encode(&ids);
        let asr = self.expand(&t_en, ids.len(), &pred_dur);
        let s_dec = &ref_s[..STYLE_DIM];
        let gen_in = self.decoder_pre(&asr, &f0, &n, s_dec);
        self.generator_forward(&gen_in, s_dec, &f0)
    }
}