From 85d0e3deb9b7f45f79c3736e5aecdc1c16f3cf34 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Wed, 15 Jul 2026 09:07:07 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Ficci::IcciNet=20=E2=80=94=20?= =?UTF-8?q?the=20eICCI=20filter,=20float-exact=20vs=20the=20ONNX?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes eICCI, the last unported JPEG-AI decode stage: three planes in, three refined planes out (wavelet transform included). The shipped ONNX is much smaller than the reference Python class implies — trust the export, not icci_models.py: class says nf=64, nbY=8, nbUV=8 ONNX is nf=48, nbY=2, nbUV=4 (145,265 params, 21 tensors) and the exporter has already folded/deduped nearly everything: the luma stem's BN is folded into onnx::Conv_8299; every conv bias is zero (hence deduped to one tensor); every block scale equals scaleY=0.1 (also deduped); BNUV is γ=1,β=0,μ=0,σ²=1, i.e. x/√(1+ε), folded into conv_firstUV here. IcciNet also drops a redundancy: the reference's process_U and process_V each rebuild the same chroma features, so it runs that trunk twice. The trunk is evaluated once and fed to both heads — identical output, half the chroma work. Validated against the ONNX itself (planes → planes, ort): max|Δ| = 0.0 / 1e-7 / 0.0 for Y/U/V. Float, not quantised: the net is float in the reference, and IEEE-754 with a fixed accumulation order is deterministic across machines — so the port is byte-identical and carries no PTQ error. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_icci.rs | 51 +++++++++ wai/wai-rs/tests/jpegai_icci_calibration.rs | 112 ++++++++++++++++++++ 2 files changed, 163 insertions(+) create mode 100644 wai/wai-rs/tests/jpegai_icci_calibration.rs diff --git a/wai/wai-rs/src/jpegai_icci.rs b/wai/wai-rs/src/jpegai_icci.rs index 08972f7e5..431209ac2 100644 --- a/wai/wai-rs/src/jpegai_icci.rs +++ b/wai/wai-rs/src/jpegai_icci.rs @@ -242,6 +242,57 @@ impl IcciPath { } } +/// The whole eICCI filter for one tile: three planes in, three refined planes out. +/// +/// Shape follows the exported model, **not** the reference Python class (which says +/// `nf=64, nbY=8, nbUV=8`): the shipped nets are `nf=48`, **2** luma blocks and **4** +/// chroma blocks. The chroma trunk is shared — the reference traces it once per chroma +/// plane, computing identical features twice; we evaluate it once and feed both heads. +/// +/// ```text +/// (Y16,U16,V16), yuv48 = dwt2(Y), dwt2(U), dwt2(V), cat(..) +/// Y_o = idwt2( conv_last_Y(luma_trunk(yuv48)) · scale + Y16 ) +/// fea = chroma_trunk(yuv48) // evaluated once +/// U_o = idwt2( conv_last_u(fea) · scale + U16 ) +/// V_o = idwt2( conv_last_v(fea) · scale + V16 ) +/// ``` +#[derive(Clone)] +pub struct IcciNet { + /// luma: `conv_first_Y` (BN pre-folded by the exporter) + its residual blocks + pub luma: IcciPath, + /// chroma trunk: `conv_firstUV` (fold `BNUV` in) + its residual blocks; `conv_last` + /// is the U head and `scale` applies to both heads. + pub chroma: IcciPath, + /// the V head — shares `chroma`'s trunk, differs only in this final convolution + pub conv_last_v: Conv, +} + +impl IcciNet { + /// Refine one tile's three planes. `y`/`u`/`v` are `[h × w]` (4:4:4); `h` and `w` + /// must be multiples of 4 (two wavelet levels). + pub fn forward(&self, y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize) + -> (Vec, Vec, Vec) { + let (y16, u16, v16, yuv48, h2, w2) = icci_dwt_planes(y, u, v, h, w); + + let y_o = self.luma.forward_plane(&yuv48, &y16, h2, w2); + + // shared chroma trunk, evaluated once + let mut fea = self.chroma.conv_first.forward(&yuv48, h2, w2); + for t in fea.iter_mut() { if *t < 0.0 { *t = 0.0; } } + for blk in &self.chroma.blocks { + fea = blk.forward(&fea, h2, w2); + } + let head = |conv: &Conv, band: &[f32]| -> Vec { + let res = conv.forward(&fea, h2, w2); + let refined: Vec = res.iter().zip(band).map(|(&r, &s)| r * self.chroma.scale + s).collect(); + haar_idwt2(&refined, h2, w2).0 + }; + let u_o = head(&self.chroma.conv_last, &u16); + let v_o = head(&self.conv_last_v, &v16); + (y_o, u_o, v_o) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/wai/wai-rs/tests/jpegai_icci_calibration.rs b/wai/wai-rs/tests/jpegai_icci_calibration.rs new file mode 100644 index 000000000..3afb86a74 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_icci_calibration.rs @@ -0,0 +1,112 @@ +//! Calibration for the eICCI post-filter ([`wai::jpegai_icci::IcciNet`]) against the +//! shipped ONNX model, which is the reference: it takes the three planes and returns the +//! three refined planes, with the wavelet transform inside the graph. +//! +//! The network is float, so this is a faithful f32 port rather than a quantised one — +//! we expect a machine-precision match, not a PTQ-floor one. +//! +//! Weight layout (traced from the export, which dedupes identical tensors and pre-folds +//! most BatchNorms): +//! • `onnx::Conv_8299` luma stem, 48→48 3×3 (BN already folded) +//! • `onnx::Conv_8302/8305`, `8308/8311` the 2 luma residual blocks ([1,3] then [3,1]) +//! • `conv_last_Y` luma head, 48→16 3×3 +//! • `conv_firstUV` + `BNUV` chroma stem (BNUV is γ=1,β=0,μ=0,σ²=1 ⇒ x/√(1+ε)) +//! • `8314/8317 … 8332/8335` the 4 chroma residual blocks +//! • `conv_last_u`, `conv_last_v` the chroma heads +//! • every bias is zero; every scale is `scaleY` (0.1), deduped by the exporter +//! +//! env: WAI_JPEGAI_ICCI_WEIGHTS = icci_model0.json { name: {dims, v:[f32]} } +//! WAI_JPEGAI_ICCI_ONNX = post_filters/model_0.onnx +#![cfg(feature = "neural")] + +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::Value; +use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock}; + +const EPS: f32 = 1e-5; // ONNX BatchNormalization default + +fn conv_from(v: &serde_json::Value, name: &str, ph: usize, pw: usize) -> Conv { + let t = &v[name]; + let dims: Vec = t["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); + let w: Vec = t["v"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(); + let (cout, cin, kh, kw) = (dims[0], dims[1], dims[2], dims[3]); + Conv { w, b: vec![0f32; cout], cin, cout, kh, kw, ph, pw } +} + +fn block(v: &serde_json::Value, a: &str, b: &str, scale: f32) -> ResBlock { + ResBlock { conv1: conv_from(v, a, 0, 1), conv2: conv_from(v, b, 1, 0), scale } +} + +#[test] +fn icci_net_matches_onnx_reference() { + let wj = match std::env::var("WAI_JPEGAI_ICCI_WEIGHTS") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_WEIGHTS to icci_model0.json"); return; } + }; + let onnx = match std::env::var("WAI_JPEGAI_ICCI_ONNX") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_ONNX to model_0.onnx"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap(); + let scale = v["scaleY"]["v"][0].as_f64().unwrap() as f32; + + // chroma stem: fold BNUV (γ=1, β=0, μ=0, σ²=1) into conv_firstUV → w/√(1+ε) + let mut chroma_stem = conv_from(&v, "conv_firstUV.weight", 1, 1); + let n = chroma_stem.cout; + chroma_stem = chroma_stem.fold_bn(&vec![1.0; n], &vec![0.0; n], &vec![0.0; n], &vec![1.0; n], EPS); + + let net = IcciNet { + luma: IcciPath { + conv_first: conv_from(&v, "onnx::Conv_8299", 1, 1), + blocks: vec![ + block(&v, "onnx::Conv_8302", "onnx::Conv_8305", scale), + block(&v, "onnx::Conv_8308", "onnx::Conv_8311", scale), + ], + conv_last: conv_from(&v, "conv_last_Y.weight", 1, 1), + scale, + }, + chroma: IcciPath { + conv_first: chroma_stem, + blocks: vec![ + block(&v, "onnx::Conv_8314", "onnx::Conv_8317", scale), + block(&v, "onnx::Conv_8320", "onnx::Conv_8323", scale), + block(&v, "onnx::Conv_8326", "onnx::Conv_8329", scale), + block(&v, "onnx::Conv_8332", "onnx::Conv_8335", scale), + ], + conv_last: conv_from(&v, "conv_last_u.weight", 1, 1), + scale, + }, + conv_last_v: conv_from(&v, "conv_last_v.weight", 1, 1), + }; + + // deterministic planes at a tile size that is a multiple of 4 + let (h, w) = (32usize, 32usize); + let plane = h * w; + let mk = |seed: usize| -> Vec { + (0..plane).map(|i| (((i * 37 + seed * 11) % 251) as f32 - 125.0) / 255.0).collect() + }; + let (y, u, vv) = (mk(0), mk(1), mk(2)); + + let (ry, ru, rv) = net.forward(&y, &u, &vv, h, w); + + // reference: the ONNX itself (planes in → refined planes out) + 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(&onnx).unwrap(); + let arr = |d: &Vec| Value::from_array(ndarray::Array4::from_shape_vec((1, 1, h, w), d.clone()).unwrap()).unwrap().into_dyn(); + let out = s.run(vec![ + ("Y".to_string(), arr(&y)), + ("U".to_string(), arr(&u)), + ("V".to_string(), arr(&vv)), + ]).unwrap(); + + for (i, (name, mine)) in [("Y_o", &ry), ("U_o", &ru), ("V_o", &rv)].iter().enumerate() { + let want: Vec = out[i].try_extract_array::().unwrap().iter().copied().collect(); + assert_eq!(mine.len(), want.len(), "{name}: size mismatch"); + let mut maxd = 0f32; + for j in 0..want.len() { maxd = maxd.max((mine[j] - want[j]).abs()); } + eprintln!("eICCI {name} vs ONNX: max|Δ| = {maxd:.7}"); + assert!(maxd < 1e-3, "{name}: eICCI diverged from the ONNX (max|Δ| = {maxd})"); + } +} -- 2.54.0 (Apple Git-157)