From d318d38967323c7575fb918325c09cbf9c9241d8 Mon Sep 17 00:00:00 2001 From: David Charlot Date: Tue, 7 Jul 2026 15:08:25 -0400 Subject: [PATCH 1/3] =?UTF-8?q?wai-rs:=20jpegai=5Fsynthesis=20=E2=80=94=20?= =?UTF-8?q?restore=20chroma=20+=20hyper=20decoders,=20add=20hyper=5Fscale?= =?UTF-8?q?=20(all=20byte-exact)=20(#53)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidated re-land of the JPEG-AI integer decode work that was rolled off main (#51/#52 were merged then force-pushed out by a concurrent main rewrite), plus the new hyper_scale_decoder. On top of the synthesis-Y decoder (#49): • decode_uv — chroma synthesis (model_uv): conv1(256->48) ‖ Slice(last96) -> Concat(144) -> conv2_t; TABs at full res; conv4_t(->8) -> DepthToSpace(CRD,2) -> 2-ch UV. 58.76 dB PSNR vs float ONNX. • decode_hyper — entropy hyper-prior decoder (common_modules): conv chain conv1(1x1)->conv2_t(4x4,up2)->ReLU6->conv3->ReLU6->conv4(->640). 69.55 dB SNR. • decode_hyper_scale — scale-index decoder (common_modules): a FULLY-INTEGER net (int64 symbols -> int32 scale indices). IntNet loader + exact integer conv + per-channel divide-by-2^n (truncation toward zero, not arithmetic shift) + QDQ clip + DepthToSpace(CRD,4) + abs + clip(0,3967). BYTE-EXACT (max|Δ| = 0) vs the onnx.reference evaluator — ORT rejects the model (no int64 Clip kernel), so the reference is the spec-compliant Python evaluator. Pure i64 => byte-identical on every machine. Structural tests on the real weights + per-net calibration tests (feature-gated on assets). Header + module docs updated. Co-authored-by: Transaction Science Co-authored-by: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_synthesis.rs | 259 +++++++++++++++++- wai/wai-rs/tests/jpegai_hyper_calibration.rs | 66 +++++ .../tests/jpegai_hyperscale_calibration.rs | 58 ++++ wai/wai-rs/tests/jpegai_uv_calibration.rs | 63 +++++ 4 files changed, 441 insertions(+), 5 deletions(-) create mode 100644 wai/wai-rs/tests/jpegai_hyper_calibration.rs create mode 100644 wai/wai-rs/tests/jpegai_hyperscale_calibration.rs create mode 100644 wai/wai-rs/tests/jpegai_uv_calibration.rs diff --git a/wai/wai-rs/src/jpegai_synthesis.rs b/wai/wai-rs/src/jpegai_synthesis.rs index 12b055bc4..ab7e5ceed 100644 --- a/wai/wai-rs/src/jpegai_synthesis.rs +++ b/wai/wai-rs/src/jpegai_synthesis.rs @@ -1,8 +1,16 @@ -//! Whole-net **byte-exact integer JPEG-AI synthesis** (model_y, tools_0) — threads -//! the real Restormer topology (traced from the ISO reference ONNX) through the -//! integer ops in [`crate::int_transformer`] + [`crate::int_hyper_synth`], loading -//! 16-bit-PTQ'd weights the sink supplies. Pure i64 ⇒ byte-identical on every -//! machine; backs `wai.neural.jpeg_ai` / `intent = replicate`. +//! Whole-net **byte-exact integer JPEG-AI synthesis** (model_y + model_uv, tools_0) +//! — threads the real Restormer topology (traced from the ISO reference ONNX) +//! through the integer ops in [`crate::int_transformer`] + [`crate::int_hyper_synth`], +//! loading 16-bit-PTQ'd weights the sink supplies. Pure i64 ⇒ byte-identical on +//! every machine; backs `wai.neural.jpeg_ai` / `intent = replicate`. +//! +//! [`decode_y`] reconstructs the luma plane [1×256×256] from the 160-channel latent; +//! [`decode_uv`] reconstructs the chroma planes [2×256×256] from the 256-channel +//! (y‖uv) latent. Both validated vs the float ONNX at ~58 dB PSNR (16-bit PTQ). +//! [`decode_hyper`] is the entropy-model hyper-prior decoder (common_modules) — +//! a plain conv chain turning the hyper-latent into the SGMM params (~69.6 dB). +//! [`decode_hyper_scale`] is the fully-integer scale-index decoder (int symbols → +//! int scale indices); it is BYTE-EXACT vs the spec reference (max|Δ| = 0), no PTQ. //! //! Topology (16×16-block trace, 256×256 image): latent[160,16,16] //! first_stage(res-block + ⇑2 + gated) → conv2_t(⇑2) → CAB(dual-branch trunk×mask @@ -218,6 +226,180 @@ pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec { }).collect() } +/// Decode the **U,V chroma planes** [2×256×256] from the 256-channel 16×16 latent +/// (the concatenated y+uv latent-space features the sink supplies as input `y`). +/// Pure-integer, deterministic; reproduces the float ONNX chroma synthesis +/// (model_uv, tools_0). +/// +/// model_uv is a **structurally distinct** net from model_y (same middle, different +/// ends): +/// • front-end: `conv1(256→48)` ‖ `Slice(input, last 96 ch)` → `Concat(144)` → +/// `conv2_t(144→64, ⇑2)` — no first_stage residual/upsample/gate; +/// • TAM: the two TABs run **at the full DepthToSpace resolution** — no ds/us conv; +/// • tail: `conv4_t(64→8, ⇑2)` → `DepthToSpace(CRD,2)` → denorm → **2-channel UV**. +/// Widths are 64 (vs model_y's 128). CAB / iact gating / TAB / denorm are identical. +pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec { + let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5 + let (h0, w0) = (16usize, 16usize); + let cin = 256usize; + let hw = h0 * w0; + // first_stage: conv1 (256→48); Concat with Slice(input, starts=-96, axis=1) — + // the last 96 channels of the input — → 144 ch; conv2_t (144→64, ⇑2) + let (c1, _, _) = wt.conv("first_stage.conv1", latent, cin, h0, w0, 3, 1, 1, 1); + let mut cat = Vec::with_capacity((48 + 96) * hw); + cat.extend_from_slice(&c1); // 48 ch (each hw contiguous) + cat.extend_from_slice(&latent[(cin - 96) * hw..]); // last 96 ch of the input + let (xc, mut h, mut w) = wt.convt("conv2_t", &cat, 144, h0, w0, 3, 2, 1, 1); // →[64,32,32] + let c = 64usize; + // CAB: trunk (2 res blocks) × sigmoid(mask branch: subscale ⇓2 → 2 res → upscale ⇑2) + // + conv2_t residual — identical shape to model_y, width 64 + let trunk = { + let t0 = res_block("CAB.residual_trunk.0", &xc, c, h, w, wt); + res_block("CAB.residual_trunk.1", &t0, c, h, w, wt) + }; + let mask = { + let (s0, sh, sw) = wt.conv("CAB.subscale", &xc, c, h, w, 3, 2, 1, 1); // ⇓2 + let m2 = res_block("CAB.residual_mask2.0", &s0, c, sh, sw, wt); + let m3 = res_block("CAB.residual_mask3.0", &m2, c, sh, sw, wt); + let (u, _, _) = wt.convt("CAB.upscale", &m3, c, sh, sw, 3, 2, 1, 1); // ⇑2 + u.iter().map(|&v| it::sigmoid(v)).collect::>() + }; + let mut x = eadd(&xc, &emul(&trunk, &mask)); + // iact2 gated: ReLU6 feeds the conv path, gate uses RAW x + let xcl = relu6_of(&x); + let (g, _, _) = wt.conv("iact2.conv", &xcl, c, h, w, 3, 1, 1, 4); + let (g, _, _) = wt.conv("iact2.conv2", &g, c, h, w, 1, 1, 0, 1); + x = gate(&x, &g); + // conv3_t (1×1) 64→256; DepthToSpace ⇑2 → [64,64,64] + let (x3, _, _) = wt.conv("conv3_t", &x, c, h, w, 1, 1, 0, 1); + let (xd, cd, nh, nw) = depth_to_space_crd(&x3, 256, h, w, 2); + (h, w) = (nh, nw); + // TAM: TAB0 → TAB1 directly at full resolution (model_uv has no ds/us conv) + let mut xt = tab("TAM.TABs.0", &xd, cd, 4, h, w, wt, eps); + xt = tab("TAM.TABs.1", &xt, cd, 4, h, w, wt, eps); + // iact3 gated: ReLU6 feeds the conv path, gate uses RAW xt + let xcl = relu6_of(&xt); + let (g, _, _) = wt.conv("iact3.conv", &xcl, c, h, w, 3, 1, 1, 4); + let (g, _, _) = wt.conv("iact3.conv2", &g, c, h, w, 1, 1, 0, 1); + x = gate(&xt, &g); + // conv4_t ⇑2 64→8 → [8,128,128]; DepthToSpace ⇑2 → [2,256,256] + let (y, yh, yw) = wt.convt("conv4_t", &x, c, h, w, 3, 2, 1, 1); + let (yd, cyd, dh, dw) = depth_to_space_crd(&y, 8, yh, yw, 2); + debug_assert_eq!((cyd, dh, dw), (2, 256, 256)); + // output denormalize: clip(·128 + 127.5, 0, 255) → u8 (2 channels) + let half = S / 2; + yd.iter().map(|&v| { + let d = (v * 128 + 127 * S + half).clamp(0, 255 * S); + ((d + half) >> F).clamp(0, 255) as u8 + }).collect() +} + +/// **hyper_decoder** (common_modules) — the entropy-model hyper-prior decoder that +/// turns the hyper-latent `z` [160×h×w] (Q16) into the 640-channel SGMM entropy +/// parameters `p` [640×2h×2w] (Q16). A plain conv chain (no transformer): +/// `conv1(1×1) → conv2_t(4×4, ⇑2) → ReLU6 → conv3(3×3) → ReLU6 → conv4(→640)`. +/// The output stays in fixed-point (it feeds the arithmetic decoder, not the +/// display), so there is no pixel denorm. Pure i64 ⇒ byte-identical everywhere. +pub fn decode_hyper(z: &[i64], cin: usize, h: usize, w: usize, wt: &Weights) -> (Vec, usize, usize) { + let (x, _, _) = wt.conv("conv1", z, cin, h, w, 1, 1, 0, 1); // 160→160, 1×1 (no bias) + let (x, h, w) = wt.convt("conv2", &x, 160, h, w, 4, 2, 1, 0); // 160→160, 4×4 ⇑2 + let x = relu6_of(&x); + let (x, _, _) = wt.conv("conv3", &x, 160, h, w, 3, 1, 1, 1); // 160→160, 3×3 + let x = relu6_of(&x); + wt.conv("conv4", &x, 160, h, w, 3, 1, 1, 1) // 160→640, 3×3 → p[640,2h,2w] +} + +/// A fully-integer weight store (int32 conv weights/biases + integer per-channel +/// power-of-two exponents) for [`decode_hyper_scale`]. Unlike [`Weights`] there is +/// NO quantization — the hyper_scale_decoder net is exact integer arithmetic, so +/// the weights are loaded verbatim as i64. +pub struct IntNet { + t: HashMap, Vec)>, +} +impl IntNet { + pub fn from_json(v: &serde_json::Value) -> Self { + let mut t = HashMap::new(); + for (k, e) in v.as_object().expect("intnet root object") { + let vals = e["v"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(); + let dims = e["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); + t.insert(k.clone(), (vals, dims)); + } + IntNet { t } + } + fn get(&self, n: &str) -> &(Vec, Vec) { + self.t.get(n).unwrap_or_else(|| panic!("missing int tensor {n}")) + } +} + +/// Exact integer conv, groups=1, i64 accumulation (no requant): output channel `o`, +/// pixel `p` = `bias[o] + Σ_i Σ_ky Σ_kx w[o,i,ky,kx]·x[i, p+ky-pad, p+kx-pad]`. +fn iconv(x: &[i64], cin: usize, h: usize, w: usize, wt: &[i64], b: &[i64], cout: usize, k: usize, stride: usize, pad: usize) -> (Vec, usize, usize) { + let ho = (h + 2 * pad - k) / stride + 1; + let wo = (w + 2 * pad - k) / stride + 1; + let mut out = vec![0i64; cout * ho * wo]; + for o in 0..cout { + let bias = b.get(o).copied().unwrap_or(0); + for oy in 0..ho { + for ox in 0..wo { + let mut acc = bias; + for i in 0..cin { + for ky in 0..k { + let iy = oy * stride + ky; + if iy < pad || iy - pad >= h { continue; } + let sy = iy - pad; + for kx in 0..k { + let ix = ox * stride + kx; + if ix < pad || ix - pad >= w { continue; } + let sx = ix - pad; + let wv = wt[((o * cin + i) * k + ky) * k + kx]; + acc += wv * x[(i * h + sy) * w + sx]; + } + } + } + out[(o * ho + oy) * wo + ox] = acc; + } + } + } + (out, ho, wo) +} + +/// **hyper_scale_decoder** (common_modules) — the scale-index hyper decoder. A +/// FULLY-INTEGER net: int symbols `z`[160×h×w] → int scale indices `I_s`[160×4h×4w]. +/// Three QDQ-style blocks (`clip(−128,127) → intconv → ÷2^n(per-channel, trunc) → +/// ReLU`) over conv1(1×1) / depthwise(3×3) / pointwise(160→2560, 1×1), then +/// `DepthToSpace(CRD, 4)` → `abs` → `clip(0, 3967)`. Byte-exact (no PTQ — the ops +/// are integer end to end); division truncates toward zero (NOT arithmetic shift). +pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNet) -> (Vec, usize, usize, usize) { + // one QDQ block: clip → intconv → per-channel ÷2^exp (trunc toward 0) → [relu] + let block = |x: &[i64], ci: usize, cname: &str, k: usize, pad: usize, ename: &str, relu_after: bool| -> (Vec, usize) { + let xq: Vec = x.iter().map(|&v| it::clip(v, -128, 127)).collect(); + let (cw, cd) = net.get(&format!("{cname}.weight")); + let (cb, _) = net.get(&format!("{cname}.bias")); + let cout = cd[0]; + let (mut y, _, _) = iconv(&xq, ci, h, w, cw, cb, cout, k, 1, pad); + let (ev, _) = net.get(ename); // per-out-channel exponent + let plane = h * w; + for o in 0..cout { + let div = 1i64 << ev[o]; // 2^n + for p in 0..plane { + let idx = o * plane + p; + let mut r = y[idx] / div; // integer division, truncates toward zero + if relu_after && r < 0 { r = 0; } + y[idx] = r; + } + } + (y, cout) + }; + let (x1, _) = block(z, cin, "conv1", 1, 0, "onnx::Pow_85", true); + let (x2, _) = block(&x1, 160, "depthwise", 3, 1, "onnx::Pow_91", true); + let (x3, c3) = block(&x2, 160, "pointwise", 1, 0, "onnx::Pow_97", false); + // DepthToSpace CRD block 4: 2560 → 160 ch, spatial ×4 + let (ds, co, ho, wo) = depth_to_space_crd(&x3, c3, h, w, 4); + // I_s = clip(abs(·), 0, 3967) + let is = ds.iter().map(|&v| v.abs().clamp(0, 3967)).collect(); + (is, co, ho, wo) +} + #[cfg(test)] mod tests { use super::*; @@ -242,4 +424,71 @@ mod tests { let (mn, mx) = (*y.iter().min().unwrap(), *y.iter().max().unwrap()); eprintln!("decode_y: {} px, value range [{}, {}]", y.len(), mn, mx); } + + /// Structural: with the real 16-bit-PTQ model_uv weights, the chroma net threads + /// end to end, is deterministic, and emits a 2×256×256 UV output. (PSNR-vs-float + /// is `tests/jpegai_uv_calibration.rs`, which measures ~58.8 dB.) + #[test] + fn decode_uv_runs_deterministically_on_real_weights() { + let path = match std::env::var("WAI_JPEGAI_UV_WEIGHTS") { + Ok(p) => p, + Err(_) => { eprintln!("skip: set WAI_JPEGAI_UV_WEIGHTS to the exported weights JSON"); return; } + }; + if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; } + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let wt = Weights::from_json(&v); + // deterministic synthetic 256-ch latent [256,16,16] in Q16 (a ramp) + let latent: Vec = (0..256 * 16 * 16).map(|i| ((i as i64 % 97) - 48) * (S / 64)).collect(); + let uv = decode_uv(&latent, &wt); + assert_eq!(uv.len(), 2 * 256 * 256, "UV output must be 2×256×256"); + assert_eq!(uv, decode_uv(&latent, &wt), "decode_uv must be deterministic"); + let (mn, mx) = (*uv.iter().min().unwrap(), *uv.iter().max().unwrap()); + eprintln!("decode_uv: {} samples (2ch), value range [{}, {}]", uv.len(), mn, mx); + } + + /// Structural: with the real 16-bit-PTQ hyper_decoder weights, the entropy + /// hyper-prior net threads end to end, is deterministic, and emits the + /// 640-channel params at 2× the hyper-latent resolution. (SNR-vs-float is + /// `tests/jpegai_hyper_calibration.rs`, which measures ~69.6 dB.) + #[test] + fn decode_hyper_runs_deterministically_on_real_weights() { + let path = match std::env::var("WAI_JPEGAI_HYPERDEC_WEIGHTS") { + Ok(p) => p, + Err(_) => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC_WEIGHTS to the exported weights JSON"); return; } + }; + if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; } + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let wt = Weights::from_json(&v); + let (c, h, w) = (160usize, 8usize, 8usize); + // integer-valued hyper-latent (entropy-decoded symbols) in Q16 + let z: Vec = (0..c * h * w).map(|i| ((i as i64 % 11) - 5) * S).collect(); + let (p, ph, pw) = decode_hyper(&z, c, h, w, &wt); + assert_eq!((ph, pw), (16, 16), "hyper output must be 2× (⇑2)"); + assert_eq!(p.len(), 640 * 16 * 16, "hyper output must be 640×16×16"); + assert_eq!(decode_hyper(&z, c, h, w, &wt).0, p, "decode_hyper must be deterministic"); + eprintln!("decode_hyper: {} params at {}×{}", p.len(), ph, pw); + } + + /// Structural: with the real integer hyper_scale_decoder weights, the fully-integer + /// scale-index net threads end to end, is deterministic, and emits `160×4h×4w` scale + /// indices in [0, 3967]. (Byte-exact-vs-reference is + /// `tests/jpegai_hyperscale_calibration.rs`, which asserts max|Δ| = 0.) + #[test] + fn decode_hyper_scale_runs_deterministically_on_real_weights() { + let path = match std::env::var("WAI_JPEGAI_HYPERSCALE_WEIGHTS") { + Ok(p) => p, + Err(_) => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_WEIGHTS to the integer weights JSON"); return; } + }; + if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; } + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let net = IntNet::from_json(&v); + let (c, h, w) = (160usize, 8usize, 8usize); + let z: Vec = (0..c * h * w).map(|i| ((i as i64 * 37) % 291) - 140).collect(); + let (is, co, ho, wo) = decode_hyper_scale(&z, c, h, w, &net); + assert_eq!((co, ho, wo), (160, 32, 32), "hyper_scale output must be 160×4h×4w"); + assert_eq!(is.len(), 160 * 32 * 32); + assert!(is.iter().all(|&v| (0..=3967).contains(&v)), "I_s must be clipped to [0,3967]"); + assert_eq!(decode_hyper_scale(&z, c, h, w, &net).0, is, "decode_hyper_scale must be deterministic"); + eprintln!("decode_hyper_scale: {} scale indices at {co}×{ho}×{wo}", is.len()); + } } diff --git a/wai/wai-rs/tests/jpegai_hyper_calibration.rs b/wai/wai-rs/tests/jpegai_hyper_calibration.rs new file mode 100644 index 000000000..a0c99c7de --- /dev/null +++ b/wai/wai-rs/tests/jpegai_hyper_calibration.rs @@ -0,0 +1,66 @@ +//! Calibration: integer `jpegai_synthesis::decode_hyper` vs the float reference — +//! run the SAME hyper-latent `z` through `hyper_decoder.onnx` (ort) and through +//! `decode_hyper`, compare the 640-channel entropy params `p`. The output is not +//! pixels but fixed-point SGMM params, so we use a scale-invariant SNR +//! (20·log10(rms(signal)/rms(error))) plus max abs error. Needs `neural` (ort). +//! +//! env: WAI_JPEGAI_HYPERDEC = path to common_modules/hyper_decoder.onnx, +//! WAI_JPEGAI_HYPERDEC_WEIGHTS = path to the exported 16-bit weights JSON. +#![cfg(feature = "neural")] + +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::Value; +use wai::jpegai_synthesis::{decode_hyper, Weights}; + +fn ort_run(path: &str, x: Vec, c: usize, h: usize, w: usize) -> (Vec, Vec) { + let mut s = Session::builder().unwrap() + .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap() + .with_optimization_level(GraphOptimizationLevel::Level3).unwrap() + .commit_from_file(path).unwrap(); + let name = s.inputs().first().unwrap().name().to_string(); + let arr = ndarray::Array4::from_shape_vec((1, c, h, w), x).unwrap(); + let v = Value::from_array(arr).unwrap(); + let out = s.run(ort::inputs![name.as_str() => v]).unwrap(); + let view = out[0].try_extract_array::().unwrap(); + (view.iter().copied().collect(), view.shape().to_vec()) +} + +#[test] +fn hyper_decoder_int_vs_float_snr() { + let onnx = match std::env::var("WAI_JPEGAI_HYPERDEC") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC to hyper_decoder.onnx"); return; } + }; + let wjson = match std::env::var("WAI_JPEGAI_HYPERDEC_WEIGHTS") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC_WEIGHTS to the weights JSON"); return; } + }; + let (c, h, w) = (160usize, 8usize, 8usize); + let n = c * h * w; + // integer-valued hyper-latent (entropy-decoded symbols), small magnitude + let lf: Vec = (0..n).map(|i| ((i % 11) as f32) - 5.0).collect(); + + let (pf, shape) = ort_run(&onnx, lf.clone(), c, h, w); + let (fmn, fmx) = pf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v))); + eprintln!("float hyper out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, pf.len(), fmn, fmx); + + let li: Vec = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect(); + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap(); + let wt = Weights::from_json(&v); + let (pi, ph, pw) = decode_hyper(&li, c, h, w, &wt); + eprintln!("int hyper out = {} vals at [{},{}]", pi.len(), ph, pw); + + assert_eq!(pi.len(), pf.len(), "output size mismatch (640×2h×2w expected)"); + let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64); + for i in 0..pi.len() { + let f = pf[i] as f64; + let q = pi[i] as f64 / 65536.0; // dequantize Q16 + let d = (q - f).abs(); + maxd = maxd.max(d); + se += d * d; + sig += f * f; + } + let snr = 10.0 * (sig / se).log10(); + eprintln!("INT vs FLOAT hyper_decoder: max|Δ| = {:.4} SNR = {:.2} dB", maxd, snr); + assert!(snr > 45.0, "hyper_decoder SNR {snr:.2} dB too low (expected high — pure conv chain)"); +} diff --git a/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs b/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs new file mode 100644 index 000000000..8b1c959f0 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs @@ -0,0 +1,58 @@ +//! Calibration: integer `jpegai_synthesis::decode_hyper_scale` vs the spec-compliant +//! reference. The hyper_scale_decoder net is FULLY INTEGER (int64 symbols in → int32 +//! scale indices out) — but ONNX Runtime rejects it (`Clip` has no int64 kernel), so +//! the reference is produced by the pure-Python `onnx.reference.ReferenceEvaluator` +//! and saved to a JSON vector; this test replays the SAME deterministic `z` through +//! `decode_hyper_scale` and asserts BYTE-EXACT equality (max|Δ| = 0). No ort needed. +//! +//! env: WAI_JPEGAI_HYPERSCALE_WEIGHTS = exported integer weights JSON, +//! WAI_JPEGAI_HYPERSCALE_REF = reference-output JSON {shape, out:[i64]}. +//! +//! Generate the reference vector (matching this test's deterministic `z`) with the +//! pure-Python ONNX reference evaluator: +//! python -c "import numpy as np,json; from onnx.reference import ReferenceEvaluator; \ +//! n=160*8*8; z=np.array([((i*37)%291)-140 for i in range(n)],dtype=np.int64).reshape(1,160,8,8); \ +//! s=ReferenceEvaluator('hyper_scale_decoder.onnx'); o=s.run(None,{s.input_names[0]:z})[0]; \ +//! json.dump({'shape':list(o.shape),'out':o.astype('int64').flatten().tolist()}, open('ref.json','w'))" +#![cfg(feature = "neural_int")] + +use wai::jpegai_synthesis::{decode_hyper_scale, IntNet}; + +#[test] +fn hyper_scale_decoder_byte_exact_vs_reference() { + let wjson = match std::env::var("WAI_JPEGAI_HYPERSCALE_WEIGHTS") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_WEIGHTS to the weights JSON"); return; } + }; + let rjson = match std::env::var("WAI_JPEGAI_HYPERSCALE_REF") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_REF to the reference JSON"); return; } + }; + let (c, h, w) = (160usize, 8usize, 8usize); + let n = c * h * w; + // SAME deterministic int symbols as the reference generator + let zi: Vec = (0..n).map(|i| ((i as i64 * 37) % 291) - 140).collect(); + + let rv: serde_json::Value = serde_json::from_slice(&std::fs::read(&rjson).unwrap()).unwrap(); + let refv: Vec = rv["out"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(); + let shape: Vec = rv["shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); + eprintln!("reference shape = {shape:?}, {} values", refv.len()); + + let wv: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap(); + let net = IntNet::from_json(&wv); + let (mine, co, ho, wo) = decode_hyper_scale(&zi, c, h, w, &net); + eprintln!("int hyper_scale out = {} vals at {co}×{ho}×{wo}", mine.len()); + + assert_eq!(mine.len(), refv.len(), "output size mismatch (160×32×32 expected)"); + let mut maxd = 0i64; + let mut mism = 0usize; + let mut first = None; + for i in 0..mine.len() { + let d = (mine[i] - refv[i]).abs(); + if d != 0 { mism += 1; if first.is_none() { first = Some((i, mine[i], refv[i])); } } + maxd = maxd.max(d); + } + eprintln!("INT vs REF hyper_scale_decoder: max|Δ| = {maxd} mismatched = {mism}/{}", mine.len()); + if let Some((i, m, r)) = first { eprintln!("first mismatch @ {i}: mine={m} ref={r}"); } + assert_eq!(maxd, 0, "hyper_scale_decoder must be BYTE-EXACT (fully-integer net); {mism} mismatches"); +} diff --git a/wai/wai-rs/tests/jpegai_uv_calibration.rs b/wai/wai-rs/tests/jpegai_uv_calibration.rs new file mode 100644 index 000000000..d82c4b07e --- /dev/null +++ b/wai/wai-rs/tests/jpegai_uv_calibration.rs @@ -0,0 +1,63 @@ +//! Calibration: my integer `jpegai_synthesis::decode_uv` vs the float reference — +//! run the SAME 256-channel latent through `synthesis16.onnx` (model_uv, ort) and +//! through `decode_uv`, compare the UV planes (max abs diff + PSNR). This is the +//! numerical-faithfulness pass for the chroma decoder. Needs `neural` (ort). +//! +//! env: WAI_JPEGAI_SYNTH16_UV = path to model_uv/synthesis16.onnx, +//! WAI_JPEGAI_UV_WEIGHTS = path to the exported 16-bit uv weights JSON. +#![cfg(feature = "neural")] + +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::Value; +use wai::jpegai_synthesis::{decode_uv, Weights}; + +fn ort_synth(path: &str, x: Vec) -> (Vec, Vec) { + let mut s = Session::builder().unwrap() + .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap() + .with_optimization_level(GraphOptimizationLevel::Level3).unwrap() + .commit_from_file(path).unwrap(); + let name = s.inputs().first().unwrap().name().to_string(); + let arr = ndarray::Array4::from_shape_vec((1, 256, 16, 16), x).unwrap(); + let v = Value::from_array(arr).unwrap(); + let out = s.run(ort::inputs![name.as_str() => v]).unwrap(); + let view = out[0].try_extract_array::().unwrap(); + (view.iter().copied().collect(), view.shape().to_vec()) +} + +#[test] +fn synthesis_uv_int_vs_float_psnr() { + let onnx = match std::env::var("WAI_JPEGAI_SYNTH16_UV") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_SYNTH16_UV to model_uv/synthesis16.onnx"); return; } + }; + let wjson = match std::env::var("WAI_JPEGAI_UV_WEIGHTS") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_UV_WEIGHTS to the weights JSON"); return; } + }; + // deterministic small latent [256,16,16] (typical post-LSBS magnitude) + let n = 256 * 16 * 16; + let lf: Vec = (0..n).map(|i| (((i % 97) as f32) - 48.0) / 96.0).collect(); + + let (uvf, shape) = ort_synth(&onnx, lf.clone()); + let (fmn, fmx) = uvf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v))); + eprintln!("float uv-synth out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, uvf.len(), fmn, fmx); + eprintln!("float sample [0..6] = {:?}", &uvf[..6]); + + let li: Vec = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect(); + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap(); + let wt = Weights::from_json(&v); + let uvi = decode_uv(&li, &wt); + + assert_eq!(uvi.len(), uvf.len(), "output size mismatch (2×256×256 expected)"); + let (mut maxd, mut mse) = (0f64, 0f64); + for i in 0..uvi.len() { + let f = uvf[i].clamp(0.0, 255.0) as f64; + let d = (uvi[i] as f64 - f).abs(); + maxd = maxd.max(d); + mse += d * d; + } + mse /= uvi.len() as f64; + let psnr = if mse > 0.0 { 20.0 * (255.0f64).log10() - 10.0 * mse.log10() } else { f64::INFINITY }; + eprintln!("INT vs FLOAT uv-synthesis: max|Δ| = {:.2} MSE = {:.2} PSNR = {:.2} dB", maxd, mse, psnr); + assert!(psnr > 45.0, "uv-synthesis PSNR {psnr:.2} dB too low (expected ~55-58)"); +} -- 2.54.0 (Apple Git-157) From 9d7dae99917a26e9a23e861066e53526de93061c Mon Sep 17 00:00:00 2001 From: David Charlot Date: Tue, 7 Jul 2026 17:26:51 -0400 Subject: [PATCH 2/3] =?UTF-8?q?wai-rs:=20jpegai=5Fsynthesis=20=E2=80=94=20?= =?UTF-8?q?add=20MCM=20context-model=20decoders=20(mcm=5Fstage0/mcm=5Fstag?= =?UTF-8?q?e=5Fn)=20(#54)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the 4-stage Multi-stage Context Model (common_modules/MCM) to the JPEG-AI integer decode path. Each stage refines the entropy context, conditioned on the stage's prior `p_i` and the accumulated decoded context `y[ conv1 -> ReLU -> conv2 -> ReLU -> conv3 -> y[0] (bias-free 1x1, width 160) • stage i>=1: cy = conv.1(conv.0(ctx)); Concat(cy, p_i) [->320] -> fusion.conv1 -> ReLU -> conv2 -> ReLU -> conv3 -> y[i] (conv.0 1x1 ctx_ch->160, conv.1 3x3 grouped-5 +bias; ctx = concat(y[i-1..0])) Float nets -> Q16 PTQ (mcm_stage0/mcm_stage_n reuse the shared Weights loader + conv path). Each stage validated independently against the float ONNX via ort: SNR 73.86 / 68.94 / 61.98 / 61.18 dB (stages 0..3) — within the 16-bit PTQ. Structural test on the real per-stage weights + tests/jpegai_mcm_calibration.rs. Co-authored-by: Transaction Science Co-authored-by: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_synthesis.rs | 58 +++++++++++++++ wai/wai-rs/tests/jpegai_mcm_calibration.rs | 85 ++++++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 wai/wai-rs/tests/jpegai_mcm_calibration.rs diff --git a/wai/wai-rs/src/jpegai_synthesis.rs b/wai/wai-rs/src/jpegai_synthesis.rs index ab7e5ceed..c693ce741 100644 --- a/wai/wai-rs/src/jpegai_synthesis.rs +++ b/wai/wai-rs/src/jpegai_synthesis.rs @@ -11,6 +11,8 @@ //! a plain conv chain turning the hyper-latent into the SGMM params (~69.6 dB). //! [`decode_hyper_scale`] is the fully-integer scale-index decoder (int symbols → //! int scale indices); it is BYTE-EXACT vs the spec reference (max|Δ| = 0), no PTQ. +//! [`mcm_stage0`]/[`mcm_stage_n`] are the 4-stage multi-stage context model (MCM) — +//! small conv fusion nets refining the entropy context per decode stage (~61–74 dB). //! //! Topology (16×16-block trace, 256×256 image): latent[160,16,16] //! first_stage(res-block + ⇑2 + gated) → conv2_t(⇑2) → CAB(dual-branch trunk×mask @@ -400,6 +402,34 @@ pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNe (is, co, ho, wo) } +/// **MCM stage 0** (common_modules/MCM) — the multi-stage context model's first +/// stage: `p0 → conv1 → ReLU → conv2 → ReLU → conv3 → y[0]` (bias-free 1×1 convs, +/// width 160). Q16 fixed-point; `wt` = this stage's own PTQ weights. +pub fn mcm_stage0(p0: &[i64], h: usize, w: usize, wt: &Weights) -> Vec { + let (x, _, _) = wt.conv("fusion_pred_net.conv1", p0, 160, h, w, 1, 1, 0, 1); + let x = relu(&x); + let (x, _, _) = wt.conv("fusion_pred_net.conv2", &x, 160, h, w, 1, 1, 0, 1); + let x = relu(&x); + wt.conv("fusion_pred_net.conv3", &x, 160, h, w, 1, 1, 0, 1).0 +} + +/// **MCM stage i≥1** — process the accumulated context then fuse with prior `p`: +/// `cy = conv.1(conv.0(ctx))` (conv.0 1×1 `ctx_ch→160`, conv.1 3×3 grouped-5 +bias), +/// then `Concat(cy, p)` [→320] → `fusion.conv1 → ReLU → conv2 → ReLU → conv3 → y[i]`. +/// `ctx` is the caller-assembled context `concat(y[i-1..0])` (160/320/480 ch). +pub fn mcm_stage_n(p: &[i64], ctx: &[i64], ctx_ch: usize, h: usize, w: usize, wt: &Weights) -> Vec { + let (c0, _, _) = wt.conv("conv.0", ctx, ctx_ch, h, w, 1, 1, 0, 1); // ctx_ch→160, 1×1 + let (cy, _, _) = wt.conv("conv.1", &c0, 160, h, w, 3, 1, 1, 5); // 3×3 grouped-5, +bias + let mut cat = Vec::with_capacity(320 * h * w); + cat.extend_from_slice(&cy); // 160 ch + cat.extend_from_slice(p); // 160 ch → 320 + let (x, _, _) = wt.conv("fusion_pred_net.conv1", &cat, 320, h, w, 1, 1, 0, 1); // 320→160 + let x = relu(&x); + let (x, _, _) = wt.conv("fusion_pred_net.conv2", &x, 160, h, w, 1, 1, 0, 1); + let x = relu(&x); + wt.conv("fusion_pred_net.conv3", &x, 160, h, w, 1, 1, 0, 1).0 +} + #[cfg(test)] mod tests { use super::*; @@ -491,4 +521,32 @@ mod tests { assert_eq!(decode_hyper_scale(&z, c, h, w, &net).0, is, "decode_hyper_scale must be deterministic"); eprintln!("decode_hyper_scale: {} scale indices at {co}×{ho}×{wo}", is.len()); } + + /// Structural: the 4 MCM context stages thread on the real per-stage weights, + /// are deterministic, and emit 160-channel context features at the input + /// resolution. (SNR-vs-float is `tests/jpegai_mcm_calibration.rs`, ~61–74 dB.) + #[test] + fn mcm_stages_run_deterministically_on_real_weights() { + let dir = match std::env::var("WAI_JPEGAI_MCM_WEIGHTS_DIR") { + Ok(p) => p, + Err(_) => { eprintln!("skip: set WAI_JPEGAI_MCM_WEIGHTS_DIR to the MCM weights dir"); return; } + }; + let load = |i: usize| { + let p = format!("{dir}/mcm_stage{i}_y.json"); + if !std::path::Path::new(&p).exists() { return None; } + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&p).unwrap()).unwrap(); + Some(Weights::from_json(&v)) + }; + let (Some(w0), Some(w1)) = (load(0), load(1)) else { eprintln!("skip: MCM weights not present"); return; }; + let (h, w) = (16usize, 16usize); + let p: Vec = (0..160 * h * w).map(|i| ((i as i64 % 71) - 35) * (S / 64)).collect(); + let y0 = mcm_stage0(&p, h, w, &w0); + assert_eq!(y0.len(), 160 * h * w); + assert_eq!(mcm_stage0(&p, h, w, &w0), y0, "mcm_stage0 must be deterministic"); + // stage1: context = y0 (160 ch) + let y1 = mcm_stage_n(&p, &y0, 160, h, w, &w1); + assert_eq!(y1.len(), 160 * h * w); + assert_eq!(mcm_stage_n(&p, &y0, 160, h, w, &w1), y1, "mcm_stage_n must be deterministic"); + eprintln!("MCM: y0 {} + y1 {} context features (160ch @ {h}×{w})", y0.len(), y1.len()); + } } diff --git a/wai/wai-rs/tests/jpegai_mcm_calibration.rs b/wai/wai-rs/tests/jpegai_mcm_calibration.rs new file mode 100644 index 000000000..e21c4e049 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_mcm_calibration.rs @@ -0,0 +1,85 @@ +//! Calibration: integer MCM stage decoders (`mcm_stage0` / `mcm_stage_n`) vs the +//! float reference — run the SAME inputs through each `MCM/stage{i}.onnx` (ort) and +//! through the Rust stage, compare the fixed-point context features `y[i]` (SNR, +//! since the output is params not pixels). Each stage validated independently. +//! Needs `neural` (ort). +//! +//! env: WAI_JPEGAI_MCM_DIR = dir with stage0..3.onnx, +//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json +#![cfg(feature = "neural")] + +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::Value; +use wai::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights}; + +fn arr(c: usize, h: usize, w: usize, seed: i64) -> (Vec, Vec) { + let n = c * h * w; + let f: Vec = (0..n).map(|i| ((((i as i64 + seed) * 31) % 197 - 98) as f32) / 64.0).collect(); + let q: Vec = f.iter().map(|&x| (x * 65536.0).round() as i64).collect(); + (f, q) +} + +fn run_stage(path: &str, named: Vec<(&str, Vec, usize)>, h: usize, w: usize) -> Vec { + let mut s = Session::builder().unwrap() + .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap() + .with_optimization_level(GraphOptimizationLevel::Level3).unwrap() + .commit_from_file(path).unwrap(); + let inputs: Vec<(String, Value)> = named.into_iter().map(|(name, data, c)| { + let a = ndarray::Array4::from_shape_vec((1, c, h, w), data).unwrap(); + (name.to_string(), Value::from_array(a).unwrap().into_dyn()) + }).collect(); + let out = s.run(inputs).unwrap(); + out[0].try_extract_array::().unwrap().iter().copied().collect() +} + +fn snr(int_q16: &[i64], flt: &[f32]) -> (f64, f64) { + let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64); + for i in 0..int_q16.len() { + let q = int_q16[i] as f64 / 65536.0; + let f = flt[i] as f64; + let d = (q - f).abs(); + maxd = maxd.max(d); se += d * d; sig += f * f; + } + (maxd, if se > 0.0 { 10.0 * (sig / se).log10() } else { f64::INFINITY }) +} + +#[test] +fn mcm_stages_int_vs_float_snr() { + let dir = match std::env::var("WAI_JPEGAI_MCM_DIR") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_MCM_DIR to the MCM onnx dir"); 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 to the weights dir"); return; } + }; + let (h, w) = (16usize, 16usize); + let load = |i: usize| -> Weights { + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap()).unwrap(); + Weights::from_json(&v) + }; + let mut worst = f64::INFINITY; + + // stage 0: single input p0 + { + let (pf, pq) = arr(160, h, w, 1); + let rf = run_stage(&format!("{dir}/stage0.onnx"), vec![("p0", pf, 160)], h, w); + let mine = mcm_stage0(&pq, h, w, &load(0)); + let (md, sn) = snr(&mine, &rf); + eprintln!("MCM stage0: max|Δ| = {md:.4} SNR = {sn:.2} dB (len {})", mine.len()); + worst = worst.min(sn); + } + // stages 1..3: inputs (p_i, context of width 160/320/480) + for (i, (pn, cn, cc)) in [("p1", "y[0]", 160usize), ("p2", "y[1-0]", 320), ("p3", "y[2-0]", 480)].into_iter().enumerate() { + let stage = i + 1; + let (pf, pq) = arr(160, h, w, 7 * stage as i64 + 1); + let (cf, cq) = arr(cc, h, w, 13 * stage as i64 + 2); + let rf = run_stage(&format!("{dir}/stage{stage}.onnx"), vec![(pn, pf, 160), (cn, cf, cc)], h, w); + let mine = mcm_stage_n(&pq, &cq, cc, h, w, &load(stage)); + let (md, sn) = snr(&mine, &rf); + eprintln!("MCM stage{stage}: max|Δ| = {md:.4} SNR = {sn:.2} dB (len {})", mine.len()); + worst = worst.min(sn); + } + eprintln!("MCM worst-stage SNR = {worst:.2} dB"); + assert!(worst > 45.0, "an MCM stage regressed (worst SNR {worst:.2} dB)"); +} -- 2.54.0 (Apple Git-157) From 4bb7467713f1b301a649ae27d55ed3db2d49ba3a Mon Sep 17 00:00:00 2001 From: David Charlot Date: Tue, 7 Jul 2026 22:38:36 -0400 Subject: [PATCH 3/3] =?UTF-8?q?wai-rs:=20byte-exact=20mANS=20entropy=20dec?= =?UTF-8?q?oder=20(jpegai=5Fmans)=20=E2=80=94=20sgm=20+=20fac=20paths=20(#?= =?UTF-8?q?55)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The entropy stage of the JPEG-AI decode path: a faithful pure-Rust port of the ISO reference mANS decoder (cpp_exts/mans/decompressor.cpp). This is the central non-neural piece — it turns the compressed bitstream into the latent symbols the synthesis/hyper nets consume. Mechanism (byte-for-byte with the reference): • two interleaved ANS states (state1/state2), advanced in a group-of-4 inbound → flush → outbound schedule; • a rolling 64-bit window read MSB-first (read/flush), consumed backward through the byte buffer, with a sentinel-bit init (initBackward + initWithStates); • sgm path (main y-latent): a per-position sigma index selects one of 32 precomputed transition tables; mask=false skips the position (value 0) without consuming the stream; the outbound/escape path (symbol == -bound) reads a sign + 3/16-bit magnitude for out-of-range values; • fac path (z hyper-latent): a per-channel cumulative CDF builds the row table (transitionsZPreprocess), symbol == MAX_Z(63) escapes to a raw 6-bit value. State threads across calls, so one MansDecoder decodes a whole stream. Validated BYTE-EXACT against the ISO reference (the golden was produced by the reference `ans.so`): decode_sgm 40960/40960 exact on the real SORP stream (160 calls, 703 coded symbols, two states, escapes), decode_factorize 128/128 exact on an oracle encode→decode roundtrip. Tests feature-gated on the fixtures. Co-authored-by: Transaction Science Co-authored-by: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_mans.rs | 239 ++++++++++++++++++++ wai/wai-rs/src/lib.rs | 6 + wai/wai-rs/tests/jpegai_mans_calibration.rs | 76 +++++++ 3 files changed, 321 insertions(+) create mode 100644 wai/wai-rs/src/jpegai_mans.rs create mode 100644 wai/wai-rs/tests/jpegai_mans_calibration.rs diff --git a/wai/wai-rs/src/jpegai_mans.rs b/wai/wai-rs/src/jpegai_mans.rs new file mode 100644 index 000000000..ed4050c23 --- /dev/null +++ b/wai/wai-rs/src/jpegai_mans.rs @@ -0,0 +1,239 @@ +//! Byte-exact **mANS** (masked-ANS) entropy decoder for JPEG-AI — a faithful port +//! of the ISO reference `cpp_exts/mans/decompressor.cpp`. Two interleaved ANS states, +//! MSB-first reads from a rolling 64-bit window (`read`/`flush`), sentinel-bit init, +//! and a group-of-4 inbound→flush→outbound schedule. Pure integer ⇒ byte-identical +//! on every machine. Backs the entropy stage of `wai.neural.jpeg_ai`. +//! +//! `decode_sgm` (main y-latent): per position a distribution `sigma_idx` selects one +//! of the 32 precomputed transition tables; `mask=false` ⇒ the position is skipped +//! (value 0) and does not consume the stream. The state threads across calls, so a +//! single [`MansDecoder`] decodes a whole stream over many `decode_sgm` invocations. + +/// Low 56..63-bit masks for `flush()` (BIT_MASK_HIGH in the reference). +const BIT_MASK_HIGH: [u64; 8] = [ + 0x00ff_ffff_ffff_ffff, 0x01ff_ffff_ffff_ffff, 0x03ff_ffff_ffff_ffff, 0x07ff_ffff_ffff_ffff, + 0x0fff_ffff_ffff_ffff, 0x1fff_ffff_ffff_ffff, 0x3fff_ffff_ffff_ffff, 0x7fff_ffff_ffff_ffff, +]; + +/// The backward-consumed 64-bit bit window + the two interleaved ANS states. +#[derive(Clone, Copy)] +struct BitStream { + ptr_end: usize, // index into `memory` (one past the current window's low byte block) + head: u64, + bit_pos: u32, + state1: u8, + state2: u8, +} + +impl BitStream { + fn le64(mem: &[u8], at: usize) -> u64 { + let mut b = [0u8; 8]; + b.copy_from_slice(&mem[at..at + 8]); + u64::from_le_bytes(b) + } + /// Read `step` bits (MSB-first) out of the window. + fn read(&mut self, step: u32) -> u64 { + self.bit_pos -= step; + let result = self.head >> self.bit_pos; + self.head ^= result << self.bit_pos; + result + } + /// Refill the window from earlier bytes (moving backward). + fn flush(&mut self, mem: &[u8]) { + self.ptr_end -= (7 ^ (self.bit_pos >> 3)) as usize; + self.head = Self::le64(mem, self.ptr_end) & BIT_MASK_HIGH[(self.bit_pos & 7) as usize]; + self.bit_pos |= 56; + } + fn init_backward(&mut self, mem: &[u8], ptr: usize) { + self.ptr_end = ptr - 8; + self.head = Self::le64(mem, self.ptr_end); + self.bit_pos = 63 ^ self.head.leading_zeros(); + self.head ^= 1u64 << self.bit_pos; + } + fn init_with_states(&mut self, mem: &[u8], ptr: usize) { + self.init_backward(mem, ptr); + self.state1 = self.read(8) as u8; + self.state2 = self.read(8) as u8; + self.flush(mem); + } +} + +const MAX_Z: usize = 63; + +/// A byte-exact mANS decoder over one memory buffer (single-thread streams). +pub struct MansDecoder { + memory: Vec, // 8 zero bytes prepended, then the stream bytes + stream: BitStream, // persists across decode calls (state threads) + trans_y: Vec, // num_distributions * 256 flat transition table (sgm) + bounds_y: Vec, // per-distribution outbound bound (sgm) + z_pre: [u16; 512], // transitionsZPreprocess (built once, for the fac path) + trans_z: [u32; 256], // per-row Z transition table (rebuilt per channel) +} + +impl MansDecoder { + /// `mem` = the raw stream bytes; `offsets` = per-thread end offsets + /// (`cumsum(thread_sizes)`). Only single-thread (`offsets.len()==1`) is supported. + pub fn new(mem: &[u8], offsets: &[u32]) -> Self { + assert_eq!(offsets.len(), 1, "only single-thread streams supported"); + let cap = mem.len(); + let mut memory = vec![0u8; cap + 8]; + memory[8..8 + cap].copy_from_slice(mem); + let thread_off = offsets[0] as usize + 8; + let mut stream = BitStream { ptr_end: 0, head: 0, bit_pos: 0, state1: 0, state2: 0 }; + stream.init_with_states(&memory, thread_off); + // transitionsZPreprocess: for i in 511..1, nBits grows so (i< u32 { + self.trans_y[dist * 256 + state as usize] + } + + #[inline] + fn decode_y_inbound(&mut self, which: u8, sigma: &[u8], out: &mut [i16], index: usize) { + let transition = self.t_y(sigma[index] as usize, if which == 1 { self.stream.state1 } else { self.stream.state2 }); + let ns = ((transition >> 16) as u64 | self.stream.read(transition >> 24)) as u8; + if which == 1 { self.stream.state1 = ns; } else { self.stream.state2 = ns; } + out[index] = transition as i16; // low 16 bits, signed = the symbol + } + + #[inline] + fn decode_y_outbound(&mut self, sigma: &[u8], out: &mut [i16], index: usize) { + let bound = self.bounds_y[sigma[index] as usize] as i32; + if out[index] as i32 + bound == 0 { + let step = if self.stream.read(1) != 0 { 3 } else { 16 }; + let value_temp = self.stream.read(step) as i32; + let sign = value_temp & 1; + out[index] = (((value_temp >> 1) + (bound - sign)) ^ (-sign)) as i16; + self.stream.flush(&self.memory); + } + } + + /// Decode one span: `out[i] = symbol` where `mask[i]`, else left untouched (0). + /// State threads across calls. Faithful port of `decodeSingleThread`. + pub fn decode_sgm(&mut self, sigma: &[u8], out: &mut [i16], mask: &[bool]) { + let length = sigma.len(); + let mut index = 0; + // main groups of 4: state1/state2/state1/state2, flush between in/out bound + while index + 3 < length { + if mask[index] { self.decode_y_inbound(1, sigma, out, index); } + if mask[index + 1] { self.decode_y_inbound(2, sigma, out, index + 1); } + if mask[index + 2] { self.decode_y_inbound(1, sigma, out, index + 2); } + if mask[index + 3] { self.decode_y_inbound(2, sigma, out, index + 3); } + self.stream.flush(&self.memory); + if mask[index] { self.decode_y_outbound(sigma, out, index); } + if mask[index + 1] { self.decode_y_outbound(sigma, out, index + 1); } + if mask[index + 2] { self.decode_y_outbound(sigma, out, index + 2); } + if mask[index + 3] { self.decode_y_outbound(sigma, out, index + 3); } + index += 4; + } + if length & 2 != 0 { + if mask[index] { + self.decode_y_inbound(1, sigma, out, index); + self.decode_y_outbound(sigma, out, index); + } + if mask[index + 1] { + self.decode_y_inbound(2, sigma, out, index + 1); + self.decode_y_outbound(sigma, out, index + 1); + } + index += 2; + self.stream.flush(&self.memory); + } + if length & 1 != 0 { + if mask[index] { + self.decode_y_inbound(1, sigma, out, index); + self.decode_y_outbound(sigma, out, index); + } + self.stream.flush(&self.memory); + } + } + + // ---- fac (factorized z hyper-latent) path ---- + + /// Build the per-row Z transition table from a channel's cumulative `cdfs` + /// (`MAX_Z` entries). Faithful port of `setZDistributions`. + fn set_z_distributions(&mut self, cdfs: &[u8]) { + let mut curr = 0usize; + for i in 0..MAX_Z { + let c = cdfs[i] as usize; + let mut k = c - curr; // = count for symbol i + for j in curr..c { + self.trans_z[j] = ((self.z_pre[k] as u32) << 16) | i as u32; + k += 1; + } + curr = c; + } + self.trans_z[255] = ((self.z_pre[1] as u32) << 16) | MAX_Z as u32; + } + + #[inline] + fn decode_z_inbound(&mut self, which: u8, out: &mut [u8], index: usize) { + let transition = self.trans_z[if which == 1 { self.stream.state1 } else { self.stream.state2 } as usize]; + let ns = ((transition >> 16) as u64 | self.stream.read(transition >> 24)) as u8; + if which == 1 { self.stream.state1 = ns; } else { self.stream.state2 = ns; } + out[index] = transition as u8; // low 8 bits = symbol (0..=MAX_Z) + } + + #[inline] + fn decode_z_outbound(&mut self, out: &mut [u8], index: usize) { + if out[index] as usize == MAX_Z { + out[index] = self.stream.read(6) as u8; + } + } + + fn decode_row(&mut self, cdfs: &[u8], out: &mut [u8]) { + self.set_z_distributions(cdfs); + let length = out.len(); + let mut index = 0; + while index + 3 < length { + self.decode_z_inbound(1, out, index); + self.decode_z_inbound(2, out, index + 1); + self.decode_z_inbound(1, out, index + 2); + self.decode_z_inbound(2, out, index + 3); + self.decode_z_outbound(out, index); + self.decode_z_outbound(out, index + 1); + self.decode_z_outbound(out, index + 2); + self.decode_z_outbound(out, index + 3); + self.stream.flush(&self.memory); + index += 4; + } + if length & 2 != 0 { + self.decode_z_inbound(1, out, index); + self.decode_z_outbound(out, index); + self.decode_z_inbound(2, out, index + 1); + self.decode_z_outbound(out, index + 1); + index += 2; + } + if length & 1 != 0 { + self.decode_z_inbound(1, out, index); + self.decode_z_outbound(out, index); + } + self.stream.flush(&self.memory); + } + + /// Decode the factorized z hyper-latent: `channels` rows, `size_per_channel` + /// symbols each, each row using its own `MAX_Z`-entry cumulative CDF (`cdfs` is + /// `channels × MAX_Z`). Faithful port of `decode_factorize` (single-thread). + pub fn decode_factorize(&mut self, cdfs: &[u8], out: &mut [u8], channels: usize, size_per_channel: usize) { + for i in 0..channels { + let row = &cdfs[i * MAX_Z..(i + 1) * MAX_Z]; + let (a, b) = (i * size_per_channel, (i + 1) * size_per_channel); + let mut buf = vec![0u8; size_per_channel]; + self.decode_row(row, &mut buf); + out[a..b].copy_from_slice(&buf); + } + } +} diff --git a/wai/wai-rs/src/lib.rs b/wai/wai-rs/src/lib.rs index 7194d8e1d..803c51328 100644 --- a/wai/wai-rs/src/lib.rs +++ b/wai/wai-rs/src/lib.rs @@ -73,6 +73,12 @@ pub mod int_transformer; #[cfg(feature = "neural_int")] pub mod jpegai_synthesis; +// Byte-exact mANS entropy decoder for JPEG-AI (port of the ISO reference +// decompressor.cpp) — two interleaved ANS states, sgm (y-latent) + fac (z-latent) +// paths. Pure integer ⇒ byte-identical everywhere. +#[cfg(feature = "neural_int")] +pub mod jpegai_mans; + // 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_mans_calibration.rs b/wai/wai-rs/tests/jpegai_mans_calibration.rs new file mode 100644 index 000000000..f585332dc --- /dev/null +++ b/wai/wai-rs/tests/jpegai_mans_calibration.rs @@ -0,0 +1,76 @@ +//! Byte-exact calibration for the mANS entropy decoder: replay the real JPEG-AI +//! `SORP` stream (main y-latent) through `MansDecoder::decode_sgm` and require an +//! EXACT match to the reference `expected` (the golden was produced by the ISO +//! reference `ans.so`, and this Rust is a faithful port of `decompressor.cpp`). +//! No ort needed — the fixture carries the transition tables + stream + goldens. +//! +//! env: WAI_JPEGAI_MANS_FIXTURE = path to mans_sorp_fixture.json +//! { dt:[u32; 32*256], bounds:[u8;32], mem:[u8], offsets:[u32], +//! calls:[{sigma:[u8], mask:[bool], expected:[i16]}] } +#![cfg(feature = "neural_int")] + +use wai::jpegai_mans::MansDecoder; + +#[test] +fn mans_sgm_byte_exact_vs_reference() { + let path = 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 to mans_sorp_fixture.json"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let arr_u32 = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect::>(); + let dt = arr_u32("dt"); + let bounds: Vec = v["bounds"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect(); + let mem: Vec = v["mem"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect(); + let offsets = arr_u32("offsets"); + + let mut dec = MansDecoder::new(&mem, &offsets); + dec.set_sgm_transitions(&dt, &bounds); + + let mut total = 0usize; + let mut mism = 0usize; + let mut first: Option<(usize, usize, i16, i16)> = None; + for (ci, call) in v["calls"].as_array().unwrap().iter().enumerate() { + 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 exp: Vec = call["expected"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i16).collect(); + let mut out = vec![0i16; sigma.len()]; + dec.decode_sgm(&sigma, &mut out, &mask); + for i in 0..out.len() { + total += 1; + if out[i] != exp[i] { + mism += 1; + if first.is_none() { first = Some((ci, i, out[i], exp[i])); } + } + } + } + eprintln!("mANS decode_sgm vs reference: {}/{} exact, {} mismatched", total - mism, total, mism); + if let Some((ci, i, got, exp)) = first { eprintln!("first mismatch: call {ci} pos {i}: got {got} exp {exp}"); } + assert_eq!(mism, 0, "mANS decoder must be BYTE-EXACT vs the ISO reference; {mism} mismatches"); +} + +/// Byte-exact `fac` (factorized z hyper-latent) path vs an oracle roundtrip vector. +/// env: WAI_JPEGAI_MANS_FAC = path to mans_fac_fixture.json +/// { cdfs:[u8; channels*63], mem:[u8], offsets:[u32], channels, spc, expected:[u8] } +#[test] +fn mans_fac_byte_exact_vs_reference() { + let path = match std::env::var("WAI_JPEGAI_MANS_FAC") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_MANS_FAC to mans_fac_fixture.json"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let u8v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::>(); + let cdfs = u8v("cdfs"); + let mem = u8v("mem"); + let offsets: Vec = v["offsets"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect(); + let channels = v["channels"].as_u64().unwrap() as usize; + let spc = v["spc"].as_u64().unwrap() as usize; + let expected = u8v("expected"); + + let mut dec = MansDecoder::new(&mem, &offsets); + let mut out = vec![0u8; channels * spc]; + dec.decode_factorize(&cdfs, &mut out, channels, spc); + let mism = out.iter().zip(&expected).filter(|(a, b)| a != b).count(); + eprintln!("mANS decode_factorize vs reference: {}/{} exact, {} mismatched", out.len() - mism, out.len(), mism); + assert_eq!(mism, 0, "mANS fac decoder must be BYTE-EXACT; {mism} mismatches"); +} -- 2.54.0 (Apple Git-157)