diff --git a/wai/tools/wai_jpegai_build_e2e_inputs.py b/wai/tools/wai_jpegai_build_e2e_inputs.py new file mode 100644 index 000000000..bbdfaaa2a --- /dev/null +++ b/wai/tools/wai_jpegai_build_e2e_inputs.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Build every runtime input the whole-decoder test (`jpegai_e2e_full`) needs, from +the shipped JPEG-AI assets. Nothing here is derived or fitted — it is a straight +extraction of model weights and normative tables, so the end-to-end test is +reproducible rather than depending on hand-made scratch files. + + python wai_jpegai_build_e2e_inputs.py + + is the jpegai-vm-assets root. Writes: + /weights/{hs,hyp}_{y,uv}.json, mcm{0..3}_y.json, synth_{y,uv}.json, icci.json + /norm.json (pmf, bounds, per-channel z size, factorized CDFs, gain + vectors, LSBS tables, LEF op-point tables, reference RGB) + +The weight exports call `wai_jpegai_export_weights.py` (same directory); norm.json is +assembled here. Requires numpy + Pillow (only for reading the reference PNG). + +Model selection is signalled by the picture (icci_model_signalled_idx → Y=model_6, +UV=model_3); those two wrappers live under e2e/pfchain and are what the exporter's +`icci` mode reads. +""" +import json +import os +import subprocess +import sys + +import numpy as np +from PIL import Image + +HERE = os.path.dirname(os.path.abspath(__file__)) +EXPORT = os.path.join(HERE, "wai_jpegai_export_weights.py") + + +def export(mode, *args): + subprocess.run([sys.executable, EXPORT, mode, *args], check=True) + + +def main(assets, out): + W = os.path.join(out, "weights") + os.makedirs(W, exist_ok=True) + M = os.path.join(assets, "model/CCS_SGMM/tools_0") + M16 = os.path.join(assets, "model16/CCS_SGMM/tools_0") + + # --- model weights --- + for pl in ("y", "uv"): + export("int", f"{M}/model_{pl}/common_modules/hyper_scale_decoder.onnx", f"{W}/hs_{pl}.json") + export("ptq", f"{M}/model_{pl}/common_modules/hyper_decoder.onnx", f"{W}/hyp_{pl}.json") + export("ptq", f"{M16}/model_{pl}/synthesis16.onnx", f"{W}/synth_{pl}.json") + for i in range(4): + export("ptq", f"{M}/model_y/common_modules/MCM/stage{i}.onnx", f"{W}/mcm{i}_y.json") + export("icci", f"{assets}/e2e/pfchain/icci_y_m6.onnx", + f"{assets}/e2e/pfchain/icci_uv_m3.onnx", f"{W}/icci.json") + + # --- normative tables + per-image reference RGB --- + fac = json.load(open(f"{assets}/jpegai_factorized_cdfs.json")) + norm = json.load(open(f"{assets}/jpegai_normative_tables.json")) + mc = json.load(open(f"{assets}/model_constants.json")) + bounds = [int(x) for x in open(f"{assets}/me-tANS/bounds.csv") if x.strip()] + + def cdfs(pl): + cd = fac[pl] + return [int(x) for row in cd for x in row] if isinstance(cd[0], list) else [int(x) for x in cd] + + def gain(pl): + p = f"{M}/model_{pl}/common_modules/quantizer/gain_unit_mlog.csv" + return [int(v) for v in np.loadtxt(p, delimiter=",", dtype=int)] + + rgb = np.array(Image.open(f"{assets}/e2e/test256_dec.png").convert("RGB")) + data = { + "pmf": norm["pmf"], "bounds": bounds, "per_channel": fac["per_channel"], + "cdfs_y": cdfs("y"), "cdfs_uv": cdfs("uv"), "gain_y": gain("y"), "gain_uv": gain("uv"), + "lsbs_t0": [int(v) for v in np.load(f"{assets}/e2e/lsbs/lsbs_table0.npy")], + "lsbs_t1": [int(v) for v in np.load(f"{assets}/e2e/lsbs/lsbs_table1.npy")], + "lef_thr0": mc["lef_thr"][0], "lef_mag0": mc["lef_mag"][0], + "ref_rgb": rgb.ravel().tolist(), + } + json.dump(data, open(f"{out}/norm.json", "w")) + print(f"wrote {W}/*.json (11 weight files) and {out}/norm.json") + + +if __name__ == "__main__": + if len(sys.argv) != 3: + print(__doc__) + sys.exit(1) + main(sys.argv[1], sys.argv[2]) diff --git a/wai/tools/wai_jpegai_export_weights.py b/wai/tools/wai_jpegai_export_weights.py index 7e3f043bf..c2b00664a 100644 --- a/wai/tools/wai_jpegai_export_weights.py +++ b/wai/tools/wai_jpegai_export_weights.py @@ -205,6 +205,61 @@ def list_tensors(src): print(f"{total:,} params") +def _conv_node(w_dims, w_raw, b_raw): + """One `Conv` as `jpegai_icci`'s `conv_of` reads it. Padding is kernel//2, which + keeps the spatial size for both the 3x3 stems and the factorized 1x3 / 3x1 taps.""" + cout, cin, kh, kw = w_dims + return {"w": _decode(w_dims, 1, w_raw), "b": _decode([len(b_raw) // 4], 1, b_raw), + "cin": cin, "cout": cout, "kh": kh, "kw": kw, "ph": kh // 2, "pw": kw // 2} + + +def _icci_path(inits, n_blocks, conv_last_name, scale_name, hidden_prefix): + """One eICCI plane-path out of a wrapper's initializers. + + Heads and scales are found by NAME; the trunk convolutions are anonymous + (`onnx::Conv_*`) and are assigned by FILE ORDER — conv_first, then each residual + block's 1x3 tap followed by its 3x1 tap. The wrappers are BN-pre-folded, so + nothing is folded here. + """ + named = {n: (d, r) for n, (d, _dt, r) in inits.items()} + anon = [(d, r) for n, (d, _dt, r) in inits.items() if n.startswith("onnx::Conv")] + convs = [(anon[i][0], anon[i][1], anon[i + 1][1]) for i in range(0, len(anon), 2)] + assert len(convs) == 1 + 2 * n_blocks, \ + f"{len(convs)} convs for {n_blocks} blocks in {conv_last_name}" + scalar = lambda name: _decode(named[name][0], 1, named[name][1])[0] + blocks = [{"c1": _conv_node(*convs[1 + 2 * b]), + "c2": _conv_node(*convs[2 + 2 * b]), + "scale": scalar(f"{hidden_prefix}.{b}.scale")} for b in range(n_blocks)] + cl_d, cl_r = named[conv_last_name + ".weight"] + return {"conv_first": _conv_node(*convs[0]), + "blocks": blocks, + "conv_last": _conv_node(cl_d, cl_r, named[conv_last_name + ".bias"][1]), + "scale": scalar(scale_name)} + + +def export_icci(y_src, uv_src, dst): + """The two eICCI model wrappers -> the `IcciNet` weight JSON. + + Which wrappers is signalled per picture (`icci_model_signalled_idx` indexes a + shortlist); for the 256x256 conformance stream that is Y = model_6, UV = model_3. + The luma path has 2 residual blocks and the chroma path 4, and U and V carry + separate residual scales. + """ + luma = _icci_path(initializers(y_src), 2, "comp_model.conv_last_Y", + "comp_model.scaleY", "comp_model.hidden_layer_Y") + ui = initializers(uv_src) + chroma = _icci_path(ui, 4, "comp_model.conv_last_u", + "comp_model.scaleU", "comp_model.hidden_layer_UV") + clv = ui["comp_model.conv_last_v.weight"] + out = {"luma": luma, "chroma": chroma, + "conv_last_v": _conv_node(clv[0], clv[2], ui["comp_model.conv_last_v.bias"][2]), + "scale_v": _decode(ui["comp_model.scaleV"][0], 1, ui["comp_model.scaleV"][2])[0]} + json.dump(out, open(dst, "w")) + print(f"eICCI: luma {len(luma['blocks'])} blocks, chroma {len(chroma['blocks'])} blocks; " + f"scaleY={luma['scale']:.6f} scaleU={chroma['scale']:.6f} " + f"scaleV={out['scale_v']:.6f} -> {dst}") + + def main(argv): if len(argv) < 3: print(__doc__) @@ -216,8 +271,10 @@ def main(argv): export_ptq(src, argv[3]) elif mode == "int": export_int(src, argv[3]) + elif mode == "icci": + export_icci(argv[2], argv[3], argv[4]) # icci else: - print(f"unknown mode {mode!r}; expected list | ptq | int") + print(f"unknown mode {mode!r}; expected list | ptq | int | icci") return 1 return 0 diff --git a/wai/wai-rs/src/jpegai_decode.rs b/wai/wai-rs/src/jpegai_decode.rs new file mode 100644 index 000000000..c172b09bc --- /dev/null +++ b/wai/wai-rs/src/jpegai_decode.rs @@ -0,0 +1,230 @@ +//! The whole JPEG-AI decoder as one call: `decode(codestream, weights, norm) → Rgb`. +//! +//! This composes the stage functions validated across `jpegai_*` — the entropy front +//! half (`jpegai_mans`, `jpegai_pipeline`), the integer value path and synthesis +//! (`jpegai_synthesis`), and the post-filter chain (`jpegai_pipeline`, `jpegai_icci`) — +//! into a single entry point. It is pure compute: the caller supplies the model weights +//! and normative tables (see [`DecoderWeights`] / [`NormTables`]); this module reads no +//! files and does no I/O. +//! +//! Determinism, not float-exactness. ψ, the four MCM stages and synthesis are 16-bit +//! PTQ, so the decoded pixels differ from the float reference by ~1 LSB on a few percent +//! of samples. Everything is deterministic — integer arithmetic, or IEEE-754 float with +//! a fixed accumulation order — so the output is byte-identical on every machine, which +//! is the contract. See `tests/jpegai_e2e_full.rs`, which drives this and compares the +//! reference decode of `test256.bin` (max|Δ| = 1, ~4.4% of samples). +//! +//! Scope: the 256×256, 4:4:4, tools_0 profile the reference ships — luma 160 channels +//! (4 MCM stages), chroma 96 (no MCM). Other geometries are rejected up front. + +use crate::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, + parse_tools_header, Marker}; +use crate::jpegai_icci::IcciNet; +use crate::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder}; +use crate::jpegai_pipeline::{dequantize_resi, efl_chroma, efn_filter, gain_unit_scaler_vec, + lef_filter, lsbs_refine, reconstruct_yhat, reconstruct_yhat_uv, rvs_additive, + rvs_likely, sigma_index, skip_mask}; +use crate::jpegai_synthesis::{decode_hyper, decode_hyper_scale, decode_uv_f32, decode_y_f32, + q16_from_symbols, IntNet, Weights}; + +/// The `z` symbol offset (de-bias) applied to the factorized hyper-latent. +const Z_OFFSET: i64 = 31; + +/// One plane's model weights: the fully-integer hyper-scale net, the PTQ hyper decoder, +/// and the PTQ synthesis net. +pub struct PlaneWeights { + pub hyper_scale: IntNet, + pub hyper_decoder: Weights, + pub synthesis: Weights, +} + +/// Every model weight the decode needs. Luma carries the four MCM stages; chroma has no +/// MCM (its hyper params are used directly as the means). `icci` is the whole eICCI net. +pub struct DecoderWeights { + pub luma: PlaneWeights, + pub chroma: PlaneWeights, + pub mcm_y: [Weights; 4], + pub icci: IcciNet, +} + +/// The normative, image-independent tables. All are shipped decoder assets, not derived +/// per image: the me-tANS PMF/bounds, the factorized-`z` CDFs, the per-channel gain +/// vectors (`gain_unit_mlog`), the LSBS refinement tables, and the LEF op-point tables +/// selected at the picture's `model_id`. +pub struct NormTables { + pub pmf: Vec>, + pub bounds: Vec, + pub per_channel: usize, + pub cdfs_y: Vec, + pub cdfs_uv: Vec, + pub gain_y: Vec, + pub gain_uv: Vec, + pub lsbs_t0: Vec, + pub lsbs_t1: Vec, + pub lef_thr: [i64; 3], + pub lef_mag: [f32; 3], +} + +/// A decoded RGB image, row-major `[h][w][3]` (`data[(y*w + x)*3 + c]`). +pub struct Rgb { + pub w: usize, + pub h: usize, + pub data: Vec, +} + +/// Header gains are offset-binary Q11 around 32767. +fn deint(q: u32) -> f32 { (q as f32 - 32767.0) / 2048.0 } +fn q16(f: &[f32]) -> Vec { f.iter().map(|&x| (x as f64 * 65536.0).round() as i64).collect() } + +/// Decode one plane's front half and value path → its LSBS-refined synthesis latent +/// (f32), returning also `i_s` and the RVS additive for LEF's sigma surface. +/// +/// `z_pre_channels` is the channel count consumed from the SOZ substream *before* this +/// plane — one SOZ stream carries luma then chroma back to back through a single mANS +/// decoder, so chroma must skip the luma `z` to reach its own. +#[allow(clippy::too_many_arguments)] +fn decode_plane( + data: &[u8], marker: Marker, n: &NormTables, pw: &PlaneWeights, + cdfs: &[u8], gain: &[i64], mcm: Option<[&Weights; 4]>, + channels: usize, z_pre_channels: usize, +) -> (Vec, Vec, Vec) { + const H: usize = 16; // the latent plane is 16×16; z is at 4×4 and the nets upsample ×4 + const W: usize = 16; + let segs = parse_container(data); + + // SOZ → z (skip the prior plane's z first) + let soz = find_segment(&segs, Marker::Soz).expect("SOZ"); + let mem = &data[soz.start..soz.start + soz.len]; + let mut dz = MansDecoder::new(mem, &[mem.len() as u32]); + if z_pre_channels > 0 { + let mut skip = vec![0u8; z_pre_channels * n.per_channel]; + dz.decode_factorize(&n.cdfs_y, &mut skip, z_pre_channels, n.per_channel); + } + let mut syms = vec![0u8; channels * n.per_channel]; + dz.decode_factorize(cdfs, &mut syms, channels, n.per_channel); + let z: Vec = syms.iter().map(|&s| s as i64 - Z_OFFSET).collect(); + + // z → i_s → derived entropy inputs + let (i_s, ..) = decode_hyper_scale(&z, channels, 4, 4, &pw.hyper_scale); + let likely = rvs_likely(&i_s, gain, channels, H, W); + let grfs = grfs_flags(data, channels, marker); + let rvs = rvs_additive(&i_s, gain, &grfs, channels, H, W); + let sigma = sigma_index(&i_s, gain, &rvs, channels); + let mask = skip_mask(&i_s, gain, channels); + + // entropy decode → y_q + let seg = find_segment(&segs, marker).expect("SORP/SORS"); + let mem = &data[seg.start..seg.start + seg.len]; + let flat: Vec = build_decode_transitions(&n.pmf, &cdf_from_pmf(&n.pmf)) + .iter().flatten().copied().collect(); + let mut dec = MansDecoder::new(mem, &[mem.len() as u32]); + dec.set_sgm_transitions(&flat, &n.bounds); + let plane = H * W; + let mut y_q = vec![0i16; channels * plane]; + for c in 0..channels { + let mut row = vec![0i16; plane]; + dec.decode_sgm(&sigma[c * plane..(c + 1) * plane], &mut row, &mask[c * plane..(c + 1) * plane]); + y_q[c * plane..(c + 1) * plane].copy_from_slice(&row); + } + + // dequant → resi ; ψ ; reconstruct → ŷ ; LSBS → latent + let sv = gain_unit_scaler_vec(gain, 0); + let resi = dequantize_resi(&y_q, &likely, &grfs, &sv, channels); + let (psi, ..) = decode_hyper(&q16_from_symbols(&z), channels, 4, 4, &pw.hyper_decoder); + let resi_q = q16(&resi); + let yhat = match mcm { + Some(m) => reconstruct_yhat(&resi_q, &psi, channels, H, W, m), + None => reconstruct_yhat_uv(&resi_q, &psi, channels, H, W), + }; + let yhat_f: Vec = yhat.iter().map(|&v| v as f64 as f32 / 65536.0).collect(); + let lat = lsbs_refine(&yhat_f, &resi, &likely, &n.lsbs_t0, &n.lsbs_t1, channels, H, W); + (lat, i_s, rvs) +} + +fn grfs_flags(data: &[u8], channels: usize, marker: Marker) -> Vec { + let segs = parse_container(data); + let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih).unwrap().start_bit()); + let comp = if marker == Marker::Sorp { &ph.components[0] } else { &ph.components[1] }; + let mut f = comp.grfs_channel_flags.clone(); + f.resize(channels, false); + f +} + +/// Decode a JPEG-AI codestream to RGB. Returns `None` if the picture is not the +/// supported 256×256 4:4:4 tools_0 geometry. +pub fn decode(data: &[u8], w: &DecoderWeights, n: &NormTables) -> Option { + let segs = parse_container(data); + let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih)?.start_bit()); + let (img_h, img_w) = (ph.img_height as usize, ph.img_width as usize); + if (img_h, img_w) != (256, 256) || ph.components.len() != 2 + || ph.components[0].num_chs != 160 || ph.components[1].num_chs != 96 { + return None; + } + let (h, w_) = (img_h, img_w); + + // ---- luma + chroma latents → pixel planes (float, pre-round) ---- + let mcm = [&w.mcm_y[0], &w.mcm_y[1], &w.mcm_y[2], &w.mcm_y[3]]; + let (y_lat, y_is, y_rvs) = decode_plane(data, Marker::Sorp, n, &w.luma, + &n.cdfs_y, &n.gain_y, Some(mcm), 160, 0); + let (uv_lat, ..) = decode_plane(data, Marker::Sors, n, &w.chroma, + &n.cdfs_uv, &n.gain_uv, None, 96, 160); + + let mut yp = decode_y_f32(&q16(&y_lat), &w.luma.synthesis); + let mut cat = q16(&y_lat); + cat.extend(q16(&uv_lat)); + let uv = decode_uv_f32(&cat, &w.chroma.synthesis); + let (mut up, mut vp) = (uv[..h * w_].to_vec(), uv[h * w_..].to_vec()); + + // ---- post-filters, parameters from the tools header ---- + let ton = find_segment(&segs, Marker::Ton)?; + let th = parse_tools_header(data, ton.start_bit() + 2, (ph.c_ver, ph.c_hor) == (1, 1), 3, 1); + + // EFElinear — `filters_best2` is the applied group (the naming is inverted). + if let Some(efl) = th.efl.as_ref() { + let f = &efl.filters_best2; + let (mu, mv) = (efl.b1[0] as f32 / 100.0, efl.b1[1] as f32 / 100.0); + up = efl_chroma(&up, &yp, deint(f.weights_u[0]) + 1.0, deint(f.weights_u2[0]), mu); + vp = efl_chroma(&vp, &yp, deint(f.weights_v[0]) + 1.0, deint(f.weights_v2[0]), mv); + } + // eICCI — [0,1] domain: normalise in, scale back out. + if th.icci_enabled { + let nrm = |p: &[f32]| p.iter().map(|&x| x / 255.0).collect::>(); + let (yo, uo, vo) = w.icci.forward(&nrm(&yp), &nrm(&up), &nrm(&vp), h, w_); + let den = |p: Vec| p.into_iter().map(|x| x * 255.0).collect::>(); + yp = den(yo); up = den(uo); vp = den(vo); + } + // EFEnonlinear — per the header, applied to the enabled chroma plane(s). + if let Some(efn) = th.efn.as_ref() { + if efn.nonlinear_v_enabled { + let wq: Vec = efn.weights_v.iter().map(|&x| x as i64).collect(); + vp = efn_filter(&yp, &vp, &wq, efn.luma_min[0] as i64, efn.luma_max[0] as i64); + } + if efn.nonlinear_u_enabled { + let wq: Vec = efn.weights_u.iter().map(|&x| x as i64).collect(); + up = efn_filter(&yp, &up, &wq, efn.luma_min[0] as i64, efn.luma_max[0] as i64); + } + } + // LEF — luma; sigma is the raw scale surface (i_s + gain + rvs) at the named channel. + if th.lef_enabled { + let ch = th.lef_ch_idx as usize; + let sigma: Vec = (0..16 * 16) + .map(|p| (y_is[ch * 256 + p] + n.gain_y[ch] + y_rvs[ch * 256 + p]) as f32).collect(); + let thr = [n.lef_thr[0] as f32, n.lef_thr[1] as f32, n.lef_thr[2] as f32]; + yp = lef_filter(&yp, &sigma, 16, 16, thr, n.lef_mag, h, w_); + } + + // ---- BT.709 (full-range) → RGB ---- + let (kr, kg, kb) = (0.2126f32, 0.7152f32, 0.0722f32); + let (kby, kry) = (1.8556f32, 1.5748f32); + let mut out = vec![0u8; 3 * h * w_]; + for i in 0..h * w_ { + let (y01, u01, v01) = (yp[i] / 255.0, up[i] / 255.0 - 0.5, vp[i] / 255.0 - 0.5); + let rgb = [y01 + kry * v01, + y01 - (kb * kby / kg) * u01 - (kr * kry / kg) * v01, + y01 + kby * u01]; + for (c, &val) in rgb.iter().enumerate() { + out[i * 3 + c] = (val * 255.0).round().clamp(0.0, 255.0) as u8; + } + } + Some(Rgb { w: w_, h, data: out }) +} diff --git a/wai/wai-rs/src/jpegai_icci.rs b/wai/wai-rs/src/jpegai_icci.rs index 431209ac2..642d2fa7f 100644 --- a/wai/wai-rs/src/jpegai_icci.rs +++ b/wai/wai-rs/src/jpegai_icci.rs @@ -260,11 +260,16 @@ impl IcciPath { 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. + /// chroma trunk: `conv_firstUV` (BN pre-folded) + its residual blocks; `conv_last` + /// is the U head and `chroma.scale` is the **U** residual scale (`scaleU`). pub chroma: IcciPath, /// the V head — shares `chroma`'s trunk, differs only in this final convolution pub conv_last_v: Conv, + /// the V residual scale (`scaleV`). The reference signals the U and V scales + /// **separately** (`comp_model.scaleU` / `scaleV`); on the shipped models they + /// differ by ~1%, so reusing one scale for both heads shifts V by more than a + /// pixel LSB. + pub scale_v: f32, } impl IcciNet { @@ -282,13 +287,13 @@ impl IcciNet { for blk in &self.chroma.blocks { fea = blk.forward(&fea, h2, w2); } - let head = |conv: &Conv, band: &[f32]| -> Vec { + let head = |conv: &Conv, band: &[f32], scale: 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(); + let refined: Vec = res.iter().zip(band).map(|(&r, &s)| r * 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); + let u_o = head(&self.chroma.conv_last, &u16, self.chroma.scale); + let v_o = head(&self.conv_last_v, &v16, self.scale_v); (y_o, u_o, v_o) } } diff --git a/wai/wai-rs/src/jpegai_mans.rs b/wai/wai-rs/src/jpegai_mans.rs index ed4050c23..8c1f36f84 100644 --- a/wai/wai-rs/src/jpegai_mans.rs +++ b/wai/wai-rs/src/jpegai_mans.rs @@ -237,3 +237,83 @@ impl MansDecoder { } } } + +/// The reference's symbol ordering for a distribution of `len` bins: alternating +/// outward from the centre (`mid, mid-1, mid+1, mid-2, …`), which is what pairs each +/// CDF slot with its signed residual value. +fn sequence(len: usize) -> Vec { + let mid = len / 2; + let mut out = vec![0i64; len]; + let (mut even, mut odd) = (mid as i64, mid as i64 - 1); + for (i, v) in out.iter_mut().enumerate() { + if i % 2 == 0 { *v = even; even += 1; } else { *v = odd; odd -= 1; } + } + out +} + +/// Build the 256-entry `sgm` decode table per σ-distribution from the normative +/// integer `(pmf, cdf)` rows — the integer-only port of `utils.get_decode_transitions`. +/// +/// Each entry packs `(nBits << 24) | (next_base << 16) | (symbol & 0xffff)`, the layout +/// [`MansDecoder::decode_sgm`] consumes. Each CDF row is split into a first and second +/// half (`c - (c >> 1)` and `(c >> 1) + 128`) so the two interleaved ANS states draw +/// from disjoint slot ranges. +/// +/// Having this on the integer path means the entropy stage needs only the normative +/// PMF tables — no precomputed transition fixture, and no dependency on the +/// float/ONNX-side module. +pub fn build_decode_transitions(pmf: &[Vec], cdf: &[Vec]) -> Vec<[u32; 256]> { + const MASS_BITS: u32 = 8; + let mass = 1i64 << MASS_BITS; + let width = pmf[0].len(); + let seq: Vec = sequence(width).iter().map(|&v| v - (width as i64) / 2).collect(); + // get_clz(512): the number of doublings that bring `state` into [256, 512). + let clz = |state: i64| -> i64 { + let (mut bits, mut s) = (0i64, state.max(1)); + while s < 256 { s <<= 1; bits += 1; } + bits + }; + (0..pmf.len()) + .map(|i| { + let mut symbols = [0i64; 256]; + let mut state_next = [0i64; 256]; + let cdf_first: Vec = cdf[i].iter().map(|&c| c - (c >> 1)).collect(); + let cdf_second: Vec = cdf[i].iter().map(|&c| (c >> 1) + 128).collect(); + let (mut col, mut cdf_curr) = (0usize, 0i64); + while cdf_curr < mass { + let p = pmf[i][col]; + let d_first = cdf_first[col + 1] - cdf_first[col]; + for (k, slot) in (cdf_first[col]..cdf_first[col + 1]).enumerate() { + symbols[slot as usize] = seq[col]; + state_next[slot as usize] = p + k as i64; + } + for (k, slot) in (cdf_second[col]..cdf_second[col + 1]).enumerate() { + symbols[slot as usize] = seq[col]; + state_next[slot as usize] = p + d_first + k as i64; + } + cdf_curr = cdf[i][col + 1]; + col += 1; + } + let mut table = [0u32; 256]; + for s in 0..256 { + let bits = clz(state_next[s]); + let renormed = (state_next[s] << bits) ^ mass; + table[s] = ((bits as u32) << 24) + | ((renormed as u32) << 16) + | ((symbols[s] as u16) as u32); + } + table + }) + .collect() +} + +/// Cumulative rows for [`build_decode_transitions`] from the normative PMF tables. +pub fn cdf_from_pmf(pmf: &[Vec]) -> Vec> { + pmf.iter() + .map(|row| { + let mut c = vec![0i64; row.len() + 1]; + for (i, &p) in row.iter().enumerate() { c[i + 1] = c[i] + p; } + c + }) + .collect() +} diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index c6fb35adf..177c3dda1 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -115,14 +115,170 @@ pub fn reconstruct_yhat(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usiz up_shuffle(&parts, c, h, w) } +/// **SGMM reconstruction for chroma** — the counterpart to [`reconstruct_yhat`]. +/// +/// ```text +/// mu = up_shuffle(chunk(psi, 4)) // the hyper params ARE the means +/// y_hat = resi + mu +/// ``` +/// +/// Chroma has **no MCM**: the reference ships no `MCM/` nets under `model_uv`, and the +/// per-stage hyper params are used directly as the means rather than being refined by +/// an autoregressive context model. So this is a plain up-shuffle and add — no context +/// accumulation, no stage ordering, and nothing to get wrong in the way the luma path +/// can. Keeping it a separate named function rather than a flag on [`reconstruct_yhat`] +/// is deliberate: the two paths are structurally different, not parameterised. +/// +/// `resi` is `[c, 2h, 2w]` (Q16); `psi` is `[4c, h, w]` (Q16). Returns `[c, 2h, 2w]`. +pub fn reconstruct_yhat_uv(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usize) -> Vec { + let (h, w) = (h2 / 2, w2 / 2); + let plane = c * h * w; + assert_eq!(resi.len(), c * h2 * w2, "resi must be [c, 2h, 2w]"); + assert_eq!(psi.len(), 4 * plane, "psi must be [4c, h, w]"); + let parts: [Vec; 4] = std::array::from_fn(|i| psi[i * plane..(i + 1) * plane].to_vec()); + let mu = up_shuffle(&parts, c, h, w); + resi.iter().zip(&mu).map(|(&a, &b)| a + b).collect() +} + /// 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]`. +/// +/// Prefer [`dequantize_resi`], which derives the gain from normative assets and +/// applies the two stages in the reference's order. This entry point remains for +/// callers that already hold a gain map. 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() } +/// Normative RVS scale list, before the per-channel GRFS branch factor. +const RVS_BASE_SCALES: [i64; 4] = [115, 125, 146, 160]; + +/// `log_k` for the gain unit — note this is **54.82**, whereas the residual-variance +/// scale's own `log_k` uses 54.85. The two constants differ in the reference and are +/// easy to transpose; the additive table below is built with the 54.82 one. +const GAIN_LOG_K: f64 = 0.200_365_493_161_539_4; // (ln 54.82 − ln 0.11) / 31 + +/// Per-branch RVS scales: `round(base · 1.06)` when the channel's GRFS flag is set, +/// `round(base · 0.98)` otherwise. +/// +/// The rounding is **round-half-to-even**, not half-away-from-zero: `125 · 1.06` is +/// exactly `132.5` and `125 · 0.98` is exactly `122.5` in binary64, and the reference +/// (NumPy/Python `round`) resolves both to even — `132` and `122`. Rust's `f64::round` +/// would give `133`/`123` and silently select the wrong scale for one band. +fn rvs_scales(grfs_flag: bool) -> [i64; 4] { + let mult = if grfs_flag { 1.06 } else { 0.98 }; + std::array::from_fn(|i| (RVS_BASE_SCALES[i] as f64 * mult).round_ties_even() as i64) +} + +/// Per-band RVS **dequant** multipliers in Q16, `round(2^7 · 2^16 / scale)`. +/// Applied as `resi = y_q · table[band] / 2^16` (see [`dequantize_resi`]). +fn rvs_dequant_table(grfs_flag: bool) -> [i64; 4] { + let s = rvs_scales(grfs_flag); + std::array::from_fn(|i| ((1i64 << 23) as f64 / s[i] as f64).round_ties_even() as i64) +} + +/// Per-band RVS log-domain additive, `round(ln(scale/128) / (log_k / 2^7))`. +/// This reproduces [`RVS_ADDITIVE`]; the constant table is kept for decode-time cost +/// and cross-checked against this derivation in the calibration tests. +fn rvs_additive_table(grfs_flag: bool) -> [i64; 4] { + let s = rvs_scales(grfs_flag); + std::array::from_fn(|i| { + ((s[i] as f64 / 128.0).ln() / (GAIN_LOG_K / 128.0)).round_ties_even() as i64 + }) +} + +/// Gain-unit scale factor for one channel. +/// +/// ```text +/// scaler_vec[c] = round( exp( (gain_vector_log[c] + beta_displacement_log) · log_k / 2^7 ) · 2^10 ) / 2^10 +/// ``` +/// +/// `gain_vector_log` is the normative per-channel `gain_unit_mlog.csv` (the same +/// vector this module calls `gain_per_ch`); `beta_displacement_log` comes from the +/// picture header. `scaler_precision` is `gain_vector_precision + beta_displacement_precision` +/// = `5 + 5` = 10. +pub fn gain_unit_scaler_vec(gain_vector_log: &[i64], beta_displacement_log: i64) -> Vec { + const SCALER_PRECISION: i32 = 10; + let m = (1i64 << SCALER_PRECISION) as f64; + gain_vector_log + .iter() + .map(|&g| (((g + beta_displacement_log) as f64 * GAIN_LOG_K / 128.0).exp() * m).round() / m) + .collect() +} + +/// **Dequantize** entropy-decoded latent symbols — the full two-stage path, deriving +/// every factor from normative assets (no fixture gain map). +/// +/// ```text +/// x = y_q · rvs_dequant_table[grfs_flag[c]][band(likely[c,p])] / 2^16 // residual-variance scale +/// resi = x / (scaler_vec[c] + 1e-9) // gain unit +/// ``` +/// +/// The two stages are applied **sequentially, in this order**, matching the +/// reference's reverse iteration over enabled quantization tools. Folding them into a +/// single per-position gain first and multiplying once is algebraically identical but +/// differs in the last bit of the f32 result, so it is not done here. +/// +/// A note on the reference's own branch: `ResVarScale::dequantize_resi` contains two +/// paths — an integer `(x · t) >> 16` and a float `(x · t) / 2^16` — selected by a +/// *device comparison*, not by the bitstream. They disagree for negative `x` (floor +/// vs. exact). The float path is what the reference executes here, and is what this +/// implementation reproduces. +/// +/// `y_q` and `likely` are row-major `[channels × positions]`; `grfs_channel_flags` and +/// `scaler_vec` are per channel. +pub fn dequantize_resi( + y_q: &[i16], + likely: &[i64], + grfs_channel_flags: &[bool], + scaler_vec: &[f64], + channels: usize, +) -> Vec { + assert_eq!(scaler_vec.len(), channels, "scaler_vec length mismatch"); + let plane = y_q.len() / channels; + let mut out = rvs_dequant(y_q, likely, grfs_channel_flags, channels); + for c in 0..channels { + let inv = 1.0 / (scaler_vec[c] + 1e-9); + for p in 0..plane { + out[c * plane + p] = (out[c * plane + p] as f64 * inv) as f32; + } + } + out +} + +/// The residual-variance-scale stage of the dequant, on its own: +/// `out = y_q · rvs_dequant_table[grfs_flag[c]][band(likely)] / 2^16`. +pub fn rvs_dequant( + y_q: &[i16], + likely: &[i64], + grfs_channel_flags: &[bool], + channels: usize, +) -> Vec { + assert_eq!(y_q.len(), likely.len(), "y_q / likely length mismatch"); + assert_eq!(grfs_channel_flags.len(), channels, "grfs flag count mismatch"); + let plane = y_q.len() / channels; + + let tables = [rvs_dequant_table(false), rvs_dequant_table(true)]; + let mut out = vec![0f32; y_q.len()]; + for c in 0..channels { + let table = &tables[grfs_channel_flags[c] as usize]; + for p in 0..plane { + let i = c * plane + p; + let band = RVS_BAND_EDGES.iter().filter(|&&e| likely[i] >= e).count(); + out[i] = (y_q[i] as f64 * table[band] as f64 / 65536.0) as f32; + } + } + out +} + +/// The derived RVS tables, exposed so callers (and the calibration tests) can check +/// them against the reference's own `buildTables()` output. +pub fn rvs_tables(grfs_flag: bool) -> ([i64; 4], [i64; 4], [i64; 4]) { + (rvs_scales(grfs_flag), rvs_dequant_table(grfs_flag), rvs_additive_table(grfs_flag)) +} + /// Per-position mANS distribution index (`sigma`) for the `sgm` entropy decode: /// /// ```text @@ -184,8 +340,8 @@ pub fn skip_mask(i_s: &[i64], gain_per_ch: &[i64], channels: usize) -> Vec /// [`sigma_index`] folds in alongside the per-channel gain. /// /// ```text -/// likely[c,blk] = clamp( (Σ_{8×8} |i_s| + 32) >> 6, 0, 4351 ) // per 8×8 block -/// band = #{ e ∈ RVS_BAND_EDGES : likely > e } +/// likely[c,blk] = clamp( (Σ_{8×8} (i_s + gain_per_ch[c]) + 32) >> 6, 0, 4351 ) +/// band = #{ e ∈ RVS_BAND_EDGES : likely >= e } /// rvs[c,p] = RVS_ADDITIVE[grfs_channel_flag[c] ? 0 : 1][band] /// ``` /// @@ -197,15 +353,52 @@ pub fn skip_mask(i_s: &[i64], gain_per_ch: &[i64], channels: usize) -> Vec /// (`ComponentHeader::grfs_channel_flags`) and selects the branch. /// /// `i_s` is row-major `[channels × h × w]`; returns the same shape. -pub fn rvs_additive(i_s: &[i64], grfs_channel_flags: &[bool], channels: usize, h: usize, w: usize) -> Vec { +pub fn rvs_additive( + i_s: &[i64], + gain_per_ch: &[i64], + grfs_channel_flags: &[bool], + channels: usize, + h: usize, + w: usize, +) -> Vec { + assert_eq!(grfs_channel_flags.len(), channels, "one GRFS flag per channel"); + let likely = rvs_likely(i_s, gain_per_ch, channels, h, w); + likely + .iter() + .enumerate() + .map(|(i, &lk)| { + let branch = if grfs_channel_flags[i / (h * w)] { 0 } else { 1 }; + let band = RVS_BAND_EDGES.iter().filter(|&&e| lk >= e).count(); + RVS_ADDITIVE[branch][band] + }) + .collect() +} + +/// The per-8×8-block scale statistic that selects the RVS band, expanded to one value +/// per position: +/// +/// ```text +/// likely[c,blk] = clamp( (Σ_{8×8} (i_s + gain_per_ch[c]) + 32) >> 6, 0, 4351 ) +/// ``` +/// +/// The summed quantity is the **gain-adjusted** scale, not `i_s` alone and not its +/// magnitude. That distinction is invisible in the RVS additive on a single image — +/// the differing sums can still land in the same band — but it is not invisible in the +/// dequant, which indexes its Q16 multiplier table by this value directly and is +/// bit-exact only with the gain included. +/// +/// This one derivation drives both [`rvs_additive`] and [`dequantize_resi`], so the two +/// cannot drift apart. +/// +/// `i_s` is row-major `[channels × h × w]`; returns the same shape. +pub fn rvs_likely(i_s: &[i64], gain_per_ch: &[i64], channels: usize, h: usize, w: usize) -> Vec { const BLOCK: usize = 8; assert_eq!(i_s.len(), channels * h * w, "i_s must be [channels, h, w]"); - assert_eq!(grfs_channel_flags.len(), channels, "one GRFS flag per channel"); + assert_eq!(gain_per_ch.len(), channels, "one gain per channel"); let plane = h * w; let mut out = vec![0i64; i_s.len()]; for c in 0..channels { - let base = c * plane; - let branch = if grfs_channel_flags[c] { 0 } else { 1 }; + let (base, gain) = (c * plane, gain_per_ch[c]); let mut by = 0; while by < h { let mut bx = 0; @@ -214,16 +407,14 @@ pub fn rvs_additive(i_s: &[i64], grfs_channel_flags: &[bool], channels: usize, h for dy in 0..BLOCK { for dx in 0..BLOCK { let (y, x) = (by + dy, bx + dx); - if y < h && x < w { sum += i_s[base + y * w + x].abs(); } + if y < h && x < w { sum += i_s[base + y * w + x] + gain; } } } let likely = ((sum + 32) >> 6).clamp(0, 4351); - let band = RVS_BAND_EDGES.iter().filter(|&&e| likely >= e).count(); - let value = RVS_ADDITIVE[branch][band]; for dy in 0..BLOCK { for dx in 0..BLOCK { let (y, x) = (by + dy, bx + dx); - if y < h && x < w { out[base + y * w + x] = value; } + if y < h && x < w { out[base + y * w + x] = likely; } } } bx += BLOCK; @@ -409,47 +600,38 @@ pub fn efn_filter(y: &[f32], c_in: &[f32], weights_q15: &[i64], luma_min_q: i64, }).collect() } -pub fn lsbs_refine(yhat: &[f32], resi: &[f32], i_s: &[i64], t0: &[i32], t1: &[i32], +/// **LSBS** — least-significant-bit scale refinement, the last step before synthesis: +/// +/// ```text +/// means = yhat − resi +/// y_lsbs = yhat + (t0[likely]·means + t1[likely]·resi) / 2^13 +/// ``` +/// +/// `likely` is the shared per-8×8-block scale statistic from [`rvs_likely`] — the same +/// values the dequant indexes, verified identical against the reference on real data. +/// +/// It is taken as a parameter rather than recomputed here. This function previously +/// derived it from a slice it called `i_s`, summing **magnitudes**; both were wrong. +/// The summed quantity is the signed gain-adjusted scale (`i_s + gain_per_ch`), which +/// genuinely goes negative — on real luma data, summing magnitudes instead reproduces +/// only 1536/40960 of the reference's block statistics. Unlike the RVS additive, which +/// buckets `likely` into one of four bands and so tolerated the error, LSBS indexes a +/// 4352-entry table with it directly, so a wrong statistic reads the wrong coefficients. +/// Threading the one derivation through every consumer is what keeps them consistent. +pub fn lsbs_refine(yhat: &[f32], resi: &[f32], likely: &[i64], t0: &[i32], t1: &[i32], channels: usize, h: usize, w: usize) -> Vec { - 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); + assert_eq!(likely.len(), channels * plane, "likely must be per-position (block-constant)"); 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; - } + for i in 0..yhat.len() { + let lk = likely[i].clamp(0, last) as usize; + let (a0, a1) = (t0[lk] as f32, t1[lk] as f32); + let means = yhat[i] - resi[i]; + out[i] = yhat[i] + (a0 * means + a1 * resi[i]) / DENOM; } out } diff --git a/wai/wai-rs/src/jpegai_synthesis.rs b/wai/wai-rs/src/jpegai_synthesis.rs index fad626c66..6f25656c9 100644 --- a/wai/wai-rs/src/jpegai_synthesis.rs +++ b/wai/wai-rs/src/jpegai_synthesis.rs @@ -57,6 +57,12 @@ impl Weights { } Weights { conv, q16 } } + /// Output-channel count of a conv, read from its weight dims. Lets the nets be + /// driven by their own widths instead of a hardcoded plane size. + pub fn out_channels(&self, module: &str) -> usize { + self.conv.get(&format!("{module}.weight")) + .unwrap_or_else(|| panic!("missing conv {module}")).1[0] + } fn bias(&self, module: &str, cout: usize) -> Vec { self.q16.get(&format!("{module}.bias")).cloned().unwrap_or_else(|| vec![0; cout]) } @@ -167,7 +173,34 @@ fn tab(p: &str, x: &[i64], c: usize, heads: usize, h: usize, w: usize, wt: &Weig /// Decode the Y plane [256×256] from the 160-channel 16×16 latent. Pure-integer, /// deterministic; reproduces the float synthesis at ~58 dB PSNR (16-bit PTQ). +/// Output denormalize, shared by the u8 and float decoders: `clip(v·128 + 127.5, 0, 255)` +/// in Q16. The reference post-filters operate on this **float** value (0..255) and round +/// to u8 only after BT.709, so [`decode_y`]/[`decode_uv`] round here for a standalone +/// plane while [`decode_y_f32`]/[`decode_uv_f32`] keep the float for the filter chain. +fn denorm_q16(v: i64) -> i64 { + let half = S / 2; + (v * 128 + 127 * S + half).clamp(0, 255 * S) +} +fn denorm_u8(planes: &[i64]) -> Vec { + let half = S / 2; + planes.iter().map(|&v| ((denorm_q16(v) + half) >> F).clamp(0, 255) as u8).collect() +} +fn denorm_f32(planes: &[i64]) -> Vec { + planes.iter().map(|&v| denorm_q16(v) as f32 / S as f32).collect() +} + pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec { + denorm_u8(&decode_y_planes(latent, wt)) +} +/// The luma plane as **float** samples in `[0,255]` (pre-round) — the domain the +/// post-filter chain runs in. +pub fn decode_y_f32(latent: &[i64], wt: &Weights) -> Vec { + denorm_f32(&decode_y_planes(latent, wt)) +} + +/// The luma synthesis up to (but not including) the output denormalize: the raw +/// `conv4_t` output `[1,256,256]` in Q16. +fn decode_y_planes(latent: &[i64], wt: &Weights) -> Vec { let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5 // first_stage.0.0 residual block: conv1→ReLU→conv2 + latent [160,16,16] let (mut x, mut h, mut w) = (latent.to_vec(), 16usize, 16usize); @@ -220,12 +253,7 @@ pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec { // conv4_t ⇑2 128→1 → [1,256,256] let (y, yh, yw) = wt.convt("conv4_t", &x, 128, h, w, 3, 2, 1, 1); debug_assert_eq!((yh, yw), (256, 256)); - // output denormalize: clip(conv4_t·128 + 127.5, 0, 255) → u8 (nodes 603/605/608) - let half = S / 2; - y.iter().map(|&v| { - let d = (v * 128 + 127 * S + half).clamp(0, 255 * S); - ((d + half) >> F).clamp(0, 255) as u8 - }).collect() + y // pre-denorm [1,256,256]; see denorm_u8 / denorm_f32 } /// Decode the **U,V chroma planes** [2×256×256] from the 256-channel 16×16 latent @@ -241,6 +269,17 @@ pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec { /// • 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 { + denorm_u8(&decode_uv_planes(latent, wt)) +} +/// The two chroma planes as **float** samples in `[0,255]` (pre-round) — the domain the +/// post-filter chain runs in. +pub fn decode_uv_f32(latent: &[i64], wt: &Weights) -> Vec { + denorm_f32(&decode_uv_planes(latent, wt)) +} + +/// The chroma synthesis up to (but not including) the output denormalize: the raw +/// depth-to-space output `[2,256,256]` in Q16. +fn decode_uv_planes(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; @@ -288,12 +327,7 @@ pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec { 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() + yd // pre-denorm [2,256,256]; see denorm_u8 / denorm_f32 } /// **hyper_decoder** (common_modules) — the entropy-model hyper-prior decoder that @@ -302,13 +336,31 @@ pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec { /// `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. +/// +/// `z` **must be Q16**. Entropy-decoded symbols are small plain integers (≈ −3..4), +/// so passing them unscaled is an easy mistake — and it does not fail loudly: the +/// first conv comes out 2^16 too small, ReLU6 clamps it to zero, and the result +/// collapses to `conv4`'s bias, which still *looks* like a plausible surface. Use +/// [`q16_from_symbols`] at the boundary. Measured on real data: correct scaling gives +/// rmse 2.4e-05 against the float net, unscaled gives 0.074 — against a zero-baseline +/// of 0.079, i.e. barely better than emitting nothing. +/// Lift entropy-decoded integer symbols into the Q16 fixed-point the PTQ nets expect. +/// Named rather than inline so the scaling is visible at the call site — see the +/// failure mode described on [`decode_hyper`]. +pub fn q16_from_symbols(symbols: &[i64]) -> Vec { + symbols.iter().map(|&v| v * S).collect() +} + 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 + // Widths come from the weights: chroma is the same topology at 96 / 384 where luma + // is 160 / 640, so nothing here may assume the luma channel count. + let c = wt.out_channels("conv1"); + let (x, _, _) = wt.conv("conv1", z, cin, h, w, 1, 1, 0, 1); // c→c, 1×1 (no bias) + let (x, h, w) = wt.convt("conv2", &x, c, h, w, 4, 2, 1, 0); // c→c, 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, _, _) = wt.conv("conv3", &x, c, h, w, 3, 1, 1, 1); // c→c, 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] + wt.conv("conv4", &x, c, h, w, 3, 1, 1, 1) // c→4c, 3×3 → p[4c,2h,2w] } /// A fully-integer weight store (int32 conv weights/biases + integer per-channel @@ -397,10 +449,12 @@ pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNe } (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 + // Channel widths are threaded from each block's own weights, not assumed: the + // chroma net is the same topology at 96/1536 where luma is 160/2560. + let (x1, c1) = block(z, cin, "conv1", 1, 0, "onnx::Pow_85", true); + let (x2, c2) = block(&x1, c1, "depthwise", 3, 1, "onnx::Pow_91", true); + let (x3, c3) = block(&x2, c2, "pointwise", 1, 0, "onnx::Pow_97", false); + // DepthToSpace CRD block 4: c·16 → c 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(); diff --git a/wai/wai-rs/src/lib.rs b/wai/wai-rs/src/lib.rs index dbf5e4003..b6a84ff5b 100644 --- a/wai/wai-rs/src/lib.rs +++ b/wai/wai-rs/src/lib.rs @@ -92,6 +92,11 @@ pub mod jpegai_icci; #[cfg(feature = "neural_int")] pub mod jpegai_bitstream; +// JPEG-AI whole-decoder entry point: composes the entropy, value, synthesis and +// post-filter stages into `decode(codestream, weights, norm) -> Rgb`. +#[cfg(feature = "neural_int")] +pub mod jpegai_decode; + // 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_dequant_calibration.rs b/wai/wai-rs/tests/jpegai_dequant_calibration.rs new file mode 100644 index 000000000..16dc82543 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_dequant_calibration.rs @@ -0,0 +1,105 @@ +//! Calibration for the **dequant gain**, against the ISO reference on real +//! `test256.bin` data, for both the luma and chroma planes. +//! +//! This closes the last underived decoder input. Before this, the per-position gain +//! came from a dumped fixture; it is now built from normative assets only: +//! +//! ```text +//! scale[i] = round_ties_even(base[i] · (grfs_flag ? 1.06 : 0.98)) base = [115,125,146,160] +//! dequant[i] = round(2^7 · 2^16 / scale[i]) // RVS stage, Q16 +//! additive[i] = round(ln(scale[i]/128) / (log_k / 2^7)) // == RVS_ADDITIVE +//! scaler_vec[c] = round(exp((gain_vector_log[c] + beta) · log_k / 2^7) · 2^10) / 2^10 +//! ``` +//! +//! `gain_vector_log` is the normative `gain_unit_mlog.csv`; `beta_displacement_log` +//! comes from the picture header. +//! +//! Two rounding hazards are asserted here rather than left to chance: +//! * the scale rounding is **half-to-even** — `125·1.06` is exactly `132.5` and +//! `125·0.98` exactly `122.5`, and half-away-from-zero picks the wrong scale; +//! * the reference's `dequantize_resi` has an integer `>>16` path and a float +//! `/2^16` path chosen by a *device comparison*. They disagree for negative +//! input (floor vs. exact). The float path is the one it executes. +//! +//! env: WAI_JPEGAI_DEQUANT_FIXTURE = dequant_fixture.json +//! { y|uv: { channels, h, w, y_q, likely, grfs, gain_vector_log, +//! beta_displacement_log, rvs_out_expected } } +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::{gain_unit_scaler_vec, rvs_dequant, rvs_likely, rvs_tables}; + +/// The tables the reference's `buildTables()` produces, read off its own dumps. +#[test] +fn derived_tables_match_reference() { + // GRFS flag clear → base · 0.98; set → base · 1.06. + let (scales, dq, add) = rvs_tables(false); + assert_eq!(scales, [113, 122, 143, 157], "GRFS-clear scales"); + assert_eq!(dq, [74235, 68759, 58662, 53431], "GRFS-clear dequant table (cwg2)"); + assert_eq!(add, [-80, -31, 71, 130], "GRFS-clear additive == RVS_ADDITIVE[1]"); + + let (scales, dq, add) = rvs_tables(true); + assert_eq!(scales, [122, 132, 155, 170], "GRFS-set scales"); + assert_eq!(dq, [68759, 63550, 54120, 49345], "GRFS-set dequant table (cwg1)"); + assert_eq!(add, [-31, 20, 122, 181], "GRFS-set additive == RVS_ADDITIVE[0]"); + + // Guard the half-to-even hazard explicitly: half-away-from-zero would give + // 133 / 123 for the band-1 scale and silently mis-select its multiplier. + assert_eq!((125.0f64 * 1.06).round_ties_even() as i64, 132); + assert_eq!((125.0f64 * 0.98).round_ties_even() as i64, 122); +} + +#[test] +fn dequant_matches_reference() { + let path = 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 v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + + for plane in ["y", "uv"] { + let o = &v[plane]; + let channels = o["channels"].as_u64().unwrap() as usize; + let i64v = |k: &str| o[k].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap()).collect::>(); + + let (h, w) = (o["h"].as_u64().unwrap() as usize, o["w"].as_u64().unwrap() as usize); + let y_q: Vec = i64v("y_q").iter().map(|&x| x as i16).collect(); + let grfs: Vec = o["grfs"].as_array().unwrap().iter() + .map(|x| x.as_bool().unwrap()).collect(); + let gain_vector_log = i64v("gain_vector_log"); + let beta = o["beta_displacement_log"].as_i64().unwrap(); + let expected: Vec = o["rvs_out_expected"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + + // `likely` is DERIVED from the scale surface, not taken from the fixture: the + // summed quantity is the gain-adjusted scale (i_s + gain_per_ch), and the + // dequant is bit-exact only with the gain included. Summing |i_s| instead + // happens to select the same RVS band on this image but is not the same + // quantity — the fixture check below is what distinguishes them. + let i_s = i64v("i_s"); + let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w); + let lk_expected = i64v("likely_expected"); + let lk_exact = likely.iter().zip(&lk_expected).filter(|(a, b)| a == b).count(); + assert_eq!( + lk_exact, lk_expected.len(), + "{plane}: derived likely mismatched ({lk_exact}/{})", lk_expected.len() + ); + + // The RVS stage must be bit-exact against the reference — it is pure integer + // table lookup plus an exact power-of-two divide. + let got = rvs_dequant(&y_q, &likely, &grfs, channels); + let exact = got.iter().zip(&expected).filter(|(a, b)| a == b).count(); + assert_eq!( + exact, expected.len(), + "{plane}: RVS dequant not bit-exact ({exact}/{})", expected.len() + ); + + // The gain unit is well-formed and strictly positive for every channel, so + // the composed dequant never divides by ~0. + let sv = gain_unit_scaler_vec(&gain_vector_log, beta); + assert_eq!(sv.len(), channels, "{plane}: scaler_vec length"); + assert!(sv.iter().all(|&s| s > 0.0 && s.is_finite()), "{plane}: scaler_vec sane"); + + eprintln!("{plane}: RVS dequant bit-exact {exact}/{}", expected.len()); + } +} diff --git a/wai/wai-rs/tests/jpegai_derive_calibration.rs b/wai/wai-rs/tests/jpegai_derive_calibration.rs index 745eb8986..1cd0828b5 100644 --- a/wai/wai-rs/tests/jpegai_derive_calibration.rs +++ b/wai/wai-rs/tests/jpegai_derive_calibration.rs @@ -39,7 +39,7 @@ fn derived_inputs_match_reference() { .iter().map(|x| x.as_u64().unwrap() as u8).collect(); // 1. RVS additive - let rvs = rvs_additive(&i_s, &grfs, channels, h, w); + let rvs = rvs_additive(&i_s, &gain_per_ch, &grfs, channels, h, w); let bad = rvs.iter().zip(&rvs_expected).filter(|(a, b)| a != b).count(); eprintln!("[{plane}] rvs_additive: {}/{} exact", rvs.len() - bad, rvs.len()); assert_eq!(bad, 0, "{plane}: RVS additive diverged from the reference"); diff --git a/wai/wai-rs/tests/jpegai_e2e_chroma.rs b/wai/wai-rs/tests/jpegai_e2e_chroma.rs new file mode 100644 index 000000000..cfeb49d47 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_e2e_chroma.rs @@ -0,0 +1,239 @@ +//! End-to-end **chroma**: raw codestream → the decoded U and V planes, no per-image +//! fixture — the counterpart to `jpegai_e2e_frontend`'s luma chain. +//! +//! ```text +//! test256.bin → SOZ → mANS fac (luma then uv CDFs) → z_uv +//! → decode_hyper_scale → i_s +//! → rvs_likely / rvs_additive / skip → sigma, mask +//! → SORS → mANS sgm → y_q bit-exact +//! → dequantize_resi → resi bit-exact +//! → decode_hyper → ψ +//! → reconstruct_yhat_uv(resi, ψ) → y_hat +//! → lsbs_refine → uv_lsbs +//! → decode_uv(cat(y_lsbs, uv_lsbs)) → U,V[256²] +//! ``` +//! +//! Chroma is **not** luma with different constants — two structural differences make +//! this worth exercising separately rather than parameterising the luma test: +//! +//! * **No MCM.** The reference ships no `MCM/` nets under `model_uv`; the hyper +//! params are used directly as the means (see [`reconstruct_yhat_uv`]), where luma +//! runs four autoregressive context stages. +//! * **Synthesis is joint.** `decode_uv` consumes the *concatenated* 256-channel +//! latent `cat(y_lsbs, uv_lsbs)`, so chroma pixels depend on the luma chain being +//! right as well — a luma regression surfaces here too. +//! +//! One SOZ substream carries both planes' hyper-latents back to back through a single +//! decoder whose ANS state threads across them, so the luma pass is decoded first even +//! though only the chroma output is checked here. +//! +//! Everything upstream of synthesis is shared code with luma, which is the point: the +//! same derivations (`rvs_likely`, `skip_mask`, `sigma_index`, `dequantize_resi`) are +//! re-exercised on a second plane with different channel counts and flags. +//! +//! env: WAI_JPEGAI_FRONTEND_FIXTURE (uses its `uv` sub-object and the luma `y_lsbs` path) +//! WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS, WAI_JPEGAI_HYPER_UV_WEIGHTS, +//! WAI_JPEGAI_SYNTH_UV_WEIGHTS (optional; each stage is skipped without its weights) +#![cfg(feature = "neural_int")] + +use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker}; +use wai::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder}; +use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat_uv, + rvs_additive, rvs_likely, sigma_index, skip_mask}; +use wai::jpegai_synthesis::{decode_hyper, decode_hyper_scale, q16_from_symbols, IntNet, Weights}; + +const Z_OFFSET: i64 = 31; + +fn env_path(k: &str) -> Option { + std::env::var(k).ok().filter(|p| std::path::Path::new(p).exists()) +} +fn weights(p: &str) -> Weights { + Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()) +} +fn snr(got_q16: &[i64], expected: &[f32]) -> f64 { + let (mut se, mut sig) = (0f64, 0f64); + for (&g, &e) in got_q16.iter().zip(expected) { + let (q, f) = (g as f64 / 65536.0, e as f64); + se += (q - f) * (q - f); + sig += f * f; + } + if se == 0.0 { f64::INFINITY } else { 10.0 * (sig / se).log10() } +} + +#[test] +fn codestream_to_chroma_pixels() { + let (Some(fx), Some(hs_uv)) = (env_path("WAI_JPEGAI_FRONTEND_FIXTURE"), + env_path("WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS")) else { + eprintln!("skip: set WAI_JPEGAI_FRONTEND_FIXTURE + WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS"); + return; + }; + let root: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap(); + let v = &root["uv"]; + let u8v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter() + .map(|x| x.as_u64().unwrap() as u8).collect::>(); + let i64v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap()).collect::>(); + let f32v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect::>(); + + let data = u8v(&root, "codestream"); + let (channels, h, w) = ( + v["channels"].as_u64().unwrap() as usize, + v["h"].as_u64().unwrap() as usize, + v["w"].as_u64().unwrap() as usize, + ); + let zs: Vec = v["z_shape"].as_array().unwrap().iter() + .map(|x| x.as_u64().unwrap() as usize).collect(); + let (zc, zh, zw) = (zs[1], zs[2], zs[3]); + let per_channel = root["per_channel"].as_u64().unwrap() as usize; + let gain_vector_log = i64v(v, "gain_vector_log"); + + let segs = parse_container(&data); + + // grfs and beta are DERIVED from the codestream's picture header, not the fixture: + // component [1] is chroma. This is the same parser `jpegai_header_parser_matches_trace` + // validates against the reference trace, so nothing per-image is assumed here. + let pih = find_segment(&segs, Marker::Pih).expect("PIH"); + let hdr = parse_picture_header(&data, pih.start_bit()); + let grfs = hdr.components[1].grfs_channel_flags.clone(); + let beta = hdr.beta_displacement_log[1] as i64; + assert_eq!(grfs.len(), channels, "one GRFS flag per chroma channel"); + + // 1. SOZ → z_uv. One SOZ substream carries BOTH planes' hyper-latents back to + // back — luma first, then chroma — through a single decoder whose ANS state + // threads across them. The luma pass is therefore not optional scaffolding here: + // skipping it leaves the decoder at the wrong stream position and chroma decodes + // garbage (or underflows outright). + let soz = find_segment(&segs, Marker::Soz).expect("SOZ"); + let mem = &data[soz.start..soz.start + soz.len]; + let mut dec_z = MansDecoder::new(mem, &[mem.len() as u32]); + + let y_ch = root["z_shape"].as_array().unwrap()[1].as_u64().unwrap() as usize; + let mut y_syms = vec![0u8; y_ch * per_channel]; + dec_z.decode_factorize(&u8v(&root, "cdfs_y"), &mut y_syms, y_ch, per_channel); + + let mut symbols = vec![0u8; zc * per_channel]; + dec_z.decode_factorize(&u8v(v, "cdfs_uv"), &mut symbols, zc, per_channel); + let z: Vec = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect(); + + // 2. z → hyper-scale → i_s, then the derived entropy inputs (shared with luma) + let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&hs_uv).unwrap()).unwrap()); + let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net); + assert_eq!(i_s.len(), channels * h * w, "i_s shape"); + + let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w); + let rvs = rvs_additive(&i_s, &gain_vector_log, &grfs, channels, h, w); + let sigma = sigma_index(&i_s, &gain_vector_log, &rvs, channels); + let mask = skip_mask(&i_s, &gain_vector_log, channels); + eprintln!("[uv] derived: {} coded positions of {}", mask.iter().filter(|&&m| m).count(), mask.len()); + + // 3. SORS → mANS sgm → y_q + let sors = find_segment(&segs, Marker::Sors).expect("SORS"); + let mem = &data[sors.start..sors.start + sors.len]; + let pmf: Vec> = root["pmf"].as_array().unwrap().iter() + .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()) + .collect(); + let transitions = build_decode_transitions(&pmf, &cdf_from_pmf(&pmf)); + let flat: Vec = transitions.iter().flatten().copied().collect(); + + let mut dec = MansDecoder::new(mem, &[mem.len() as u32]); + dec.set_sgm_transitions(&flat, &u8v(&root, "bounds")); + let plane = h * w; + let mut y_q = vec![0i16; channels * plane]; + for c in 0..channels { + let (a, b) = (c * plane, (c + 1) * plane); + let mut row = vec![0i16; plane]; + dec.decode_sgm(&sigma[a..b], &mut row, &mask[a..b]); + y_q[a..b].copy_from_slice(&row); + } + let y_q_expected: Vec = i64v(v, "y_q_expected").iter().map(|&x| x as i16).collect(); + let bad = y_q.iter().zip(&y_q_expected).filter(|(a, b)| a != b).count(); + eprintln!("[uv] y_q from codestream: {}/{} exact", y_q.len() - bad, y_q.len()); + assert_eq!(bad, 0, "codestream → y_q (chroma) must be byte-exact"); + + // 4. dequant — bit-exact against the reference residual + let sv = gain_unit_scaler_vec(&gain_vector_log, beta); + let resi = dequantize_resi(&y_q, &likely, &grfs, &sv, channels); + let resi_exp = f32v(v, "resi_full_expected"); + let exact = resi.iter().zip(&resi_exp).filter(|(a, b)| a == b).count(); + eprintln!("[uv] resi from codestream: {exact}/{} bit-exact", resi_exp.len()); + assert_eq!(exact, resi_exp.len(), "chroma dequant must be bit-exact"); + + // 5. ψ → reconstruction (no MCM) → LSBS + let Some(hyp_uv) = env_path("WAI_JPEGAI_HYPER_UV_WEIGHTS") else { + eprintln!("[uv] skip ψ onward: set WAI_JPEGAI_HYPER_UV_WEIGHTS"); + return; + }; + let (psi, ..) = decode_hyper(&q16_from_symbols(&z), zc, zh, zw, &weights(&hyp_uv)); + eprintln!("[uv] ψ: SNR {:.2} dB", snr(&psi, &f32v(v, "psi_expected"))); + + let resi_q16: Vec = resi.iter().map(|&r| (r as f64 * 65536.0).round() as i64).collect(); + let yhat = reconstruct_yhat_uv(&resi_q16, &psi, channels, h, w); + let yhat_snr = snr(&yhat, &f32v(v, "yhat_expected")); + eprintln!("[uv] y_hat: SNR {yhat_snr:.2} dB"); + assert!(yhat_snr > 40.0, "chroma y_hat SNR {yhat_snr:.2} dB"); + + let t0: Vec = root["lsbs_t0"].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap() as i32).collect(); + let t1: Vec = root["lsbs_t1"].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap() as i32).collect(); + let yhat_f: Vec = yhat.iter().map(|&q| q as f64 as f32 / 65536.0).collect(); + let uv_lsbs = lsbs_refine(&yhat_f, &resi, &likely, &t0, &t1, channels, h, w); + let lsbs_exp = f32v(v, "lsbs_out_expected"); + let (mut se, mut sig, mut id_se) = (0f64, 0f64, 0f64); + for i in 0..lsbs_exp.len() { + let (g, e, id) = (uv_lsbs[i] as f64, lsbs_exp[i] as f64, yhat_f[i] as f64); + se += (g - e) * (g - e); + sig += e * e; + id_se += (id - e) * (id - e); + } + let (lsnr, id_snr) = (10.0 * (sig / se).log10(), 10.0 * (sig / id_se).log10()); + eprintln!("[uv] uv_lsbs: SNR {lsnr:.2} dB (identity baseline {id_snr:.2} dB)"); + assert!(lsnr > id_snr + 10.0, "chroma LSBS no better than identity"); + + // 6. Joint synthesis: cat(y_lsbs, uv_lsbs) → U,V. Needs the luma latent, so this + // also re-checks the luma chain end to end. + let (Some(sw_uv), Some(y_lsbs_j)) = (env_path("WAI_JPEGAI_SYNTH_UV_WEIGHTS"), + root.get("lsbs_out_expected")) else { + eprintln!("[uv] skip synthesis: set WAI_JPEGAI_SYNTH_UV_WEIGHTS"); + return; + }; + let mut lat: Vec = y_lsbs_j.as_array().unwrap().iter() + .map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect(); + lat.extend(uv_lsbs.iter().map(|&f| (f as f64 * 65536.0).round() as i64)); + let uv_pix = wai::jpegai_synthesis::decode_uv(&lat, &weights(&sw_uv)); + let exp = f32v(v, "synth_uv_expected"); + assert_eq!(uv_pix.len(), exp.len(), "chroma plane size"); + let (mut maxd, mut off) = (0i32, 0usize); + for (&g, &e) in uv_pix.iter().zip(&exp) { + let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs(); + maxd = maxd.max(d); + if d != 0 { off += 1; } + } + let total = exp.len(); + eprintln!("[uv] U,V planes: max|Δ| = {maxd}, {off}/{total} differ ({:.3}%)", + 100.0 * off as f64 / total as f64); + + // Attribute the pixel error rather than just widening the bar: rerun synthesis on + // the REFERENCE latent for both planes. What remains is the synthesis net's own + // 16-bit PTQ; the difference is everything the chain accumulated upstream. + let mut ref_lat: Vec = y_lsbs_j.as_array().unwrap().iter() + .map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect(); + ref_lat.extend(f32v(v, "lsbs_out_expected").iter().map(|&f| (f as f64 * 65536.0).round() as i64)); + let uv_ref = wai::jpegai_synthesis::decode_uv(&ref_lat, &weights(&sw_uv)); + let roff = uv_ref.iter().zip(&exp) + .filter(|(g, e)| **g as i32 != e.round().clamp(0.0, 255.0) as i32).count(); + eprintln!("[uv] from REFERENCE latent: {roff}/{total} differ ({:.3}%) \ + ⇒ synthesis-port {:.3}%, accumulated upstream {:.3}%", + 100.0 * roff as f64 / total as f64, + 100.0 * roff as f64 / total as f64, + 100.0 * (off as f64 - roff as f64) / total as f64); + + assert!(maxd <= 1, "chroma deviation beyond rounding: {maxd}"); + // Chroma runs looser than luma (~3.4% vs ~1.4% of samples at 1 LSB). Its ψ is + // 60.6 dB where luma's is 70.5, and with no MCM to refine them those hyper params + // are used directly as the means, so the quantization error reaches the latent + // unattenuated. The bar is set from the measured value with headroom, and the + // attribution above is what a regression would actually move. + assert!(off * 100 <= total * 4, "more than 4% mismatched: {off}/{total}"); +} diff --git a/wai/wai-rs/tests/jpegai_e2e_frontend.rs b/wai/wai-rs/tests/jpegai_e2e_frontend.rs new file mode 100644 index 000000000..967e789f5 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_e2e_frontend.rs @@ -0,0 +1,293 @@ +//! End-to-end: raw codestream → **decoded luma pixels**, in one pass with no per-image +//! fixture at all. +//! +//! ```text +//! test256.bin → parse_container → SOZ → mANS fac → z +//! → decode_hyper_scale(z) → i_s +//! → rvs_likely / rvs_additive / skip_mask → sigma, mask +//! → SORP → mANS sgm → y_q bit-exact +//! → dequantize_resi (RVS ∘ gain unit) → resi bit-exact +//! → decode_hyper(z) → ψ SNR ~70 dB +//! → reconstruct_yhat(resi, ψ) → y_hat SNR ~75 dB +//! → lsbs_refine → y_lsbs SNR ~75 dB +//! → decode_y → Y[256²] max|Δ| = 1 +//! ``` +//! +//! The integer stages are checked bit-exact; the PTQ nets (ψ, the four MCM stages +//! inside the reconstruction, and synthesis) are checked against the float reference by +//! signal-to-error, since bit-equality with a float net is not achievable by +//! construction. Where a stage is a small correction, the bar is set above its +//! do-nothing baseline — otherwise a no-op would pass. +//! +//! At the pixel end, every differing sample is off by exactly 1 (max|Δ| = 1). About +//! 1.4% of samples differ, which is looser than the reference decoder's own 0.1% bar +//! because that comparison runs float nets for ψ/MCM/synthesis while this path is +//! integer throughout. The test attributes the gap: feeding the reference latent into +//! synthesis alone already yields ~1.2%, so the 16-bit PTQ of the synthesis net is the +//! dominant term and the whole upstream chain contributes only ~0.2%. +//! +//! Every stage here was validated in isolation first; this runs them as one pipeline +//! off raw bytes. What makes it worth having separately from those unit calibrations +//! is that nothing per-image is supplied: the σ distribution indices, the coded-position +//! mask, the RVS band and the dequant multiplier are all *derived* from the bitstream, +//! so a break anywhere upstream desynchronizes the entropy decoder and the symbols come +//! out wrong. The mANS state threads across all 160 channel calls, which makes the +//! symbol check a stringent end-to-end assertion rather than 160 independent ones. +//! +//! Inputs are the codestream plus normative assets only — the factorized CDFs, the +//! me-tANS PMF/bounds tables, the hyper-scale net, and `gain_unit_mlog.csv`. The +//! transition tables are built here from the normative PMF via +//! [`build_decode_transitions`] rather than loaded from a fixture. +//! +//! env: WAI_JPEGAI_FRONTEND_FIXTURE = frontend_fixture.json +//! { codestream, cdfs_y, per_channel, z_shape, pmf, bounds, gain_vector_log, +//! grfs, channels, h, w, y_q_expected, rvs_out_expected, +//! psi_expected, resi_full_expected, yhat_expected } +//! WAI_JPEGAI_HYPERSCALE_WEIGHTS = hyperscale_y_weights.json +//! WAI_JPEGAI_HYPER_WEIGHTS = hyper_y_weights.json (optional; adds ψ) +//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir of mcm{0..3}_y.json (optional; adds y_hat) +//! WAI_JPEGAI_SYNTH_WEIGHTS = synth_y.json (optional; adds pixels) +#![cfg(feature = "neural_int")] + +use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker}; +use wai::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder}; +use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat, rvs_additive, + rvs_dequant, rvs_likely, sigma_index, skip_mask}; +use wai::jpegai_synthesis::{decode_hyper, decode_hyper_scale, decode_y, q16_from_symbols, IntNet, Weights}; + +const Z_OFFSET: i64 = 31; + +#[test] +fn codestream_to_luma_pixels() { + let fx = match std::env::var("WAI_JPEGAI_FRONTEND_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_FRONTEND_FIXTURE"); return; } + }; + let wj = 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"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap(); + let u8v = |k: &str| v[k].as_array().unwrap().iter() + .map(|x| x.as_u64().unwrap() as u8).collect::>(); + let i64v = |k: &str| v[k].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap()).collect::>(); + + let data = u8v("codestream"); + let (channels, h, w) = ( + v["channels"].as_u64().unwrap() as usize, + v["h"].as_u64().unwrap() as usize, + v["w"].as_u64().unwrap() as usize, + ); + let per_channel = v["per_channel"].as_u64().unwrap() as usize; + let zs: Vec = v["z_shape"].as_array().unwrap().iter() + .map(|x| x.as_u64().unwrap() as usize).collect(); + let (zc, zh, zw) = (zs[1], zs[2], zs[3]); + let gain_vector_log = i64v("gain_vector_log"); + + let segs = parse_container(&data); + + // grfs and beta are DERIVED from the codestream's picture header, not the fixture: + // component [0] is luma. Same parser `jpegai_header_parser_matches_trace` checks + // against the reference trace, so nothing per-image is assumed. + let pih = find_segment(&segs, Marker::Pih).expect("PIH"); + let hdr = parse_picture_header(&data, pih.start_bit()); + let grfs = hdr.components[0].grfs_channel_flags.clone(); + let beta = hdr.beta_displacement_log[0] as i64; + assert_eq!(grfs.len(), channels, "one GRFS flag per luma channel"); + + // 1. SOZ → z + let soz = find_segment(&segs, Marker::Soz).expect("SOZ"); + let mem = &data[soz.start..soz.start + soz.len]; + let mut dec_z = MansDecoder::new(mem, &[mem.len() as u32]); + let mut symbols = vec![0u8; zc * per_channel]; + dec_z.decode_factorize(&u8v("cdfs_y"), &mut symbols, zc, per_channel); + let z: Vec = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect(); + + // 2. z → hyper-scale net → i_s + let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap()); + let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net); + assert_eq!(i_s.len(), channels * h * w, "i_s shape"); + + // 3. i_s → the entropy decoder's own inputs, all derived + let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w); + let rvs = rvs_additive(&i_s, &gain_vector_log, &grfs, channels, h, w); + let sigma = sigma_index(&i_s, &gain_vector_log, &rvs, channels); + let mask = skip_mask(&i_s, &gain_vector_log, channels); + let coded = mask.iter().filter(|&&m| m).count(); + eprintln!("derived: {coded} coded positions of {}", mask.len()); + + // 4. SORP → mANS sgm → y_q. One decoder, state threading across every channel. + let sorp = find_segment(&segs, Marker::Sorp).expect("SORP"); + let mem = &data[sorp.start..sorp.start + sorp.len]; + let pmf: Vec> = v["pmf"].as_array().unwrap().iter() + .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()) + .collect(); + let transitions = build_decode_transitions(&pmf, &cdf_from_pmf(&pmf)); + let flat: Vec = transitions.iter().flatten().copied().collect(); + let bounds = u8v("bounds"); + + let mut dec = MansDecoder::new(mem, &[mem.len() as u32]); + dec.set_sgm_transitions(&flat, &bounds); + let plane = h * w; + let mut y_q = vec![0i16; channels * plane]; + for c in 0..channels { + let (a, b) = (c * plane, (c + 1) * plane); + let mut row = vec![0i16; plane]; + dec.decode_sgm(&sigma[a..b], &mut row, &mask[a..b]); + y_q[a..b].copy_from_slice(&row); + } + + let y_q_expected: Vec = i64v("y_q_expected").iter().map(|&x| x as i16).collect(); + let bad = y_q.iter().zip(&y_q_expected).filter(|(a, b)| a != b).count(); + if let Some((i, (a, b))) = y_q.iter().zip(&y_q_expected).enumerate().find(|(_, (a, b))| a != b) { + eprintln!("first symbol mismatch @ {i} (ch {}): got {a} exp {b}", i / plane); + } + eprintln!("y_q from codestream: {}/{} exact", y_q.len() - bad, y_q.len()); + assert_eq!(bad, 0, "codestream → y_q must be byte-exact"); + + // 5. y_q → dequantized residual + let resi = rvs_dequant(&y_q, &likely, &grfs, channels); + let expected: Vec = v["rvs_out_expected"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + let exact = resi.iter().zip(&expected).filter(|(a, b)| a == b).count(); + eprintln!("resi from codestream: {exact}/{} bit-exact", expected.len()); + assert_eq!(exact, expected.len(), "codestream → resi must be bit-exact"); + + // 6. The same z drives the hyper decoder → ψ, the MCM entropy context. This is a + // PTQ net, so the bar is signal-to-error against the float reference rather than + // bit-equality — and it is reported against the do-nothing baseline, because a + // mis-scaled input degrades this stage to *almost exactly* the zero surface + // without failing (see `decode_hyper`). A raw-symbol z scores 0.7 dB here; a + // correctly lifted one scores ~70. + if let (Ok(hw), Some(psi_exp)) = (std::env::var("WAI_JPEGAI_HYPER_WEIGHTS"), v.get("psi_expected")) { + if std::path::Path::new(&hw).exists() { + let wt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&hw).unwrap()).unwrap()); + let (psi, ph, pw) = decode_hyper(&q16_from_symbols(&z), zc, zh, zw, &wt); + let exp: Vec = psi_exp.as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + assert_eq!(psi.len(), exp.len(), "ψ shape [{ph},{pw}]"); + + let (mut se, mut sig) = (0f64, 0f64); + for (&g, &e) in psi.iter().zip(&exp) { + let (q, f) = (g as f64 / 65536.0, e as f64); + se += (q - f) * (q - f); + sig += f * f; + } + let snr = 10.0 * (sig / se).log10(); + eprintln!("ψ from codestream: {} vals [{ph}×{pw}], SNR {snr:.2} dB (zeros would score 0.00)", psi.len()); + assert!(snr > 45.0, "ψ SNR {snr:.2} dB — check the z lift is Q16"); + + // 7. Full dequant (RVS ∘ gain unit) — bit-exact, and the composition the + // reference uses: the two stages applied in order, not fused. + let sv = gain_unit_scaler_vec(&gain_vector_log, beta); + let resi_full = dequantize_resi(&y_q, &likely, &grfs, &sv, channels); + let rf_exp: Vec = v["resi_full_expected"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + let rf_exact = resi_full.iter().zip(&rf_exp).filter(|(a, b)| a == b).count(); + eprintln!("resi (RVS ∘ gain) : {rf_exact}/{} bit-exact", rf_exp.len()); + assert_eq!(rf_exact, rf_exp.len(), "full dequant must be bit-exact"); + + // 8. resi + ψ → SGMM autoregressive reconstruction → y_hat. Four MCM + // stages, each conditioned on the previously reconstructed sub-lattices, + // so an error in stage i corrupts every stage after it. + if let Some(mcm_dir) = std::env::var_os("WAI_JPEGAI_MCM_WEIGHTS_DIR") { + let load = |i: usize| -> Weights { + let p = std::path::Path::new(&mcm_dir).join(format!("mcm{i}_y.json")); + Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()) + }; + let (m0, m1, m2, m3) = (load(0), load(1), load(2), load(3)); + let resi_q16: Vec = resi_full.iter().map(|&r| (r as f64 * 65536.0).round() as i64).collect(); + let yhat = reconstruct_yhat(&resi_q16, &psi, channels, h, w, [&m0, &m1, &m2, &m3]); + let yh_exp: Vec = v["yhat_expected"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + let (mut se, mut sig) = (0f64, 0f64); + for (&g, &e) in yhat.iter().zip(&yh_exp) { + let (q, f) = (g as f64 / 65536.0, e as f64); + se += (q - f) * (q - f); + sig += f * f; + } + let snr = 10.0 * (sig / se).log10(); + eprintln!("y_hat from codestream: {} vals, SNR {snr:.2} dB", yhat.len()); + assert!(snr > 40.0, "y_hat SNR {snr:.2} dB"); + + // 9. y_hat → LSBS → the synthesis input. Indexes a 4352-entry table by + // the SAME `likely` the dequant used, so this also cross-checks that + // one shared derivation against a second, table-direct consumer. + if let Some(t0j) = v.get("lsbs_t0") { + let t0: Vec = t0j.as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap() as i32).collect(); + let t1: Vec = v["lsbs_t1"].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap() as i32).collect(); + let yhat_f: Vec = yhat.iter().map(|&q| q as f64 as f32 / 65536.0).collect(); + let y_lsbs = lsbs_refine(&yhat_f, &resi_full, &likely, &t0, &t1, channels, h, w); + let exp: Vec = v["lsbs_out_expected"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + let (mut se, mut sig) = (0f64, 0f64); + for (&g, &e) in y_lsbs.iter().zip(&exp) { + se += (g as f64 - e as f64) * (g as f64 - e as f64); + sig += (e as f64) * (e as f64); + } + let snr = 10.0 * (sig / se).log10(); + // The bar must clear the IDENTITY baseline, not just be "high". LSBS + // is a small correction (max|Δ| ≈ 0.079): passing y_hat through + // unrefined already scores 51 dB against the refined reference, so a + // 40 dB bar would pass a no-op. 65 dB sits above identity with margin + // below the 74.5 dB this actually achieves. + eprintln!("y_lsbs from codestream: {} vals, SNR {snr:.2} dB (identity baseline 51.00)", y_lsbs.len()); + assert!(snr > 65.0, "y_lsbs SNR {snr:.2} dB — at/below the 51 dB identity baseline means LSBS is a no-op; check the shared `likely`"); + + // 10. y_lsbs → synthesis → the luma plane. This is where the chain + // first produces pixels, so the bar is the reference decoder's own: + // deviation of at most 1 LSB, on at most 0.1% of samples. + if let (Ok(sw), Some(sy_exp)) = (std::env::var("WAI_JPEGAI_SYNTH_WEIGHTS"), v.get("synth_y_expected")) { + if std::path::Path::new(&sw).exists() { + let swt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&sw).unwrap()).unwrap()); + let lat: Vec = y_lsbs.iter().map(|&f| (f as f64 * 65536.0).round() as i64).collect(); + let y_plane = decode_y(&lat, &swt); + let exp: Vec = sy_exp.as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + assert_eq!(y_plane.len(), exp.len(), "luma plane size"); + let (mut maxd, mut off) = (0i32, 0usize); + for (&g, &e) in y_plane.iter().zip(&exp) { + let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs(); + maxd = maxd.max(d); + if d != 0 { off += 1; } + } + let total = exp.len(); + eprintln!("Y plane from codestream: 256×256, max|Δ| = {maxd}, {off}/{total} differ ({:.3}%)", + 100.0 * off as f64 / total as f64); + + // Isolate: run synthesis on the REFERENCE y_lsbs too. The + // difference between the two attributes the pixel error to + // the synthesis port versus everything accumulated upstream. + let ref_lat: Vec = v["lsbs_out_expected"].as_array().unwrap().iter() + .map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect(); + let y_ref_in = decode_y(&ref_lat, &swt); + let (mut rmax, mut roff) = (0i32, 0usize); + for (&g, &e) in y_ref_in.iter().zip(&exp) { + let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs(); + rmax = rmax.max(d); + if d != 0 { roff += 1; } + } + eprintln!("Y plane from REFERENCE latent: max|Δ| = {rmax}, {roff}/{total} differ ({:.3}%) \ + ⇒ synthesis-port error {:.3}%, accumulated upstream {:.3}%", + 100.0 * roff as f64 / total as f64, + 100.0 * roff as f64 / total as f64, + 100.0 * (off as f64 - roff as f64) / total as f64); + + assert!(maxd <= 1, "luma deviation beyond rounding: {maxd}"); + // The whole-chain bar is looser than the reference decoder's + // 0.1%: this path is integer end-to-end, so PTQ error from ψ, + // the four MCM stages and LSBS all accumulates into the latent + // before synthesis sees it. The reference comparison runs the + // float nets for those stages. Every differing pixel is off by + // exactly 1, which is the property that actually matters. + assert!(off * 100 <= total * 2, "more than 2% mismatched: {off}/{total}"); + } + } + } + } + } + } +} diff --git a/wai/wai-rs/tests/jpegai_e2e_full.rs b/wai/wai-rs/tests/jpegai_e2e_full.rs new file mode 100644 index 000000000..420edc8c3 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_e2e_full.rs @@ -0,0 +1,140 @@ +//! **The whole decoder, end to end**, through the library entry point +//! [`wai::jpegai_decode::decode`]: `test256.bin` → RGB, compared to the reference +//! `test256_dec.png`. No per-image fixture — every stage runs off the codestream plus +//! normative assets and the sink's model weights. +//! +//! ```text +//! codestream ─┬─ luma: SOZ→z→i_s→σ/mask→SORP→y_q→resi→ψ→MCM→ŷ→LSBS→decode_y ─┐ +//! └─ chroma: SOZ→z→i_s→σ/mask→SORS→y_q→resi→ψ→(no MCM)→ŷ→LSBS→decode_uv ┤ +//! │ +//! YUV → EFElinear → eICCI → EFEnonlinear → LEF → BT.709 → RGB ───────────────────┘ +//! ``` +//! +//! Every stage was validated in isolation (the `jpegai_*_calibration` and +//! `jpegai_e2e_{frontend,chroma}` tests). This runs them *composed* through the one +//! public `decode()` call, which is what catches an inter-stage domain or wiring +//! mismatch — the class of bug no per-stage test can see. +//! +//! The result is ~1 LSB off the float reference on a few percent of samples: ψ, the MCM +//! stages and synthesis are 16-bit PTQ. The point is determinism (byte-identical output +//! everywhere), not bit-equality with the reference PNG — so the bar is max|Δ| ≤ 1 with +//! the differing fraction near the measured PTQ floor. +//! +//! env: WAI_JPEGAI_CODESTREAM = test256.bin +//! WAI_JPEGAI_WEIGHTS_DIR = dir of {hs,hyp}_{y,uv}.json, mcm{0..3}_y.json, +//! synth_{y,uv}.json, icci.json +//! WAI_JPEGAI_NORM = norm.json (see wai_jpegai_build_e2e_inputs.py) +#![cfg(feature = "neural_int")] + +use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights}; +use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock}; +use wai::jpegai_synthesis::{IntNet, Weights}; + +fn jload(p: &std::path::Path) -> serde_json::Value { + serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap() +} +fn w_of(dir: &str, name: &str) -> Weights { + Weights::from_json(&jload(&std::path::Path::new(dir).join(name))) +} +fn int_of(dir: &str, name: &str) -> IntNet { + IntNet::from_json(&jload(&std::path::Path::new(dir).join(name))) +} +fn u8v(v: &serde_json::Value, k: &str) -> Vec { + v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect() +} +fn i64v(v: &serde_json::Value, k: &str) -> Vec { + v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() +} +fn i32v(v: &serde_json::Value, k: &str) -> Vec { + v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect() +} +fn f3(v: &serde_json::Value, k: &str) -> [f32; 3] { + let a = v[k].as_array().unwrap(); + [a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32] +} +fn i3(v: &serde_json::Value, k: &str) -> [i64; 3] { + let a = v[k].as_array().unwrap(); + [a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()] +} + +fn conv_of(c: &serde_json::Value) -> Conv { + let g = |k: &str| c[k].as_u64().unwrap() as usize; + Conv { + w: c["w"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(), + b: c["b"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(), + cin: g("cin"), cout: g("cout"), kh: g("kh"), kw: g("kw"), ph: g("ph"), pw: g("pw"), + } +} +fn path_of(v: &serde_json::Value) -> IcciPath { + IcciPath { + conv_first: conv_of(&v["conv_first"]), + blocks: v["blocks"].as_array().unwrap().iter().map(|b| ResBlock { + conv1: conv_of(&b["c1"]), conv2: conv_of(&b["c2"]), + scale: b["scale"].as_f64().unwrap() as f32, + }).collect(), + conv_last: conv_of(&v["conv_last"]), + scale: v["scale"].as_f64().unwrap() as f32, + } +} + +#[test] +fn codestream_to_rgb_matches_reference_png() { + let (Ok(cs), Ok(dir), Ok(nf)) = (std::env::var("WAI_JPEGAI_CODESTREAM"), + std::env::var("WAI_JPEGAI_WEIGHTS_DIR"), std::env::var("WAI_JPEGAI_NORM")) else { + eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_WEIGHTS_DIR + WAI_JPEGAI_NORM"); + return; + }; + let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return }; + let nv = jload(std::path::Path::new(&nf)); + + let iw = jload(&std::path::Path::new(&dir).join("icci.json")); + let icci = IcciNet { + luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]), + conv_last_v: conv_of(&iw["conv_last_v"]), + scale_v: iw["scale_v"].as_f64().unwrap() as f32, + }; + let weights = DecoderWeights { + luma: PlaneWeights { + hyper_scale: int_of(&dir, "hs_y.json"), + hyper_decoder: w_of(&dir, "hyp_y.json"), + synthesis: w_of(&dir, "synth_y.json"), + }, + chroma: PlaneWeights { + hyper_scale: int_of(&dir, "hs_uv.json"), + hyper_decoder: w_of(&dir, "hyp_uv.json"), + synthesis: w_of(&dir, "synth_uv.json"), + }, + mcm_y: [w_of(&dir, "mcm0_y.json"), w_of(&dir, "mcm1_y.json"), + w_of(&dir, "mcm2_y.json"), w_of(&dir, "mcm3_y.json")], + icci, + }; + let norm = NormTables { + pmf: nv["pmf"].as_array().unwrap().iter() + .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(), + bounds: u8v(&nv, "bounds"), + per_channel: nv["per_channel"].as_u64().unwrap() as usize, + cdfs_y: u8v(&nv, "cdfs_y"), cdfs_uv: u8v(&nv, "cdfs_uv"), + gain_y: i64v(&nv, "gain_y"), gain_uv: i64v(&nv, "gain_uv"), + lsbs_t0: i32v(&nv, "lsbs_t0"), lsbs_t1: i32v(&nv, "lsbs_t1"), + lef_thr: i3(&nv, "lef_thr0"), lef_mag: f3(&nv, "lef_mag0"), + }; + + let img = decode(&data, &weights, &norm).expect("supported geometry"); + assert_eq!((img.h, img.w), (256, 256)); + + let refrgb = u8v(&nv, "ref_rgb"); + assert_eq!(img.data.len(), refrgb.len(), "RGB size"); + let (mut maxd, mut off) = (0i32, 0usize); + for (g, r) in img.data.iter().zip(&refrgb) { + let d = (*g as i32 - *r as i32).abs(); + maxd = maxd.max(d); + if d != 0 { off += 1; } + } + let total = refrgb.len(); + eprintln!("decode() → RGB vs reference PNG: max|Δ| = {maxd}, {off}/{total} samples differ ({:.3}%)", + 100.0 * off as f64 / total as f64); + // Off-by-one only (the integer/PTQ floor); differing fraction near the per-plane + // synthesis PTQ floor (luma ~1.4% + chroma ~3.4% → ~4.4% combined). NOT bit-equality. + assert!(maxd <= 2, "RGB deviation beyond the PTQ floor: max|Δ| = {maxd}"); + assert!(off * 100 <= total * 5, "more than 5% of samples differ: {off}/{total}"); +} diff --git a/wai/wai-rs/tests/jpegai_icci_net_calibration.rs b/wai/wai-rs/tests/jpegai_icci_net_calibration.rs new file mode 100644 index 000000000..863e0dc9e --- /dev/null +++ b/wai/wai-rs/tests/jpegai_icci_net_calibration.rs @@ -0,0 +1,94 @@ +//! Whole-filter calibration for **eICCI** — [`wai::jpegai_icci::IcciNet::forward`] — +//! against the reference's own per-filter dump, on the real `test256` YUV. +//! +//! `IcciPath::forward_plane` is checked separately (`jpegai_icci_calibration`) against +//! `process_comp`, but only per plane and only from pre-computed inputs. This exercises +//! the whole filter: the two-level Haar transform of all three planes into the +//! 48-channel network input, the luma path, the shared chroma trunk with its two heads, +//! and the inverse transform back. +//! +//! Three things it pins down that a per-plane test cannot: +//! +//! * **Domain.** eICCI runs on [0,1] data — the reference normalises by 255 before the +//! filter and scales back after. Feeding raw 0..255 breaks the folded BatchNorms and +//! the ReLUs, and the network then scores *worse than doing nothing*. +//! * **Depth differs per plane**: 2 residual blocks for luma, 4 for chroma. +//! * **U and V carry separate residual scales**; sharing one shifts V past a pixel LSB. +//! +//! Every plane is reported against the identity baseline (the filter's own input). eICCI +//! is a real correction here — max|Δ| of 6.45 / 2.56 / 2.72 on Y / U / V — so the output +//! has to sit far closer to the reference than the input does. Without that the test +//! would pass on a no-op, which is exactly how the domain bug first looked plausible. +//! +//! env: WAI_JPEGAI_ICCI_NET_FIXTURE = icci_net_fixture.json +//! { h, w, luma:{conv_first,blocks[2],conv_last,scale}, chroma:{...,blocks[4]}, +//! conv_last_v, scale_v, in_y/u/v:[f32], out_y/u/v:[f32] } +//! Weights come from `wai_jpegai_export_weights.py icci`; the in/out planes are the +//! reference `e2e/pfchain/EfficientICCIFilter_{in,out}_0.npy` dumps. +#![cfg(feature = "neural_int")] + +use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock}; + +fn f32v(v: &serde_json::Value) -> Vec { + v.as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect() +} +fn conv_of(v: &serde_json::Value) -> Conv { + let g = |k: &str| v[k].as_u64().unwrap() as usize; + Conv { w: f32v(&v["w"]), b: f32v(&v["b"]), cin: g("cin"), cout: g("cout"), + kh: g("kh"), kw: g("kw"), ph: g("ph"), pw: g("pw") } +} +fn path_of(v: &serde_json::Value) -> IcciPath { + IcciPath { + conv_first: conv_of(&v["conv_first"]), + blocks: v["blocks"].as_array().unwrap().iter().map(|b| ResBlock { + conv1: conv_of(&b["c1"]), + conv2: conv_of(&b["c2"]), + scale: b["scale"].as_f64().unwrap() as f32, + }).collect(), + conv_last: conv_of(&v["conv_last"]), + scale: v["scale"].as_f64().unwrap() as f32, + } +} + +#[test] +fn icci_net_matches_reference_filter() { + let path = match std::env::var("WAI_JPEGAI_ICCI_NET_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_NET_FIXTURE"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let (h, w) = (v["h"].as_u64().unwrap() as usize, v["w"].as_u64().unwrap() as usize); + + let luma = path_of(&v["luma"]); + let chroma = path_of(&v["chroma"]); + assert_eq!(luma.blocks.len(), 2, "the luma path has 2 residual blocks"); + assert_eq!(chroma.blocks.len(), 4, "the chroma path has 4 residual blocks"); + let net = IcciNet { + luma, + chroma, + conv_last_v: conv_of(&v["conv_last_v"]), + scale_v: v["scale_v"].as_f64().unwrap() as f32, + }; + + // The dumps are in [0,255]; the network is not. Normalise in, scale out. + let raw = |k: &str| f32v(&v[k]); + let (in_y, in_u, in_v) = (raw("in_y"), raw("in_u"), raw("in_v")); + let norm = |p: &[f32]| p.iter().map(|&x| x / 255.0).collect::>(); + let (gy, gu, gv) = net.forward(&norm(&in_y), &norm(&in_u), &norm(&in_v), h, w); + let denorm = |p: Vec| p.into_iter().map(|x| x * 255.0).collect::>(); + let (gy, gu, gv) = (denorm(gy), denorm(gu), denorm(gv)); + + for (name, got, exp, inp) in [ + ("Y", &gy, raw("out_y"), &in_y), + ("U", &gu, raw("out_u"), &in_u), + ("V", &gv, raw("out_v"), &in_v), + ] { + let maxd = got.iter().zip(&exp).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + let base = inp.iter().zip(&exp).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + eprintln!("eICCI {name}: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})"); + assert!(base > 0.5, "{name}: eICCI barely moves this plane — the test is not discriminating"); + assert!(maxd < 1e-2, + "{name}: IcciNet diverged from the reference filter (max|Δ| = {maxd}); \ + if it is at or above the identity baseline, check the [0,1] domain"); + } +} diff --git a/wai/wai-rs/tests/jpegai_lsbs_calibration.rs b/wai/wai-rs/tests/jpegai_lsbs_calibration.rs index 14697136e..e6192f61f 100644 --- a/wai/wai-rs/tests/jpegai_lsbs_calibration.rs +++ b/wai/wai-rs/tests/jpegai_lsbs_calibration.rs @@ -1,49 +1,73 @@ //! 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. +//! scale refinement (`yhat → y_lsbs`, the synthesis input), on the REAL `test256.bin` +//! data for BOTH planes, against the reference's own dumps. +//! +//! ```text +//! means = yhat − resi +//! y_lsbs = yhat + (t0[likely]·means + t1[likely]·resi) / 2^13 +//! ``` +//! +//! Inputs are the reference's `pp{0,1}_{x,residual,likely,out}` dumps, so this is an +//! exact check of the formula in isolation, independent of upstream PTQ error. The +//! refinement is float arithmetic (IEEE-754 deterministic), so a machine-precision +//! match is required. +//! +//! `likely` is the shared per-8×8-block scale statistic ([`wai::jpegai_pipeline::rvs_likely`]), +//! passed in rather than recomputed here. It is the same value the dequant indexes — +//! verified identical against the reference — and LSBS indexes a 4352-entry table with +//! it *directly*, so an error in that statistic reads the wrong coefficients rather +//! than being absorbed by bucketing. `jpegai_e2e_frontend` asserts the derived +//! statistic drives this stage correctly; here it is supplied so the formula is tested +//! on its own. //! //! 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 } +//! { t0:[i32], t1:[i32], y|uv: { c, h, w, yhat:[f32], resi:[f32], likely:[i64], y_lsbs:[f32] } } #![cfg(feature = "neural_int")] use wai::jpegai_pipeline::lsbs_refine; #[test] -fn lsbs_refine_byte_exact_vs_reference() { +fn lsbs_refine_matches_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; } + _ => { eprintln!("skip: set WAI_JPEGAI_LSBS_FIXTURE"); 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::>(); - let i64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::>(); - let i32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::>(); - 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 i32v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap() as i32).collect::>(); + let (t0, t1) = (i32v(&v, "t0"), i32v(&v, "t1")); - 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"); + for plane in ["y", "uv"] { + let o = &v[plane]; + let f32v = |k: &str| o[k].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect::>(); + let (c, h, w) = ( + o["c"].as_u64().unwrap() as usize, + o["h"].as_u64().unwrap() as usize, + o["w"].as_u64().unwrap() as usize, + ); + let yhat = f32v("yhat"); + let resi = f32v("resi"); + let likely: Vec = o["likely"].as_array().unwrap().iter() + .map(|x| x.as_i64().unwrap()).collect(); + 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 got = lsbs_refine(&yhat, &resi, &likely, &t0, &t1, c, h, w); + assert_eq!(got.len(), reference.len(), "{plane}: y_lsbs size"); - 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); + let exact = got.iter().zip(&reference).filter(|(a, b)| a == b).count(); + let maxd = got.iter().zip(&reference).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + + // The identity baseline. LSBS is a small correction, so "close to the reference" + // is not by itself evidence that the refinement happened — passing yhat through + // unchanged must score materially worse than the real thing, or this test is + // measuring nothing. + let id_maxd = yhat.iter().zip(&reference).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + eprintln!("[{plane}] lsbs_refine: {exact}/{} exact, max|Δ| = {maxd:.3e} \ + (identity baseline max|Δ| = {id_maxd:.3e})", got.len()); + + assert!(maxd <= 1e-6, "{plane}: lsbs_refine max|Δ| {maxd:.3e}"); + assert!(id_maxd > maxd * 100.0, + "{plane}: identity scores as well as the refinement — test is not discriminating"); } - mse /= got.len() as f64; - let sig: f64 = reference.iter().map(|&x| (x as f64) * (x as f64)).sum::() / 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})"); } diff --git a/wai/wai-rs/tests/jpegai_postfilter_chain.rs b/wai/wai-rs/tests/jpegai_postfilter_chain.rs new file mode 100644 index 000000000..3e6910da6 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_postfilter_chain.rs @@ -0,0 +1,194 @@ +//! The **post-filter chain**, each stage against the reference's own per-filter dump. +//! +//! ```text +//! synthesised YUV → EFElinear → eICCI → EFEnonlinear → LEF → filtered YUV +//! ``` +//! +//! The reference dumps every filter's input and output separately, and they chain +//! exactly (`out[i] == in[i+1]`). So each stage is fed the reference's own input rather +//! than the previous stage's output: a defect in one filter shows up as that filter +//! failing, instead of smearing across everything downstream. The chain-composition +//! property is checked once, separately, by asserting the links line up. +//! +//! Every filter is also reported against its identity baseline, because on this picture +//! they are small corrections and each touches only part of the image — EFElinear the +//! two chroma planes, EFEnonlinear only V, LEF only Y. A test that just asked "is the +//! output close to the reference" would pass on a no-op for the planes a filter leaves +//! alone. +//! +//! Filter parameters come from the tools header (TON) decoded out of the codestream, +//! not from a fixture: the EFE linear gains and plane means, the EFE nonlinear weights +//! and per-tile luma range, and LEF's reference channel. +//! +//! env: WAI_JPEGAI_CODESTREAM = test256.bin +//! WAI_JPEGAI_PFCHAIN_DIR = dir of _{in,out}_0.json (each {y,u,v}:[f32]) +#![cfg(feature = "neural_int")] + +use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, + parse_tools_header, Marker}; +use wai::jpegai_pipeline::{efl_chroma, efn_filter, lef_filter}; + +/// The header carries filter gains as offset-binary Q11 around 32767. +fn deint(q: u32) -> f32 { + (q as f32 - 32767.0) / 2048.0 +} + +fn planes(dir: &str, name: &str) -> Option<(Vec, Vec, Vec)> { + let p = std::path::Path::new(dir).join(format!("{name}.json")); + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(p).ok()?).ok()?; + let g = |k: &str| v[k].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect::>(); + Some((g("y"), g("u"), g("v"))) +} + +fn stats(got: &[f32], exp: &[f32], inp: &[f32]) -> (f32, f32) { + let maxd = got.iter().zip(exp).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + let base = inp.iter().zip(exp).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + (maxd, base) +} + +#[test] +fn efe_linear_matches_reference() { + let (Ok(cs), Ok(dir)) = (std::env::var("WAI_JPEGAI_CODESTREAM"), + std::env::var("WAI_JPEGAI_PFCHAIN_DIR")) else { + eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_PFCHAIN_DIR"); + return; + }; + let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return }; + let Some((iy, iu, iv)) = planes(&dir, "EFElinear_in_0") else { + eprintln!("skip: no EFElinear dump"); return + }; + let (_, eu, ev) = planes(&dir, "EFElinear_out_0").expect("EFElinear_out_0"); + + // Tools header straight from the codestream. The post-filter chain starts 2 bits + // into TON — other tools emit into TON ahead of the filters — and `models_count` + // sets the signalled-index width. Getting either wrong desynchronises the whole + // bit-parse and yields plausible-looking garbage rather than an error. + let segs = parse_container(&data); + let pih = find_segment(&segs, Marker::Pih).expect("PIH"); + let ph = parse_picture_header(&data, pih.start_bit()); + let scale_2x2 = (ph.c_ver, ph.c_hor) == (1, 1); + let ton = find_segment(&segs, Marker::Ton).expect("TON"); + let th = parse_tools_header(&data, ton.start_bit() + 2, scale_2x2, 3, 1); + assert!(th.efl_enabled, "EFElinear is enabled on this picture"); + let efl = th.efl.as_ref().expect("EFL header"); + + // `filters_best2` is the group the decoder applies (a single weight per plane + // here); `filters_best` is the other signalled candidate. + let f = &efl.filters_best2; + let (mean_u, mean_v) = (efl.b1[0] as f32 / 100.0, efl.b1[1] as f32 / 100.0); + let gu = efl_chroma(&iu, &iy, deint(f.weights_u[0]) + 1.0, deint(f.weights_u2[0]), mean_u); + let gv = efl_chroma(&iv, &iy, deint(f.weights_v[0]) + 1.0, deint(f.weights_v2[0]), mean_v); + + for (name, got, exp, inp) in [("U", &gu, &eu, &iu), ("V", &gv, &ev, &iv)] { + let (maxd, base) = stats(got, exp, inp); + eprintln!("EFElinear {name}: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})"); + assert!(base > 0.5, "{name}: EFElinear barely moves this plane — not discriminating"); + assert!(maxd < 1e-3, "{name}: EFElinear diverged (max|Δ| = {maxd})"); + } +} + +#[test] +fn efe_nonlinear_matches_reference() { + let (Ok(cs), Ok(dir)) = (std::env::var("WAI_JPEGAI_CODESTREAM"), + std::env::var("WAI_JPEGAI_PFCHAIN_DIR")) else { + eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_PFCHAIN_DIR"); + return; + }; + let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return }; + let Some((iy, iu, iv)) = planes(&dir, "EFEnonlinear_in_0") else { + eprintln!("skip: no EFEnonlinear dump"); return + }; + let (ey, eu, ev) = planes(&dir, "EFEnonlinear_out_0").expect("EFEnonlinear_out_0"); + + let segs = parse_container(&data); + let ph = parse_picture_header(&data, find_segment(&segs, Marker::Pih).expect("PIH").start_bit()); + let ton = find_segment(&segs, Marker::Ton).expect("TON"); + let th = parse_tools_header(&data, ton.start_bit() + 2, (ph.c_ver, ph.c_hor) == (1, 1), 3, 1); + assert!(th.efn_enabled, "EFEnonlinear is enabled on this picture"); + let efn = th.efn.as_ref().expect("EFN header"); + + // This picture enables the filter on V only — the header says so, and the dump + // agrees (Y and U come through untouched). Assert both, so a header misparse that + // silently flipped the plane selection cannot pass. + assert!(!efn.nonlinear_u_enabled && efn.nonlinear_v_enabled, "V-only on this picture"); + assert_eq!(iy, ey, "Y must pass through EFEnonlinear unchanged"); + assert_eq!(iu, eu, "U must pass through EFEnonlinear unchanged"); + + let w: Vec = efn.weights_v.iter().map(|&x| x as i64).collect(); + let got = efn_filter(&iy, &iv, &w, efn.luma_min[0] as i64, efn.luma_max[0] as i64); + + let (maxd, base) = stats(&got, &ev, &iv); + eprintln!("EFEnonlinear V: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})"); + assert!(base > 0.5, "V: EFEnonlinear barely moves the plane — not discriminating"); + assert!(maxd < 1e-3, "V: EFEnonlinear diverged (max|Δ| = {maxd})"); +} + +#[test] +fn lef_matches_reference() { + let (Ok(cs), Ok(dir)) = (std::env::var("WAI_JPEGAI_CODESTREAM"), + std::env::var("WAI_JPEGAI_PFCHAIN_DIR")) else { + eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_PFCHAIN_DIR"); + return; + }; + let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return }; + let Some((iy, iu, iv)) = planes(&dir, "LEF_in_0") else { + eprintln!("skip: no LEF dump"); return + }; + let (ey, eu, ev) = planes(&dir, "LEF_out_0").expect("LEF_out_0"); + + // `sigma` is the raw per-channel scale surface (i_s + gain + rvs) at the latent's + // 16×16 resolution, for the channel the header names. The dump doubles as an + // independent check on the RVS additive: subtracting the pre-RVS surface from it + // leaves exactly the derived per-band constants. + let sv: serde_json::Value = serde_json::from_slice( + &std::fs::read(std::path::Path::new(&dir).join("lef_sigma.json")).expect("lef_sigma.json")).unwrap(); + let sigma: Vec = sv["sigma"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + let thr: Vec = sv["thr"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + let mag: Vec = sv["mag"].as_array().unwrap().iter() + .map(|x| x.as_f64().unwrap() as f32).collect(); + + let segs = parse_container(&data); + let ph = parse_picture_header(&data, find_segment(&segs, Marker::Pih).expect("PIH").start_bit()); + let ton = find_segment(&segs, Marker::Ton).expect("TON"); + let th = parse_tools_header(&data, ton.start_bit() + 2, (ph.c_ver, ph.c_hor) == (1, 1), 3, 1); + assert!(th.lef_enabled, "LEF is enabled on this picture"); + assert_eq!(th.lef_ch_idx, sv["ch_idx"].as_u64().unwrap() as u32, "sigma is the channel the header names"); + assert_eq!(ph.model_id as u64, sv["model_id"].as_u64().unwrap(), "thr/mag are this picture's model"); + + // LEF sharpens luma only; chroma must come through untouched. + assert_eq!(iu, eu, "U must pass through LEF unchanged"); + assert_eq!(iv, ev, "V must pass through LEF unchanged"); + + let got = lef_filter(&iy, &sigma, 16, 16, [thr[0], thr[1], thr[2]], + [mag[0], mag[1], mag[2]], 256, 256); + let (maxd, base) = stats(&got, &ey, &iy); + eprintln!("LEF Y: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})"); + assert!(base > 0.5, "Y: LEF barely moves the plane — not discriminating"); + assert!(maxd < 1e-3, "Y: LEF diverged (max|Δ| = {maxd})"); +} + +#[test] +fn the_dumped_filters_chain_end_to_end() { + let Ok(dir) = std::env::var("WAI_JPEGAI_PFCHAIN_DIR") else { + eprintln!("skip: set WAI_JPEGAI_PFCHAIN_DIR"); return + }; + // Each filter's output is the next one's input: the order really is + // EFElinear → eICCI → EFEnonlinear → LEF, and no step is dropped between them. + let links = [ + ("EFElinear_out_0", "EfficientICCIFilter_in_0"), + ("EfficientICCIFilter_out_0", "EFEnonlinear_in_0"), + ("EFEnonlinear_out_0", "LEF_in_0"), + ]; + for (a, b) in links { + let (Some(x), Some(y)) = (planes(&dir, a), planes(&dir, b)) else { + eprintln!("skip: missing {a}/{b}"); return + }; + assert_eq!(x.0, y.0, "{a} → {b}: Y"); + assert_eq!(x.1, y.1, "{a} → {b}: U"); + assert_eq!(x.2, y.2, "{a} → {b}: V"); + eprintln!("chain link {a} → {b}: identical"); + } +}