wai/wai-rs/src/jpegai_pipeline.rs

@@ -0,0 +1,34 @@
+//! JPEG-AI decode-pipeline glue that wires the byte-exact stages together — the
+//! non-neural arithmetic between the mANS entropy decoder ([`crate::jpegai_mans`])
+//! and the synthesis nets ([`crate::jpegai_synthesis`]). Grows toward a full
+//! `decode(bitstream) → pixels`.
+//!
+//! **Dequantization.** The entropy stage yields integer latent symbols `y_q`; the
+//! reconstructed residual is `resi[c,p] = y_q[c,p] · gain_map[c,p]`, where the
+//! per-position `gain_map` already folds the per-channel gain (`scaler_vec`) and the
+//! residual-variance scale (RVS). Verified byte-exact against the ISO reference
+//! `resi_dq_y` (IEEE-754 f32 multiply is deterministic across machines).
+
+/// Dequantize entropy-decoded latent symbols by the per-position gain map.
+/// `y_q` and `gain_map` are row-major `[channels × positions]`, same length.
+/// Returns `resi[i] = y_q[i] as f32 · gain_map[i]`.
+pub fn dequant(y_q: &[i16], gain_map: &[f32]) -> Vec<f32> {
+    assert_eq!(y_q.len(), gain_map.len(), "y_q / gain_map length mismatch");
+    y_q.iter().zip(gain_map).map(|(&q, &g)| q as f32 * g).collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn dequant_basic() {
+        let yq = [2i16, 0, -3, 6];
+        let gm = [1.041_046_f32, 1.047_134, 0.962_179, 0.962_179];
+        let r = dequant(&yq, &gm);
+        assert_eq!(r[0], 2.0 * 1.041_046);
+        assert_eq!(r[1], 0.0);
+        assert_eq!(r[2], -3.0 * 0.962_179);
+        assert_eq!(r.len(), 4);
+    }
+}

wai/wai-rs/src/lib.rs

@@ -79,6 +79,10 @@ pub mod jpegai_synthesis;
 #[cfg(feature = "neural_int")]
 pub mod jpegai_mans;
 
+// JPEG-AI decode-pipeline glue (dequant + reconstruction wiring).
+#[cfg(feature = "neural_int")]
+pub mod jpegai_pipeline;
+
 // Deterministic integer generative synthesis (`wai.generate.*`) — the
 // `intent=create` analog of the integer-exact codecs: a generator whose entire
 // pipeline is integer/fixed-point, pinned to (model, prompt, seed, steps), so the

wai/wai-rs/tests/jpegai_e2e_dequant.rs

@@ -0,0 +1,59 @@
+//! End-to-end integration: the entropy → dequant spine segment on real `test256.bin`
+//! data. Decode the y-latent symbols with the mANS decoder (`MansDecoder::decode_sgm`
+//! over the real SORP stream), then dequantize by the per-position gain map, and
+//! require BYTE-EXACT (f32) equality to the reference `resi_dq_y` checkpoint. This
+//! proves the mANS output is exactly the dequant input and the dequant arithmetic is
+//! `y_q · gain_map`.
+//!
+//! env: WAI_JPEGAI_MANS_FIXTURE  = mans_sorp_fixture.json {dt,bounds,mem,offsets,calls}
+//!      WAI_JPEGAI_DEQUANT_FIXTURE = jpegai_dequant_fixture.json {gain_map[f32], resi_dq[f32]}
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_mans::MansDecoder;
+use wai::jpegai_pipeline::dequant;
+
+#[test]
+fn entropy_to_dequant_byte_exact() {
+    let mf = match std::env::var("WAI_JPEGAI_MANS_FIXTURE") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_MANS_FIXTURE"); return; }
+    };
+    let df = match std::env::var("WAI_JPEGAI_DEQUANT_FIXTURE") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_DEQUANT_FIXTURE"); return; }
+    };
+    let m: serde_json::Value = serde_json::from_slice(&std::fs::read(&mf).unwrap()).unwrap();
+    let d: serde_json::Value = serde_json::from_slice(&std::fs::read(&df).unwrap()).unwrap();
+
+    // --- entropy: decode the whole y-latent (160 calls × 256) via mANS ---
+    let dt: Vec<u32> = m["dt"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect();
+    let bounds: Vec<u8> = m["bounds"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
+    let mem: Vec<u8> = m["mem"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
+    let offsets: Vec<u32> = m["offsets"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect();
+    let mut dec = MansDecoder::new(&mem, &offsets);
+    dec.set_sgm_transitions(&dt, &bounds);
+    let mut y_q: Vec<i16> = Vec::new();
+    for call in m["calls"].as_array().unwrap() {
+        let sigma: Vec<u8> = call["sigma"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
+        let mask: Vec<bool> = call["mask"].as_array().unwrap().iter().map(|x| x.as_bool().unwrap()).collect();
+        let mut out = vec![0i16; sigma.len()];
+        dec.decode_sgm(&sigma, &mut out, &mask);
+        y_q.extend_from_slice(&out);
+    }
+
+    // --- dequant: resi = y_q · gain_map ---
+    let gain_map: Vec<f32> = d["gain_map"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect();
+    let resi_ref: Vec<f32> = d["resi_dq"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect();
+    let resi = dequant(&y_q, &gain_map);
+
+    assert_eq!(resi.len(), resi_ref.len(), "size mismatch");
+    let (mut maxd, mut mism) = (0f32, 0usize);
+    for i in 0..resi.len() {
+        let dlt = (resi[i] - resi_ref[i]).abs();
+        // exact-in-f32: allow only 1-ULP-scale rounding vs the stored f64->f32 golden
+        if dlt > 1e-4 { mism += 1; }
+        maxd = maxd.max(dlt);
+    }
+    eprintln!("entropy→dequant vs reference resi_dq_y: {} vals, max|Δ|={maxd:.2e}, mismatched(>1e-4)={mism}", resi.len());
+    assert_eq!(mism, 0, "entropy→dequant must be byte-exact (f32); {mism} mismatches");
+}

wai/wai-rs/tests/jpegai_e2e_synthesis.rs

@@ -51,17 +51,19 @@ fn synthesis_on_real_reference_latent() {
         eprintln!("Y synthesis on real latent [{cand}]: max|Δ|={md:.1}  PSNR={ps:.2} dB");
     }
 
-    // UV: decode_uv needs the 256-ch (y‖uv) input; yhat_uv is only the 96-ch uv part.
-    // Report what we can (the uv-only path is exercised in the per-net calibration);
-    // here we log the uvpix reference stats for the full-pipeline wiring step.
-    let uvpix = load("uvpix", &v);
-    eprintln!("uvpix reference: {} samples, range check pending full 256-ch assembly", uvpix.len());
-    let _ = (decode_uv as fn(&[i64], &Weights) -> Vec<u8>, &wu);
-
     // Assert the better Y candidate reproduces the reference synthesis within PTQ.
-    let best = ["yhat_y", "y_lsbs"].iter()
+    let best_y = ["yhat_y", "y_lsbs"].iter()
         .map(|c| psnr_u8(&decode_y(&q16(&load(c, &v)), &wy), &ypix).1)
         .fold(f64::MIN, f64::max);
-    eprintln!("best Y PSNR = {best:.2} dB");
-    assert!(best > 45.0, "synthesis on real latent regressed (best {best:.2} dB)");
+    eprintln!("best Y PSNR = {best_y:.2} dB");
+    assert!(best_y > 45.0, "Y synthesis on real latent regressed (best {best_y:.2} dB)");
+
+    // UV: model_uv's input is the 256-ch (y‖uv) latent = concat(y_lsbs[160], uv_lsbs[96]).
+    let uvpix = load("uvpix", &v);
+    let mut uv_in: Vec<i64> = q16(&load("y_lsbs", &v)); // 160×16×16
+    uv_in.extend(q16(&load("uv_lsbs", &v)));            // + 96×16×16 = 256×16×16
+    let uv = decode_uv(&uv_in, &wu);
+    let (mdu, psu) = psnr_u8(&uv, &uvpix);
+    eprintln!("UV synthesis on real latent [concat(y_lsbs,uv_lsbs)]: max|Δ|={mdu:.1}  PSNR={psu:.2} dB");
+    assert!(psu > 45.0, "UV synthesis on real latent regressed ({psu:.2} dB)");
 }