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

// wai-conformance: int-kernels
//! `wai.generate.spectral` — **generation-equivalence**: a deterministic integer
//! generator produces byte-identical bytes on every machine from the same
//! `(model, prompt, seed, steps)`, so its hash is portable (the create-side of
//! "conformance is a hash"). Committed vectors under `tests/vectors/generate/`:
//! `generate.json` (96×96), `sizes.json` (256×128 and 640×320), and `instr-v0.json`,
//! the instruction object of form v0 (SPEC §5) with the SHA-256 of each of its four
//! listed sizes. Nothing skips except the two larger listed sizes in a debug build,
//! which the release step of the conformance workflow runs; the objects a sink
//! refuses are run through `wai` in `tests/cli_generate_score.rs`.
#![cfg(feature = "generate")]

use std::path::PathBuf;

use wai::generate::{fnv1a, generate, SpectralModel};

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

fn dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/generate")
}
fn jload(name: &str) -> serde_json::Value {
    serde_json::from_slice(&std::fs::read(dir().join(name)).unwrap()).unwrap()
}
fn hash(rgb: &[u8]) -> String {
    format!("{:016x}", fnv1a(rgb))
}

#[test]
fn generation_is_deterministic_and_pinned() {
    let m = SpectralModel::from_value(&jload("generate_spectral_model.json"));
    let pin = jload("generate.json");
    let (prompt, seed, steps, h, w) = (
        pin["prompt"].as_str().unwrap(),
        pin["seed"].as_u64().unwrap(),
        pin["steps"].as_u64().unwrap() as usize,
        pin["h"].as_u64().unwrap() as usize,
        pin["w"].as_u64().unwrap() as usize,
    );

    // 1. PIN: the committed (prompt, seed, steps) reproduces the committed hash,
    //    byte-for-byte — the portable receipt of a generated image.
    let rgb = generate(&m, prompt, seed, steps, h, w);
    assert_eq!(rgb.len(), h * w * 3);
    let got = hash(&rgb);
    println!("generate.spectral: {h}x{w} '{prompt}' seed={seed} steps={steps} -> fnv1a {got}");
    assert_eq!(got, pin["rgb_hash"].as_str().unwrap(), "pinned generation must reproduce the committed hash");
    sha256_sidecar::one(&dir(), "generate", "RGB8: rows top to bottom, pixels left to right, R, G, B bytes", &rgb);

    // 2. DETERMINISM: regenerating the same inputs is byte-identical.
    assert_eq!(generate(&m, prompt, seed, steps, h, w), rgb, "regeneration must be byte-identical");

    // 3. SENSITIVITY: a different prompt or seed draws a different image (the
    //    generation is conditioned, not constant).
    assert_ne!(hash(&generate(&m, "a different prompt", seed, steps, h, w)), got, "prompt must change the output");
    assert_ne!(hash(&generate(&m, prompt, seed ^ 1, steps, h, w)), got, "seed must change the output");
    assert_ne!(hash(&generate(&m, prompt, seed, steps + 1, h, w)), got, "steps must change the output");
}

fn model() -> SpectralModel {
    SpectralModel::from_value(&jload("generate_spectral_model.json"))
}

/// `sizes.json`: the committed vector's prompt, seed and steps at 256×128 and 640×320,
/// non-square, so the `(h, w)` order of `generate` is pinned too. Each SHA-256 was
/// made natively by `wai generate` and equals the shared browser binding's output.
#[test]
fn sizes_reproduce_their_sha256() {
    let (m, pin) = (model(), jload("sizes.json"));
    let model_bytes = std::fs::read(dir().join("generate_spectral_model.json")).unwrap();
    assert_eq!(sha256_sidecar::sha256_hex(&model_bytes), pin["model_sha256"].as_str().unwrap(), "the committed model");
    let (prompt, seed, steps) =
        (pin["prompt"].as_str().unwrap(), pin["seed"].as_u64().unwrap(), pin["steps"].as_u64().unwrap() as usize);
    let sizes = pin["sizes"].as_array().unwrap();
    assert_eq!(sizes.len(), 2);
    for s in sizes {
        let (w, h) = (s["w"].as_u64().unwrap() as usize, s["h"].as_u64().unwrap() as usize);
        assert_ne!(w, h, "a non-square size pins the order of w and h");
        let rgb = generate(&m, prompt, seed, steps, h, w);
        assert_eq!(rgb.len(), w * h * 3);
        assert_eq!(sha256_sidecar::sha256_hex(&rgb), s["sha256"].as_str().unwrap(), "{w}x{h}");
        // h and w swapped make another picture.
        assert_ne!(generate(&m, prompt, seed, steps, w, h), rgb, "{w}x{h}: swapping w and h");
    }
}

/// The members of the vector's instruction object, read with serde_json, and the
/// object written again in form v0's canonical serialization (SPEC §5, *Payload of
/// `wai.generate.spectral`*), stated here independently of `wai`'s reader.
fn instr_v0() -> (serde_json::Value, Vec<(usize, usize, String)>) {
    let v = jload("instr-v0.json");
    let object = v["object"].as_str().unwrap();
    let o: serde_json::Value = serde_json::from_str(object).unwrap();
    let sizes: Vec<String> = o["sizes"].as_array().unwrap().iter()
        .map(|p| format!("[{},{}]", p[0].as_u64().unwrap(), p[1].as_u64().unwrap())).collect();
    let canonical = format!(
        "{{\"model\":\"{}\",\"prompt\":\"{}\",\"seed\":{},\"sizes\":[{}],\"steps\":{}}}",
        o["model"].as_str().unwrap(), o["prompt"].as_str().unwrap(), o["seed"].as_u64().unwrap(),
        sizes.join(","), o["steps"].as_u64().unwrap());
    assert_eq!(canonical, object, "the vector's object is in canonical form");
    let outputs = v["outputs"].as_array().unwrap().iter().map(|e| (
        e["w"].as_u64().unwrap() as usize, e["h"].as_u64().unwrap() as usize, e["sha256"].as_str().unwrap().to_string(),
    )).collect();
    (o, outputs)
}

fn make_listed(min_pixels: usize, max_pixels: usize) -> usize {
    let (o, outputs) = instr_v0();
    let model_bytes = std::fs::read(dir().join(jload("instr-v0.json")["model_file"].as_str().unwrap())).unwrap();
    assert_eq!(format!("sha256:{}", sha256_sidecar::sha256_hex(&model_bytes)), o["model"].as_str().unwrap(),
               "the object names the committed model by its SHA-256");
    let listed: Vec<(usize, usize)> = o["sizes"].as_array().unwrap().iter()
        .map(|p| (p[0].as_u64().unwrap() as usize, p[1].as_u64().unwrap() as usize)).collect();
    assert_eq!(listed, outputs.iter().map(|(w, h, _)| (*w, *h)).collect::<Vec<_>>(), "one output per listed size");
    let m = SpectralModel::from_value(&serde_json::from_slice(&model_bytes).unwrap());
    let (prompt, seed, steps) =
        (o["prompt"].as_str().unwrap(), o["seed"].as_u64().unwrap(), o["steps"].as_u64().unwrap() as usize);
    let mut made = 0;
    for (w, h, sha) in &outputs {
        if !(min_pixels..=max_pixels).contains(&(w * h)) {
            continue;
        }
        let rgb = generate(&m, prompt, seed, steps, *h, *w);
        assert_eq!(rgb.len(), w * h * 3);
        assert_eq!(&sha256_sidecar::sha256_hex(&rgb), sha, "{w}x{h}");
        made += 1;
    }
    made
}

/// `instr-v0.json` at its two smaller listed sizes. Its 256×128 and 640×320 outputs are
/// `sizes.json`'s: the object's prompt, seed and steps are the committed vector's.
#[test]
fn instruction_object_v0_small_sizes_reproduce() {
    assert_eq!(make_listed(0, 640 * 320), 2);
    let (_, outputs) = instr_v0();
    let sizes = jload("sizes.json");
    for (s, (w, h, sha)) in sizes["sizes"].as_array().unwrap().iter().zip(&outputs) {
        assert_eq!((s["w"].as_u64().unwrap() as usize, s["h"].as_u64().unwrap() as usize, s["sha256"].as_str().unwrap()),
                   (*w, *h, sha.as_str()));
    }
}

/// `instr-v0.json` at 1280×720 and 1920×1080: under a second in a release build, many
/// seconds in a debug one, so the release step of the conformance workflow runs it.
#[test]
#[cfg_attr(debug_assertions, ignore = "1280x720 and 1920x1080 take seconds in a debug build; the release step runs them")]
fn instruction_object_v0_large_sizes_reproduce() {
    assert_eq!(make_listed(640 * 320 + 1, usize::MAX), 2);
}

/// SPEC §5, *Payload of `wai.generate.spectral`*, stated here a second time and
/// independently of `wai`'s reader: the bytes are one JSON object with exactly the
/// five members, each value in range, and the bytes are exactly that object written
/// canonically (members in ascending byte order, no whitespace, integers in decimal,
/// the prompt unescaped). A member given twice fails the last test, since the object
/// written again names each member once. `Some(object)` when the bytes are form v0.
fn form_v0(bytes: &[u8]) -> Option<serde_json::Value> {
    let o: serde_json::Value = serde_json::from_slice(bytes).ok()?;
    let m = o.as_object()?;
    let keys: Vec<&str> = m.keys().map(String::as_str).collect();
    let mut sorted = keys.clone();
    sorted.sort();
    if sorted != ["model", "prompt", "seed", "sizes", "steps"] {
        return None;
    }
    let hex = o["model"].as_str()?.strip_prefix("sha256:")?;
    let hex_ok = hex.len() == 64 && hex.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b));
    let prompt = o["prompt"].as_str()?;
    let prompt_ok = prompt.len() <= 256 && prompt.bytes().all(|b| (0x20..=0x7e).contains(&b) && b != b'"' && b != b'\\');
    let seed = o["seed"].as_u64()?;
    let steps = o["steps"].as_u64()?;
    let sizes: Vec<(u64, u64)> = o["sizes"].as_array()?.iter()
        .map(|p| p.as_array().filter(|a| a.len() == 2).and_then(|a| Some((a[0].as_u64()?, a[1].as_u64()?))))
        .collect::<Option<_>>()?;
    let sizes_ok = (1..=16).contains(&sizes.len())
        && sizes.iter().all(|&(w, h)| (1..=4096).contains(&w) && (1..=4096).contains(&h) && w * h <= 1 << 23)
        && sizes.windows(2).all(|p| p[0] < p[1]);
    if !(hex_ok && prompt_ok && seed <= u32::MAX as u64 && (1..=64).contains(&steps) && sizes_ok) {
        return None;
    }
    let listed: Vec<String> = sizes.iter().map(|(w, h)| format!("[{w},{h}]")).collect();
    let canonical = format!("{{\"model\":\"sha256:{hex}\",\"prompt\":\"{prompt}\",\"seed\":{seed},\"sizes\":[{}],\"steps\":{steps}}}",
                            listed.join(","));
    (canonical.as_bytes() == bytes).then_some(o)
}

/// `instr-v0.json`, read by the statement above: its object and every object it lists
/// as accepted (one at each inclusive bound) are form v0, and each accepted object's
/// first listed size, made by the reference generator, has the SHA-256 the vector
/// states; every object it lists as refused is not form v0. `wai`'s own reader is held
/// to the same lists, with the reason for each refusal, in `tests/cli_generate_score.rs`.
#[test]
fn instruction_object_v0_bounds_and_refusals() {
    let v = jload("instr-v0.json");
    let model_bytes = std::fs::read(dir().join(v["model_file"].as_str().unwrap())).unwrap();
    let m = SpectralModel::from_value(&serde_json::from_slice(&model_bytes).unwrap());
    let object = form_v0(v["object"].as_str().unwrap().as_bytes()).expect("the vector's object is form v0");
    let accepted = v["accepted"].as_array().unwrap();
    assert!(accepted.len() >= 12);
    for c in accepted {
        let why = c["why"].as_str().unwrap();
        let o = form_v0(c["object"].as_str().unwrap().as_bytes()).unwrap_or_else(|| panic!("{why}: not form v0"));
        assert_eq!(o["model"], object["model"], "{why}: the committed model");
        let first = &o["sizes"][0];
        for e in c["outputs"].as_array().unwrap() {
            let (w, h) = (e["w"].as_u64().unwrap() as usize, e["h"].as_u64().unwrap() as usize);
            assert_eq!((first[0].as_u64(), first[1].as_u64()), (Some(w as u64), Some(h as u64)), "{why}: the first listed size");
            let rgb = generate(&m, o["prompt"].as_str().unwrap(), o["seed"].as_u64().unwrap(), o["steps"].as_u64().unwrap() as usize, h, w);
            assert_eq!(rgb.len(), w * h * 3);
            assert_eq!(&sha256_sidecar::sha256_hex(&rgb), e["sha256"].as_str().unwrap(), "{why}: {w}x{h}");
        }
    }
    let refused = v["refused"].as_array().unwrap();
    assert!(refused.len() >= 30);
    for c in refused {
        let why = c["why"].as_str().unwrap();
        assert!(form_v0(c["object"].as_str().unwrap().as_bytes()).is_none(), "{why}: the statement of form v0 accepts it");
    }
}