open-standards/wai/wai-rs/tests/int_hyper_synth.rs

//! The Rust integer decode of `wai.neural.int_hyper` — the pretrained mbt2018_mean
//! mean-scale hyperprior (CompressAI zoo, quality 6), exported to integers —
//! reproduces the Python export tool's reconstruction BYTE-FOR-BYTE, and costs at
//! most 1 dB against the same model's float decode on the photograph the vectors
//! were built from. The decode is all-integer by construction; its
//! cross-architecture byte-equality is checked on the committed tiny vector
//! (`tests/vectors/int_hyper/`, `wai-byte-equality.yml`), not on these. The
//! pretrained-model tests below take their expected hashes from the export run
//! that wrote the vector directory, not from the repository (only the tiny
//! vector's test reads a committed hash). That export encodes with the float model,
//! and its files are not portable: the synthesis-only test's latent and its
//! `rgb_hash` were measured to differ between an arm64 macOS machine and an x86_64
//! Linux runner, and two runs on that runner wrote different `cond.json` and
//! `int_synth_sota.json` (most likely in `psnr_float`, which only the quality
//! test reads, against thresholds). None of these hashes is pinned for use by
//! another run (`tools/wai_sota_vectors.sh`).
//!
//! The model + bitstream + expected FNV-1a checksums are regenerated by
//! `tools/wai_sota_vectors.sh` (pretrained mbt2018_mean exported to integers by
//! `tools/wai_sota_int_synth_export.py`) into `WAI_SOTA_INT_DIR` (default
//! `/tmp/sota_int`). CI (`wai-sota-vectors.yml`) regenerates them from a pinned
//! public-domain photograph and runs this file with `WAI_REQUIRE_SOTA_VECTORS=1`,
//! so an absent vector fails there; a local run without them skips loudly. A
//! `WAI_SOTA_INT_DIR` naming a directory that does not exist, or one missing a
//! file a test reads, fails either way: only the default may be absent and skip.
//!
//! Quality is image-dependent: the 42.6 dB development figure for this codec was
//! measured on a photograph that is not in the repository. What is asserted here,
//! on whatever image produced the vectors, is the integer decode's cost against the
//! float codec on that same image. Nothing here compares the model with any other
//! codec.
#![cfg(feature = "neural_int")]

#[path = "support/require_vectors.rs"]
mod require_vectors;

use require_vectors::{load_all, load_json, InputDir};
use wai::int_hyper_synth::{fnv1a, IntModel};

const PRODUCE: &str = "tools/wai_sota_vectors.sh <dir> (then WAI_SOTA_INT_DIR=<dir>, default /tmp/sota_int)";

/// `WAI_SOTA_INT_DIR`, else `/tmp/sota_int`. Set to a directory that does not
/// exist, or to one missing a file the test reads, it fails the test instead of
/// letting it skip.
fn dir() -> InputDir {
    require_vectors::env_dir("int_hyper_synth", "WAI_SOTA_INT_DIR", "/tmp/sota_int", PRODUCE)
}

/// Every named file from the vector dir, or `None` after reporting what is absent.
fn inputs<const N: usize>(test: &str, d: &InputDir, names: [&str; N]) -> Option<[serde_json::Value; N]> {
    load_all(test, d, names, PRODUCE)
}

#[test]
fn rust_integer_synthesis_reproduces_the_export_hash() {
    let d = dir();
    let Some([conf, model, input]) = inputs(
        "int_hyper_synth::rust_integer_synthesis_reproduces_the_export_hash",
        &d,
        ["int_synth_sota.json", "int_synth_sota_model.json", "int_synth_sota_input.json"],
    ) else {
        return;
    };

    let (c, h, w) = (
        input["c"].as_u64().unwrap() as usize,
        input["h"].as_u64().unwrap() as usize,
        input["w"].as_u64().unwrap() as usize,
    );
    let latent: Vec<i64> =
        input["data"].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect();

    let model = IntModel::from_json(&model);
    // The synthesis through its lowering (whm1/1 with the RGB8 step), as the decode a sink
    // runs computes it, row-major HWC.
    let t = through_the_lowering(&model, true, latent, c, h, w);
    let (oh, ow) = (t.h as usize, t.w as usize);
    assert_eq!((oh, ow), (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize));
    assert_eq!(t.c as u64, conf["out_c"].as_u64().unwrap());
    let plane = oh * ow;
    let mut rgb = vec![0u8; t.data.len()];
    for (k, &v) in t.data.iter().enumerate() {
        rgb[(k % plane) * t.c as usize + k / plane] = v as u8;
    }
    let got = format!("{:016x}", fnv1a(&rgb));
    let want = conf["rgb_hash"].as_str().unwrap();
    println!("int_hyper integer synthesis: {oh}x{ow}, fnv1a {got} (psnr {:.2} dB)", conf["psnr_int"].as_f64().unwrap_or(0.0));
    assert_eq!(got, want, "the integer synthesis through the integer operator set must reproduce the export's reconstruction byte-for-byte");
}

/// One network of the pretrained set through its lowering `whm1/1` (with the RGB8 step
/// where `rgb8` is set), loaded under `acc64` and run by the integer operator set
/// (`int_ops`, its fast kernels behind it) on `x`, `[c, h, w]`.
fn through_the_lowering(m: &IntModel, rgb8: bool, x: Vec<i64>, c: usize, h: usize, w: usize) -> wai_ops::Tensor {
    let g = wai_ops::lower::whm1_net(&wai::int_ops::whm1_net(m).expect("the model converts"), rgb8).expect("the network lowers");
    let l = wai::int_ops::load(&g).expect("the network loads under acc64");
    wai::int_ops::run1(&l, wai_ops::Tensor::new(c as u32, h as u32, w as u32, x), &[]).expect("the input is within its check range")
}

#[test]
fn rust_integer_hyper_decoder_reproduces_the_export_hash() {
    let d = dir();
    let Some([conf, model, input]) = inputs(
        "int_hyper_synth::rust_integer_hyper_decoder_reproduces_the_export_hash",
        &d,
        ["int_hs.json", "int_hs_model.json", "int_hs_input.json"],
    ) else {
        return;
    };
    let (c, h, w) = (
        input["c"].as_u64().unwrap() as usize,
        input["h"].as_u64().unwrap() as usize,
        input["w"].as_u64().unwrap() as usize,
    );
    let z: Vec<i64> = input["data"].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect();
    let model = IntModel::from_json(&model);
    // The hyper-decoder through its lowering (whm1/1), as the decode a sink runs.
    let t = through_the_lowering(&model, false, z, c, h, w);
    let (oc, oh, ow) = (t.c as usize, t.h as usize, t.w as usize);
    assert_eq!((oc, oh, ow), (
        conf["out_c"].as_u64().unwrap() as usize,
        conf["out_h"].as_u64().unwrap() as usize,
        conf["out_w"].as_u64().unwrap() as usize));
    let bytes: Vec<u8> = t.data.iter().flat_map(|v| v.to_le_bytes()).collect();
    let got = format!("{:016x}", fnv1a(&bytes));
    println!("integer h_s: z_hat {c}x{h}x{w} -> sigma||mu {oc}x{oh}x{ow}, fnv1a {got}");
    assert_eq!(got, conf["params_hash"].as_str().unwrap(), "h_s through the integer operator set must reproduce the export byte-for-byte");
}

#[test]
fn decodes_the_committed_tiny_conformance_vector() {
    // A SMALL synthetic int_hyper codec, committed under tests/vectors/int_hyper/,
    // drives the full decode path (factorized z → integer h_s → integer-σ entropy
    // → integer g_s) with NO external artifacts — so CI verifies the int_hyper
    // capability without the 23 MB pretrained model. Both the JSON and the compact
    // binary model must reproduce the committed FNV-1a, byte-for-byte. Never skips.
    use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
    let d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_hyper");
    let conf = load_json(&d.join("tiny.json")).expect("committed tiny.json");
    let ent = CondEntropy::from_value(&load_json(&d.join("tiny_entropy.json")).unwrap());
    let bs = CondBitstream::from_value(&load_json(&d.join("tiny_bitstream.json")).unwrap());
    let want = conf["rgb_hash"].as_str().unwrap();
    let (oh, ow) = (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize);

    let g_j = IntModel::from_json(&load_json(&d.join("tiny_g_s.json")).unwrap());
    let h_j = IntModel::from_json(&load_json(&d.join("tiny_h_s.json")).unwrap());
    let (rgb, h, w) = decode_conditional_checked(&g_j, &h_j, &ent, &bs.stream()).expect("tiny json decode");
    assert_eq!((h, w), (oh, ow));
    assert_eq!(format!("{:016x}", fnv1a(&rgb)), want, "tiny int_hyper (json model) must reproduce the committed hash");

    let g_b = IntModel::from_bytes(&std::fs::read(d.join("tiny_g_s.bin")).unwrap()).expect("tiny g_s bin");
    let h_b = IntModel::from_bytes(&std::fs::read(d.join("tiny_h_s.bin")).unwrap()).expect("tiny h_s bin");
    let (rgb2, _, _) = decode_conditional_checked(&g_b, &h_b, &ent, &bs.stream()).expect("tiny binary decode");
    assert_eq!(format!("{:016x}", fnv1a(&rgb2)), want, "tiny int_hyper (binary model) must reproduce the committed hash");
    println!("tiny int_hyper conformance: {oh}x{ow}, {} z + {} y rANS bytes, fnv1a {want}", bs.z_rans.len(), bs.rans.len());
}

#[test]
fn rust_decodes_the_full_conditional_bitstream_to_the_export_image() {
    // The whole NeuralIntegerExact codec, end to end: a real bitstream (factorized
    // z + rANS-coded y residuals) decodes — factorized z → integer h_s → σ,μ;
    // σ → integer bucket → rANS-decode r → y_hat = r·2^afrac + μ → integer g_s —
    // to the SAME image the export tool produced, byte-for-byte (fully integer, so
    // identical on every machine). This is the integer-entropy contract assembled.
    use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
    let d = dir();
    let test = "int_hyper_synth::rust_decodes_the_full_conditional_bitstream_to_the_export_image";
    let Some([conf, g_s, h_s, ent, bs]) = inputs(
        test,
        &d,
        ["cond.json", "int_synth_sota_model.json", "int_hs_model.json", "cond_entropy.json", "cond_bitstream.json"],
    ) else {
        return;
    };
    let (g_s, h_s) = (IntModel::from_json(&g_s), IntModel::from_json(&h_s));
    let ent = CondEntropy::from_value(&ent);
    let bs = CondBitstream::from_value(&bs);
    let (rgb, oh, ow) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).expect("conditional decode");
    assert_eq!((oh, ow), (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize));
    let got = format!("{:016x}", fnv1a(&rgb));
    println!(
        "conditional bitstream: {} z + {} y rANS bytes -> {oh}x{ow}, fnv1a {got} ({:.2} dB)",
        bs.z_rans.len(), bs.rans.len(), conf["psnr_cond"].as_f64().unwrap_or(0.0)
    );
    assert_eq!(got, conf["rgb_hash"].as_str().unwrap(), "Rust must decode the bitstream to the export image byte-for-byte");

    // The COMPACT BINARY model (~½ the size) must decode bit-identically to the
    // JSON — same codec, smaller download (what makes the in-browser demo practical).
    // The export writes both, so an absent binary is an incomplete vector dir.
    let (gb, hb) = (d.join("int_synth_sota_model.bin"), d.join("int_hs_model.bin"));
    if !gb.exists() || !hb.exists() {
        d.absent(&format!("{test} (binary model)"), &format!("int_synth_sota_model.bin / int_hs_model.bin not in {}", d.display()), PRODUCE);
    } else {
        let (g_bin, h_bin) = (std::fs::read(&gb).unwrap(), std::fs::read(&hb).unwrap());
        let g_s2 = IntModel::from_bytes(&g_bin).expect("g_s binary");
        let h_s2 = IntModel::from_bytes(&h_bin).expect("h_s binary");
        let (rgb2, _, _) = decode_conditional_checked(&g_s2, &h_s2, &ent, &bs.stream()).expect("binary conditional decode");
        assert_eq!(format!("{:016x}", fnv1a(&rgb2)), conf["rgb_hash"].as_str().unwrap(),
            "the compact binary model must decode byte-identically to the JSON model");
        println!("compact binary model ({} + {} bytes) decodes byte-identically", g_bin.len(), h_bin.len());
        // Through the integer operator set (`int_hyper_synth::ops`: both networks lowered by
        // whm1/1, loaded under acc64 and run with their check ranges): the same image.
        let (rgb3, ..) = wai::int_hyper_synth::ops::decode(&g_s2, &h_s2, &ent, &bs.stream()).expect("the decode through the lowering");
        assert_eq!(format!("{:016x}", fnv1a(&rgb3)), conf["rgb_hash"].as_str().unwrap(),
            "the decode through the integer operator set must give the export image");
        println!("the decode through the integer operator set gives the same image");
    }
}

/// Gate 1 of the migration to the integer operator set (`extensions/integer-operators.md`
/// §13.2), for the pretrained set, which is too large to commit: the regenerated set is
/// the one whose certificate is committed by its digest
/// (`ops-conformance/certificates/external.json`), and this crate's lowering of the
/// models it parses (`int_hyper_synth::ops`) writes the two graphs that certificate
/// records, `h_s` then `g_s`, which load under `acc64` and not `acc32`.
/// (`wai-sota-vectors.yml` checks the certificate itself against the set, with both
/// certifiers.)
#[test]
fn the_pretrained_set_lowers_to_its_certificates_graphs() {
    let d = dir();
    let test = "int_hyper_synth::the_pretrained_set_lowers_to_its_certificates_graphs";
    let names = ["cond_entropy.json", "int_synth_sota_model.bin", "int_hs_model.bin"];
    let missing: Vec<&str> = names.iter().copied().filter(|n| !d.join(n).exists()).collect();
    if !missing.is_empty() {
        d.absent(test, &format!("{} not in {}", missing.join(", "), d.display()), PRODUCE);
        return;
    }
    let [e, g, h] = names.map(|n| std::fs::read(d.join(n)).unwrap());
    let digest = wai_ops::generate::sets::set_digest(&[("entropy.json", e), ("g_s.whm", g.clone()), ("h_s.whm", h.clone())]);
    let certs = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance/certificates");
    let ext = load_json(&certs.join("external.json")).expect("the external certificates' listing");
    let listed = ext["sets"].as_array().unwrap().iter().find(|s| s["capability"] == "wai.neural.int_hyper").expect("the pretrained set is listed");
    assert_eq!(listed["set_sha256"].as_str(), Some(digest.as_str()), "the regenerated set is the one whose certificate is committed");
    let cert = load_json(&certs.join(format!("wai.neural.int_hyper/{digest}.json"))).expect("its certificate");
    let (g_s, h_s) = (IntModel::from_bytes(&g).unwrap(), IntModel::from_bytes(&h).unwrap());
    let (hg, gg) = wai::int_hyper_synth::ops::lower(&g_s, &h_s).expect("the pretrained set lowers");
    for (k, (role, graph)) in [("h_s", &hg), ("g_s", &gg)].into_iter().enumerate() {
        let sha: String = wai_ops::sha256::sha256(&graph.to_bytes()).iter().map(|x| format!("{x:02x}")).collect();
        assert_eq!((cert["graphs"][k]["name"].as_str(), cert["graphs"][k]["graph_sha256"].as_str()), (Some(role), Some(sha.as_str())), "{role}");
    }
    let nets = wai::int_hyper_synth::ops::load(&g_s, &h_s).expect("the pretrained set loads under acc64");
    assert!(!nets.h_s.meets(wai_ops::Profile::Acc32) && !nets.g_s.meets(wai_ops::Profile::Acc32), "the pretrained set meets acc64, not acc32");
    println!("pretrained set {digest}: lowers to its certificate's graphs and loads under acc64");
}

/// The largest quality cost the integer decode may have against the float codec on
/// the same image. Measured costs: 0.19 dB for the full conditional decode on the
/// pinned photograph (arm64 macOS machine); 0.62 dB for the integer synthesis alone
/// in development, on a photograph that is not in the repository. A larger gap
/// means the quantization has regressed, even when every byte-exactness check
/// above still passes (they compare the integer decode only with itself).
const MAX_INTEGER_COST_DB: f64 = 1.0;
/// A sanity floor for the float codec on the source image: below it the pretrained
/// checkpoint is not coding natural content at all, and a matching integer decode
/// of it proves nothing. A sanity floor only, not a quality claim.
const MIN_FLOAT_PSNR_DB: f64 = 30.0;

#[test]
fn integer_decode_keeps_the_float_codec_quality() {
    // Measured by this decoder on its own output, against the exact source pixels
    // the export tool coded — not the PSNR the tool reports about itself.
    use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
    let d = dir();
    let test = "int_hyper_synth::integer_decode_keeps_the_float_codec_quality";
    let Some([conf, g_s, h_s, ent, bs]) = inputs(
        test,
        &d,
        ["cond.json", "int_synth_sota_model.json", "int_hs_model.json", "cond_entropy.json", "cond_bitstream.json"],
    ) else {
        return;
    };
    let src_path = d.join("source.rgb");
    let Some(psnr_float) = conf.get("psnr_float").and_then(|v| v.as_f64()) else {
        d.absent(test, "cond.json predates psnr_float (regenerate the vectors)", PRODUCE);
        return;
    };
    if !src_path.exists() {
        d.absent(test, &format!("source.rgb not in {}", d.display()), PRODUCE);
        return;
    }
    let source = std::fs::read(&src_path).unwrap();

    let (g_s, h_s) = (IntModel::from_json(&g_s), IntModel::from_json(&h_s));
    let (rgb, h, w) = decode_conditional_checked(&g_s, &h_s, &CondEntropy::from_value(&ent), &CondBitstream::from_value(&bs).stream())
        .expect("conditional decode");
    assert_eq!(source.len(), rgb.len(), "source.rgb is not the {h}x{w} RGB image the bitstream decodes to");

    let sse: f64 = rgb.iter().zip(&source).map(|(a, b)| (f64::from(*a) - f64::from(*b)).powi(2)).sum();
    let mse = sse / rgb.len() as f64;
    let psnr = 10.0 * (255.0f64 * 255.0 / mse).log10();
    let cost = psnr_float - psnr;
    println!("integer decode {psnr:.3} dB vs float codec {psnr_float:.3} dB on the same {h}x{w} image (cost {cost:.3} dB)");
    assert!(psnr_float >= MIN_FLOAT_PSNR_DB, "float codec {psnr_float:.2} dB < {MIN_FLOAT_PSNR_DB} dB: the checkpoint or source is wrong");
    assert!(cost <= MAX_INTEGER_COST_DB, "integer decode costs {cost:.2} dB against the float codec (limit {MAX_INTEGER_COST_DB} dB)");
}