e6453a9566
Gate 7 made these seven tests require `legacy_int` in Cargo.toml, but their files were still gated on `neural_int` or `neural` alone, which `check-tests-in-ci.py` reports: a file's `#![cfg]` must name what its `[[test]]` requires. Each file is now gated on `all(..., feature = "legacy_int")`, as `jpegai_ops.rs` is. Three test comments named the design's appendix and a plan item. They now name the text's sections (§11's `role_profiles`, §13.6.2's table). Co-Authored-By: Claude Opus 5.5 <[email protected]>
69 lines
3.6 KiB
Rust
69 lines
3.6 KiB
Rust
//! 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(all(feature = "neural_int", feature = "legacy_int"))]
|
||
|
||
use wai::jpegai_synthesis::{decode_y, decode_uv, Weights};
|
||
|
||
fn load(k: &str, v: &serde_json::Value) -> Vec<f64> {
|
||
v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect()
|
||
}
|
||
fn q16(f: &[f64]) -> Vec<i64> { 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()).unwrap();
|
||
let wu = Weights::from_json(&serde_json::from_slice(&std::fs::read(&uw).unwrap()).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");
|
||
}
|
||
|
||
// Assert the better Y candidate reproduces the reference synthesis within PTQ.
|
||
let best_y = ["yhat_y", "y_lsbs"].iter()
|
||
.map(|c| psnr_u8(&decode_y(&q16(&load(c, &v)), &wy), &ypix).1)
|
||
.fold(f64::MIN, f64::max);
|
||
eprintln!("best Y PSNR = {best_y:.2} dB");
|
||
assert!(best_y > 45.0, "Y synthesis on real latent regressed (best {best_y:.2} dB)");
|
||
|
||
// UV: model_uv's input is the 256-ch (y‖uv) latent = concat(y_lsbs[160], uv_lsbs[96]).
|
||
let uvpix = load("uvpix", &v);
|
||
let mut uv_in: Vec<i64> = q16(&load("y_lsbs", &v)); // 160×16×16
|
||
uv_in.extend(q16(&load("uv_lsbs", &v))); // + 96×16×16 = 256×16×16
|
||
let uv = decode_uv(&uv_in, &wu);
|
||
let (mdu, psu) = psnr_u8(&uv, &uvpix);
|
||
eprintln!("UV synthesis on real latent [concat(y_lsbs,uv_lsbs)]: max|Δ|={mdu:.1} PSNR={psu:.2} dB");
|
||
assert!(psu > 45.0, "UV synthesis on real latent regressed ({psu:.2} dB)");
|
||
}
|