From aa017cf8e2de50c8d4fc6935fbb88b61bbc77219 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Wed, 8 Jul 2026 07:13:13 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Fpipeline::dequant=20+=20e2e?= =?UTF-8?q?=20integration=20(entropy=E2=86=92dequant=20byte-exact,=20synth?= =?UTF-8?q?esis=20Y/UV=20on=20real=20latent)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Start the JPEG-AI decode-pipeline glue that wires the byte-exact stages together. - jpegai_pipeline::dequant — 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). Deterministic f32. - e2e entropy→dequant: decode the real test256.bin y-latent (40960 symbols) with the mANS decoder, dequantize, and match the reference resi_dq_y checkpoint — BYTE-EXACT (max|Δ| = 0). Proves the mANS output IS the dequant input. - e2e synthesis on the real reference latent: decode_y(y_lsbs)→ypix at 58.85 dB and decode_uv(concat(y_lsbs,uv_lsbs))→uvpix at 58.88 dB — the full synthesis stage validated on real pipeline data, and the UV input assembly (y‖uv) pinned. All feature-gated on the checkpoint/fixtures. Builds toward decode(bitstream)→pixels. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_pipeline.rs | 34 ++++++++++++++ wai/wai-rs/src/lib.rs | 4 ++ wai/wai-rs/tests/jpegai_e2e_dequant.rs | 59 ++++++++++++++++++++++++ wai/wai-rs/tests/jpegai_e2e_synthesis.rs | 22 +++++---- 4 files changed, 109 insertions(+), 10 deletions(-) create mode 100644 wai/wai-rs/src/jpegai_pipeline.rs create mode 100644 wai/wai-rs/tests/jpegai_e2e_dequant.rs diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs new file mode 100644 index 000000000..f03b9b70b --- /dev/null +++ b/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 { + 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); + } +} diff --git a/wai/wai-rs/src/lib.rs b/wai/wai-rs/src/lib.rs index 803c51328..287ff1090 100644 --- a/wai/wai-rs/src/lib.rs +++ b/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 diff --git a/wai/wai-rs/tests/jpegai_e2e_dequant.rs b/wai/wai-rs/tests/jpegai_e2e_dequant.rs new file mode 100644 index 000000000..2fcba7db5 --- /dev/null +++ b/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 = m["dt"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect(); + let bounds: Vec = m["bounds"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect(); + let mem: Vec = m["mem"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect(); + let offsets: Vec = 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 = Vec::new(); + for call in m["calls"].as_array().unwrap() { + let sigma: Vec = call["sigma"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect(); + let mask: Vec = 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 = d["gain_map"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(); + let resi_ref: Vec = 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"); +} diff --git a/wai/wai-rs/tests/jpegai_e2e_synthesis.rs b/wai/wai-rs/tests/jpegai_e2e_synthesis.rs index 8e4b4601d..3fbcf4522 100644 --- a/wai/wai-rs/tests/jpegai_e2e_synthesis.rs +++ b/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, &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 = 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)"); } -- 2.54.0 (Apple Git-157)