48becebcab
8a7492b23 moved ort to 2.0.0-rc.15 and gave ort_session.rs the environment rc.15 asks every session builder for; five test files build sessions of their own (jpegai_e2e and the four calibration tests) and still called `Session::builder()` with none. They build only with `neural` (and `legacy_int`), which no push-triggered job builds, and a checked list cut at ten lines hid them. rc.15 allows one environment per process, and these tests call into wai in the same process, so they take wai's: `ort_session::environment` is public as `wai::ort_environment` (a new path, compatible). All five build and run under neural,semantic,legacy_int; without their JPEG AI fixtures they return early, as before. jpegai_decode_path builds and passes under its features. clippy reports nothing in the files this touches. Changelog: none Co-Authored-By: Claude Opus 5.5 <[email protected]>
66 lines
3.3 KiB
Rust
66 lines
3.3 KiB
Rust
//! Calibration: my integer `jpegai_synthesis::decode_y` vs the float reference —
|
|
//! run the SAME latent through `synthesis16.onnx` (ort) and through `decode_y`,
|
|
//! compare the Y planes (max abs diff + PSNR). This is the PSNR pass that tells
|
|
//! whether the whole-net integer executor is numerically faithful (it exposes the
|
|
//! approximated gates). Needs `neural` (ort). Skipped if assets absent.
|
|
//!
|
|
//! env: WAI_JPEGAI_SYNTH16_Y = path to synthesis16.onnx (model_y),
|
|
//! WAI_JPEGAI_Y_WEIGHTS = path to the exported 16-bit weights JSON.
|
|
#![cfg(all(feature = "neural", feature = "legacy_int"))]
|
|
|
|
use ort::session::{builder::GraphOptimizationLevel, Session};
|
|
use ort::value::Value;
|
|
use wai::jpegai_synthesis::{decode_y, Weights};
|
|
|
|
fn ort_synth(path: &str, x: Vec<f32>) -> (Vec<f32>, Vec<usize>) {
|
|
let mut s = Session::builder(wai::ort_environment().unwrap()).unwrap()
|
|
.with_intra_threads(1).unwrap().with_inter_threads(1).unwrap()
|
|
.with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
|
|
.commit_from_file(path).unwrap();
|
|
let name = s.inputs().first().unwrap().name().to_string();
|
|
let arr = ndarray::Array4::from_shape_vec((1, 160, 16, 16), x).unwrap();
|
|
let v = Value::from_array(arr).unwrap();
|
|
let out = s.run(ort::inputs![name.as_str() => v]).unwrap();
|
|
let view = out[0].try_extract_array::<f32>().unwrap();
|
|
(view.iter().copied().collect(), view.shape().to_vec())
|
|
}
|
|
|
|
#[test]
|
|
fn synthesis_int_vs_float_psnr() {
|
|
let onnx = match std::env::var("WAI_JPEGAI_SYNTH16_Y") {
|
|
Ok(p) if std::path::Path::new(&p).exists() => p,
|
|
_ => { eprintln!("skip: set WAI_JPEGAI_SYNTH16_Y to synthesis16.onnx"); return; }
|
|
};
|
|
let wjson = match std::env::var("WAI_JPEGAI_Y_WEIGHTS") {
|
|
Ok(p) if std::path::Path::new(&p).exists() => p,
|
|
_ => { eprintln!("skip: set WAI_JPEGAI_Y_WEIGHTS to the weights JSON"); return; }
|
|
};
|
|
// deterministic small latent [160,16,16] (typical post-LSBS magnitude)
|
|
let n = 160 * 16 * 16;
|
|
let lf: Vec<f32> = (0..n).map(|i| (((i % 97) as f32) - 48.0) / 96.0).collect();
|
|
|
|
let (yf, shape) = ort_synth(&onnx, lf.clone());
|
|
let (fmn, fmx) = yf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
|
|
eprintln!("float synth out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, yf.len(), fmn, fmx);
|
|
eprintln!("float sample [0..6] = {:?}", &yf[..6]);
|
|
|
|
let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
|
|
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
|
|
let wt = Weights::from_json(&v).unwrap();
|
|
let yi = decode_y(&li, &wt);
|
|
|
|
assert_eq!(yi.len(), yf.len(), "output size mismatch");
|
|
let (mut maxd, mut mse) = (0f64, 0f64);
|
|
for i in 0..yi.len() {
|
|
let f = yf[i].clamp(0.0, 255.0) as f64;
|
|
let d = (yi[i] as f64 - f).abs();
|
|
maxd = maxd.max(d);
|
|
mse += d * d;
|
|
}
|
|
mse /= yi.len() as f64;
|
|
let psnr = if mse > 0.0 { 20.0 * (255.0f64).log10() - 10.0 * mse.log10() } else { f64::INFINITY };
|
|
eprintln!("INT vs FLOAT synthesis: max|Δ| = {:.2} MSE = {:.2} PSNR = {:.2} dB", maxd, mse, psnr);
|
|
// the integer decode reproduces the float ONNX synthesis to within the 16-bit
|
|
// PTQ quantization (~1 LSB); PSNR ≈ 58 dB. Guard against regressions.
|
|
assert!(psnr > 45.0, "synthesis PSNR {psnr:.2} dB regressed (expected ~58)");
|
|
}
|