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

//! The Rust integer synthesis of a SOTA-class learned codec reproduces the
//! Python export tool's reconstruction BYTE-FOR-BYTE — proving the decode is
//! cross-language / cross-hardware exact at SOTA quality (mbt2018_mean, 42.6 dB,
//! fully integer). The model + input + expected FNV-1a checksum are emitted by
//! `tools/wai_sota_int_synth_export.py` into a directory pointed to by
//! `WAI_SOTA_INT_DIR` (default `/tmp/sota_int`); skipped if absent.
#![cfg(feature = "neural_int")]

use std::path::PathBuf;

use wai::int_hyper_synth::{decode, fnv1a, IntModel};

fn dir() -> PathBuf {
    PathBuf::from(std::env::var("WAI_SOTA_INT_DIR").unwrap_or_else(|_| "/tmp/sota_int".into()))
}

fn jload(p: &std::path::Path) -> Option<serde_json::Value> {
    serde_json::from_slice(&std::fs::read(p).ok()?).ok()
}

#[test]
fn rust_integer_synthesis_reproduces_the_export_hash() {
    let d = dir();
    let (Some(conf), Some(model), Some(input)) = (
        jload(&d.join("int_synth_sota.json")),
        jload(&d.join("int_synth_sota_model.json")),
        jload(&d.join("int_synth_sota_input.json")),
    ) else {
        eprintln!("skipping SOTA int-synth: artifacts not in {} (run tools/wai_sota_int_synth_export.py)", d.display());
        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!("SOTA 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), Some(model), Some(input)) = (
        jload(&d.join("int_hs.json")),
        jload(&d.join("int_hs_model.json")),
        jload(&d.join("int_hs_input.json")),
    ) else {
        eprintln!("skipping SOTA int h_s: artifacts not in {}", d.display());
        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 SOTA 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 = jload(&d.join("tiny.json")).expect("committed tiny.json");
    let ent = CondEntropy::from_value(&jload(&d.join("tiny_entropy.json")).unwrap());
    let bs = CondBitstream::from_value(&jload(&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(&jload(&d.join("tiny_g_s.json")).unwrap());
    let h_j = IntModel::from_json(&jload(&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 (Some(conf), Some(g_s), Some(h_s), Some(ent), Some(bs)) = (
        jload(&d.join("cond.json")),
        jload(&d.join("int_synth_sota_model.json")),
        jload(&d.join("int_hs_model.json")),
        jload(&d.join("cond_entropy.json")),
        jload(&d.join("cond_bitstream.json")),
    ) else {
        eprintln!("skipping conditional bitstream: artifacts not in {}", d.display());
        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).
    if let (Ok(g_bin), Ok(h_bin)) =
        (std::fs::read(d.join("int_synth_sota_model.bin")), std::fs::read(d.join("int_hs_model.bin")))
    {
        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());
    }
}