diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index 92422fe1a..c302a537d 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -13,6 +13,10 @@ //! hyper-scale decoder's index `i_s` ([`crate::jpegai_synthesis::decode_hyper_scale`]) //! into the per-position distribution index the mANS `sgm` decoder consumes. //! +//! **LSBS refinement.** [`lsbs_refine`] applies the least-significant-bit scale +//! correction that turns the reconstructed latent `yhat` into the synthesis input +//! `y_lsbs` — a per-8×8-block table lookup keyed on the same `i_s`. +//! //! **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 @@ -138,6 +142,66 @@ pub fn sigma_index(i_s: &[i64], gain_per_ch: &[i64], rvs: &[i64], channels: usiz out } +/// **LSBS refinement** — the least-significant-bit scale correction that turns the +/// reconstructed latent `yhat` into the synthesis-input latent `y_lsbs`: +/// +/// ```text +/// likely[c,block] = clamp( (Σ_{8×8}|i_s| + 32) >> 6, 0, 4351 ) // per 8×8 block, broadcast +/// y_lsbs[c,p] = yhat[c,p] + ( t0[likely]·(yhat−resi) + t1[likely]·resi ) / 2^13 +/// ``` +/// +/// `i_s` is the hyper-scale decoder output (the log-domain scale). `t0`/`t1` are the +/// two LSBS tables (4352 entries, indexed by `likely`). Blocks are 8×8 with the +/// reference's edge padding value 1411 for incomplete blocks; `>>6` after `+32` is +/// round-to-nearest of the block-sum (an arithmetic shift). Reproduces the reference +/// `y_lsbs` to float precision (IEEE-754 deterministic). +/// +/// `yhat`/`resi`/`i_s` are row-major `[channels × h × w]`. +pub fn lsbs_refine(yhat: &[f32], resi: &[f32], i_s: &[i64], t0: &[i32], t1: &[i32], + channels: usize, h: usize, w: usize) -> Vec { + const BLOCK: usize = 8; + const PAD_VAL: i64 = 1411; + const DENOM: f32 = (1u32 << 13) as f32; // lsbs_scale_precision = 13 + let plane = h * w; + assert_eq!(yhat.len(), channels * plane); + assert_eq!(resi.len(), channels * plane); + assert_eq!(i_s.len(), channels * plane); + let last = (t0.len() - 1) as i64; // 4351 + let mut out = vec![0f32; yhat.len()]; + for c in 0..channels { + let base = c * plane; + let mut by = 0; + while by < h { + let mut bx = 0; + while bx < w { + // round-to-nearest of the 8×8 block sum of |i_s|, edge-padded with 1411 + let mut sum = 0i64; + for dy in 0..BLOCK { + for dx in 0..BLOCK { + let (y, x) = (by + dy, bx + dx); + sum += if y < h && x < w { i_s[base + y * w + x].abs() } else { PAD_VAL }; + } + } + let likely = (((sum + 32) >> 6).clamp(0, last)) as usize; + let (a0, a1) = (t0[likely] as f32, t1[likely] as f32); + for dy in 0..BLOCK { + for dx in 0..BLOCK { + let (y, x) = (by + dy, bx + dx); + if y < h && x < w { + let i = base + y * w + x; + let means = yhat[i] - resi[i]; + out[i] = yhat[i] + (a0 * means + a1 * resi[i]) / DENOM; + } + } + } + bx += BLOCK; + } + by += BLOCK; + } + } + out +} + #[cfg(test)] mod tests { use super::*; @@ -188,4 +252,27 @@ mod tests { let s1 = sigma_index(&[100, 100], &[28], &[0, 64], 1); assert_eq!(s1, vec![1, 2]); } + + #[test] + fn lsbs_refine_block_index_and_additive() { + // one 8×8 block (h=w=8, 1 channel). Constant |i_s| = 100 → sum = 6400, + // likely = (6400 + 32) >> 6 = 100. Pick tables so the additive is exact. + let n = 64; + let yhat = vec![10.0f32; n]; + let resi = vec![2.0f32; n]; + let i_s = vec![100i64; n]; + let mut t0 = vec![0i32; 4352]; + let mut t1 = vec![0i32; 4352]; + t0[100] = 8192; // ·means/2^13 = 1·means + t1[100] = 8192; // ·resi /2^13 = 1·resi + let out = lsbs_refine(&yhat, &resi, &i_s, &t0, &t1, 1, 8, 8); + // additive = 1·(10-2) + 1·2 = 10 ; y_lsbs = 10 + 10 = 20 + assert!(out.iter().all(|&v| (v - 20.0).abs() < 1e-6), "additive/likely wrong: {}", out[0]); + + // clamp: huge |i_s| would overflow the table index → clamped to 4351 + let big = vec![100_000i64; n]; + let mut t0b = vec![0i32; 4352]; t0b[4351] = 8192; + let o2 = lsbs_refine(&yhat, &resi, &big, &t0b, &vec![0i32; 4352], 1, 8, 8); + assert!((o2[0] - (10.0 + (10.0 - 2.0))).abs() < 1e-6, "table index must clamp to 4351"); + } } diff --git a/wai/wai-rs/tests/jpegai_lsbs_calibration.rs b/wai/wai-rs/tests/jpegai_lsbs_calibration.rs new file mode 100644 index 000000000..14697136e --- /dev/null +++ b/wai/wai-rs/tests/jpegai_lsbs_calibration.rs @@ -0,0 +1,49 @@ +//! Calibration for [`wai::jpegai_pipeline::lsbs_refine`] — the least-significant-bit +//! scale refinement (`yhat → y_lsbs`, the synthesis input). Runs on the REAL +//! test256.bin checkpoint (reconstructed latent `yhat_y`, dequant residual +//! `resi_dq_y`, hyper-scale index `i_s`, the two LSBS tables) and compares against the +//! reference `y_lsbs`. The refinement is float arithmetic (IEEE-754 deterministic), +//! so we require a machine-precision match. +//! +//! env: WAI_JPEGAI_LSBS_FIXTURE = lsbs_fixture.json +//! { yhat:[f64], resi:[f64], i_s:[i64], t0:[i32], t1:[i32], y_lsbs:[f64], c, h, w } +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::lsbs_refine; + +#[test] +fn lsbs_refine_byte_exact_vs_reference() { + let path = match std::env::var("WAI_JPEGAI_LSBS_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_LSBS_FIXTURE to lsbs_fixture.json"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let f32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect::>(); + let i64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::>(); + let i32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::>(); + let (c, h, w) = (v["c"].as_u64().unwrap() as usize, v["h"].as_u64().unwrap() as usize, v["w"].as_u64().unwrap() as usize); + + let yhat = f32v("yhat"); + let resi = f32v("resi"); + let i_s = i64v("i_s"); + let t0 = i32v("t0"); + let t1 = i32v("t1"); + let reference = f32v("y_lsbs"); + + let got = lsbs_refine(&yhat, &resi, &i_s, &t0, &t1, c, h, w); + assert_eq!(got.len(), reference.len(), "y_lsbs size mismatch"); + + let mut maxd = 0f32; + let mut mse = 0f64; + for i in 0..got.len() { + let d = (got[i] - reference[i]).abs(); + maxd = maxd.max(d); + mse += (d as f64) * (d as f64); + } + mse /= got.len() as f64; + let sig: f64 = reference.iter().map(|&x| (x as f64) * (x as f64)).sum::() / got.len() as f64; + let snr = if mse > 0.0 { 10.0 * (sig / mse).log10() } else { f64::INFINITY }; + eprintln!("lsbs_refine on real data: max|Δ| = {maxd:.6} SNR = {snr:.2} dB"); + // float refinement — match to well within an LSB of the reconstructed latent + assert!(maxd < 1e-3, "LSBS refinement diverged from the reference (max|Δ| = {maxd})"); +}