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::{decode, 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 (rgb, oh, ow) = decode(&IntModel::from_json(&model), &latent, c, h, w);
    assert_eq!((oh, ow), (conf["out_h"].as_u64().unwrap() as usize, conf["out_w"].as_u64().unwrap() as usize));
    assert_eq!(rgb.len(), oh * ow * conf["out_c"].as_u64().unwrap() as usize);

    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, "Rust integer synthesis must reproduce the export's reconstruction byte-for-byte");
}

#[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 (out, oc, oh, ow) = wai::int_hyper_synth::run_net(&IntModel::from_json(&model), &z, c, h, w);
    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> = out.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(), "integer h_s 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, 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(&g_j, &h_j, &ent, &bs).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(&g_b, &h_b, &ent, &bs).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, 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(&g_s, &h_s, &ent, &bs).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(&g_s2, &h_s2, &ent, &bs).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());
    }
}

/// 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, 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(&g_s, &h_s, &CondEntropy::from_value(&ent), &CondBitstream::from_value(&bs))
        .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)");
}