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<f32> {
+    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");
+    }
 }

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::<Vec<f32>>();
+    let i64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
+    let i32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::<Vec<i32>>();
+    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::<f64>() / 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})");
+}