diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index 3f63c9fa5..92422fe1a 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -12,6 +12,94 @@ //! **Sigma selection.** [`sigma_index`] closes the scale path: it turns the //! hyper-scale decoder's index `i_s` ([`crate::jpegai_synthesis::decode_hyper_scale`]) //! into the per-position distribution index the mANS `sgm` decoder consumes. +//! +//! **Latent reconstruction.** [`reconstruct_yhat`] is the SGMM autoregressive spine: +//! it turns the dequantized residual back into the synthesis-input latent `y_hat` by +//! running the four MCM context stages over a 2×2 checkerboard, each stage predicting +//! a mean from the hyper params plus the already-reconstructed stages. + +use crate::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights}; + +/// The 2×2 sub-lattice each MCM stage owns, in reference stage order: +/// stage 0 → (0,0), stage 1 → (1,1), stage 2 → (0,1), stage 3 → (1,0). +/// (The reference's `down_shuffle` returns `(k00, k11, k01, k10)`; `up_shuffle` is +/// its exact inverse. The diagonal pair is decoded first — a checkerboard.) +const SUBPOS: [(usize, usize); 4] = [(0, 0), (1, 1), (0, 1), (1, 0)]; + +/// Split `[c, 2h, 2w]` into the four `[c, h, w]` sub-lattices (reference `down_shuffle`). +fn down_shuffle(y: &[i64], c: usize, h2: usize, w2: usize) -> [Vec; 4] { + let (h, w) = (h2 / 2, w2 / 2); + let mut parts = std::array::from_fn(|_| vec![0i64; c * h * w]); + for (pi, &(ky, kx)) in SUBPOS.iter().enumerate() { + for ch in 0..c { + for i in 0..h { + for j in 0..w { + parts[pi][(ch * h + i) * w + j] = y[(ch * h2 + 2 * i + ky) * w2 + 2 * j + kx]; + } + } + } + } + parts +} + +/// Interleave the four `[c, h, w]` sub-lattices back into `[c, 2h, 2w]` (`up_shuffle`). +fn up_shuffle(parts: &[Vec; 4], c: usize, h: usize, w: usize) -> Vec { + let (h2, w2) = (h * 2, w * 2); + let mut out = vec![0i64; c * h2 * w2]; + for (pi, &(ky, kx)) in SUBPOS.iter().enumerate() { + for ch in 0..c { + for i in 0..h { + for j in 0..w { + out[(ch * h2 + 2 * i + ky) * w2 + 2 * j + kx] = parts[pi][(ch * h + i) * w + j]; + } + } + } + } + out +} + +/// **SGMM autoregressive reconstruction** — rebuild the synthesis-input latent `y_hat` +/// from the dequantized residual and the hyper params. +/// +/// `resi` is `[c, 2h, 2w]` (Q16); `psi` is `[4c, h, w]` (Q16, the hyper decoder's +/// output); `mcm` are the four MCM stages' own weights. Returns `y_hat` as `[c, 2h, 2w]`. +/// +/// ```text +/// p = chunk(psi, 4) // per-stage hyper params, ascending +/// r = down_shuffle(resi) // 2×2 checkerboard sub-lattices +/// for i in 0..4: +/// mu_i = MCM_i(p_i, cat(y_0 .. y_{i-1})) // ctx grows: 0 / c / 2c / 3c channels +/// y_i = r_i + mu_i +/// y_hat = up_shuffle(y_0 .. y_3) +/// ``` +/// +/// Stage 0 has no context (the reference seeds `spatial_params` with zeros and its +/// stage-0 model ignores them). The context concatenation is **ascending** (`y_0` +/// first) — the reference does `cat(y_hat_list)`, despite the ONNX input being +/// labelled `y[1-0]`. +pub fn reconstruct_yhat(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usize, mcm: [&Weights; 4]) -> Vec { + let (h, w) = (h2 / 2, w2 / 2); + let plane = c * h * w; + assert_eq!(resi.len(), c * h2 * w2, "resi must be [c, 2h, 2w]"); + assert_eq!(psi.len(), 4 * plane, "psi must be [4c, h, w]"); + + let r = down_shuffle(resi, c, h2, w2); + let mut ctx: Vec = Vec::with_capacity(3 * plane); // accumulated y_hat, ascending + let mut parts: [Vec; 4] = std::array::from_fn(|_| Vec::new()); + + for i in 0..4 { + let p = &psi[i * plane..(i + 1) * plane]; + let mu = if i == 0 { + mcm_stage0(p, h, w, mcm[0]) + } else { + mcm_stage_n(p, &ctx, i * c, h, w, mcm[i]) + }; + let y_i: Vec = r[i].iter().zip(&mu).map(|(&a, &b)| a + b).collect(); + if i < 3 { ctx.extend_from_slice(&y_i); } + parts[i] = y_i; + } + up_shuffle(&parts, c, h, w) +} /// Dequantize entropy-decoded latent symbols by the per-position gain map. /// `y_q` and `gain_map` are row-major `[channels × positions]`, same length. @@ -65,6 +153,23 @@ mod tests { assert_eq!(r.len(), 4); } + #[test] + fn shuffle_round_trips_and_uses_checkerboard_order() { + // [1ch, 4x4] ramp + let (c, h2, w2) = (1usize, 4usize, 4usize); + let y: Vec = (0..16).collect(); + let parts = down_shuffle(&y, c, h2, w2); + // stage 0 owns (0,0) of each 2x2 block: indices 0, 2, 8, 10 + assert_eq!(parts[0], vec![0, 2, 8, 10]); + // stage 1 owns (1,1): 5, 7, 13, 15 (the diagonal partner — decoded second) + assert_eq!(parts[1], vec![5, 7, 13, 15]); + // stage 2 owns (0,1): 1, 3, 9, 11 ; stage 3 owns (1,0): 4, 6, 12, 14 + assert_eq!(parts[2], vec![1, 3, 9, 11]); + assert_eq!(parts[3], vec![4, 6, 12, 14]); + // up_shuffle is the exact inverse + assert_eq!(up_shuffle(&parts, c, h2 / 2, w2 / 2), y); + } + #[test] fn sigma_index_thresholds_and_clamp() { // 1 channel, gain 0, rvs 0 → boundaries land on multiples of 128 offset by 64. diff --git a/wai/wai-rs/tests/jpegai_reconstruct_calibration.rs b/wai/wai-rs/tests/jpegai_reconstruct_calibration.rs new file mode 100644 index 000000000..10ab75983 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_reconstruct_calibration.rs @@ -0,0 +1,63 @@ +//! Calibration for [`wai::jpegai_pipeline::reconstruct_yhat`] — the SGMM +//! autoregressive spine — on REAL `test256.bin` data: rebuild the synthesis-input +//! latent `y_hat` from the reference's dequantized residual (`resi_dq_y`) and hyper +//! params (`psi_y`), and compare against the reference `yhat_y`. +//! +//! This closes the last structural gap in the decode chain: entropy → dequant → +//! **reconstruct** → synthesis. The float reference loop reproduces `yhat_y` to +//! ~139 dB (float round-off); our Q16 integer MCM stages carry the 16-bit PTQ floor +//! of the MCM nets themselves (~61–74 dB per stage), so we assert well above that. +//! +//! env: WAI_JPEGAI_RECON_FIXTURE = recon_fixture.json {resi_dq_y, psi_y, yhat_y, c, h2, w2} +//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::reconstruct_yhat; +use wai::jpegai_synthesis::Weights; + +fn q16(f: &[f64]) -> Vec { f.iter().map(|&x| (x * 65536.0).round() as i64).collect() } + +#[test] +fn reconstruct_yhat_on_real_reference_data() { + let fx = match std::env::var("WAI_JPEGAI_RECON_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_RECON_FIXTURE"); return; } + }; + let wdir = match std::env::var("WAI_JPEGAI_MCM_WEIGHTS_DIR") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_MCM_WEIGHTS_DIR"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap(); + let f64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect::>(); + let (c, h2, w2) = ( + v["c"].as_u64().unwrap() as usize, + v["h2"].as_u64().unwrap() as usize, + v["w2"].as_u64().unwrap() as usize, + ); + + let resi = q16(&f64v("resi_dq_y")); + let psi = q16(&f64v("psi_y")); + let yhat_ref = f64v("yhat_y"); + + let load = |i: usize| -> Weights { + let raw = std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap(); + Weights::from_json(&serde_json::from_slice(&raw).unwrap()) + }; + let (w0, w1, w2w, w3) = (load(0), load(1), load(2), load(3)); + + let got = reconstruct_yhat(&resi, &psi, c, h2, w2, [&w0, &w1, &w2w, &w3]); + assert_eq!(got.len(), yhat_ref.len(), "y_hat size mismatch"); + + let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64); + for i in 0..got.len() { + let g = got[i] as f64 / 65536.0; + let r = yhat_ref[i]; + let d = (g - r).abs(); + maxd = maxd.max(d); + se += d * d; + sig += r * r; + } + let snr = 10.0 * (sig / se).log10(); + eprintln!("reconstruct_yhat on real data: max|Δ| = {maxd:.5} SNR = {snr:.2} dB"); + assert!(snr > 45.0, "SGMM reconstruction regressed (SNR {snr:.2} dB)"); +}