diff --git a/wai/wai-rs/tests/jpegai_e2e_synthesis.rs b/wai/wai-rs/tests/jpegai_e2e_synthesis.rs new file mode 100644 index 000000000..8e4b4601d --- /dev/null +++ b/wai/wai-rs/tests/jpegai_e2e_synthesis.rs @@ -0,0 +1,67 @@ +//! End-to-end integration check: run the synthesis decoders on the REAL reference +//! latent from `test256.bin` (checkpoint `yhat_y`/`yhat_uv`) and compare the pixels +//! to the reference synthesis output (`ypix`/`uvpix`). Unlike the per-net calibration +//! (synthetic latent vs the ONNX), this validates the synthesis stage on real +//! pipeline data — the first integration milestone. PSNR ≈ the 16-bit PTQ (~58 dB). +//! +//! env: WAI_JPEGAI_E2E_CKPT = e2e_synth_checkpoint.json {yhat_y,yhat_uv,ypix,uvpix,y_lsbs} +//! WAI_JPEGAI_Y_WEIGHTS = synth16_y_weights.json +//! WAI_JPEGAI_UV_WEIGHTS = synth16_uv_weights.json +#![cfg(feature = "neural_int")] + +use wai::jpegai_synthesis::{decode_y, decode_uv, Weights}; + +fn load(k: &str, v: &serde_json::Value) -> Vec { + v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect() +} +fn q16(f: &[f64]) -> Vec { f.iter().map(|&x| (x * 65536.0).round() as i64).collect() } +fn psnr_u8(got: &[u8], refpix: &[f64]) -> (f64, f64) { + let (mut maxd, mut mse) = (0f64, 0f64); + for i in 0..got.len() { + let r = refpix[i].clamp(0.0, 255.0); + let d = (got[i] as f64 - r).abs(); + maxd = maxd.max(d); mse += d * d; + } + mse /= got.len() as f64; + (maxd, if mse > 0.0 { 20.0 * 255f64.log10() - 10.0 * mse.log10() } else { f64::INFINITY }) +} + +#[test] +fn synthesis_on_real_reference_latent() { + let ckpt = match std::env::var("WAI_JPEGAI_E2E_CKPT") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_E2E_CKPT"); return; } + }; + let (yw, uw) = (std::env::var("WAI_JPEGAI_Y_WEIGHTS"), std::env::var("WAI_JPEGAI_UV_WEIGHTS")); + let (yw, uw) = match (yw, uw) { + (Ok(a), Ok(b)) if std::path::Path::new(&a).exists() && std::path::Path::new(&b).exists() => (a, b), + _ => { eprintln!("skip: set WAI_JPEGAI_Y_WEIGHTS + WAI_JPEGAI_UV_WEIGHTS"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&ckpt).unwrap()).unwrap(); + let wy = Weights::from_json(&serde_json::from_slice(&std::fs::read(&yw).unwrap()).unwrap()); + let wu = Weights::from_json(&serde_json::from_slice(&std::fs::read(&uw).unwrap()).unwrap()); + + // Y: decode_y(yhat_y) vs ypix. Try both the pre-lsbs (yhat_y) and post-lsbs + // (y_lsbs) latent — whichever the synthesis actually consumes will hit ~58 dB. + let ypix = load("ypix", &v); + for cand in ["yhat_y", "y_lsbs"] { + let lat = q16(&load(cand, &v)); + let y = decode_y(&lat, &wy); + let (md, ps) = psnr_u8(&y, &ypix); + eprintln!("Y synthesis on real latent [{cand}]: max|Δ|={md:.1} PSNR={ps:.2} dB"); + } + + // UV: decode_uv needs the 256-ch (y‖uv) input; yhat_uv is only the 96-ch uv part. + // Report what we can (the uv-only path is exercised in the per-net calibration); + // here we log the uvpix reference stats for the full-pipeline wiring step. + let uvpix = load("uvpix", &v); + eprintln!("uvpix reference: {} samples, range check pending full 256-ch assembly", uvpix.len()); + let _ = (decode_uv as fn(&[i64], &Weights) -> Vec, &wu); + + // Assert the better Y candidate reproduces the reference synthesis within PTQ. + let best = ["yhat_y", "y_lsbs"].iter() + .map(|c| psnr_u8(&decode_y(&q16(&load(c, &v)), &wy), &ypix).1) + .fold(f64::MIN, f64::max); + eprintln!("best Y PSNR = {best:.2} dB"); + assert!(best > 45.0, "synthesis on real latent regressed (best {best:.2} dB)"); +}