zimage_vaecheck/
zimage_vaecheck.rs

1//! The resident wgpu VAE (`gpu::vae_decode_chain`) against the fp32 oracle.
2//!
3//! zimage_vaecheck <container.cmf> <oracle vae_*.safetensors> [reps]
4//!
5//! Decodes the oracle's `z` with `VaeDecoder::decode_fast` and prints rel
6//! against the oracle `img` (fp32 diffusers VAE), the u8 PSNR, and the
7//! in-process time of every repetition (the first one builds the weights
8//! and compiles the kernels).
9use std::collections::HashMap;
10
11fn read_st(path: &str) -> HashMap<String, (Vec<usize>, Vec<f32>)> {
12    let b = std::fs::read(path).unwrap_or_else(|e| panic!("{path}: {e}"));
13    let n = u64::from_le_bytes(b[..8].try_into().unwrap()) as usize;
14    let h: serde_json::Value = serde_json::from_slice(&b[8..8 + n]).unwrap();
15    let mut out = HashMap::new();
16    for (k, v) in h.as_object().unwrap() {
17        if v["dtype"].as_str() != Some("F32") {
18            continue;
19        }
20        let o = v["data_offsets"].as_array().unwrap();
21        let raw = &b[8 + n + o[0].as_u64().unwrap() as usize..8 + n + o[1].as_u64().unwrap() as usize];
22        let shape = v["shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
23        out.insert(k.clone(), (shape, raw.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect()));
24    }
25    out
26}
27
28fn main() {
29    let a: Vec<String> = std::env::args().collect();
30    let model = cortiq_core::CmfModel::open(&a[1]).unwrap();
31    let o = read_st(&a[2]);
32    let reps: usize = a.get(3).and_then(|v| v.parse().ok()).unwrap_or(3);
33    let (zs, z) = &o["z"];
34    let (_, img) = &o["img"];
35    let (h, w) = (zs[zs.len() - 2], zs[zs.len() - 1]);
36    let vae = cortiq_engine::vae::VaeDecoder::from_cmf(&model).unwrap();
37    for r in 0..reps {
38        let t = std::time::Instant::now();
39        let got = vae.decode_fast(z, h, w);
40        let dt = t.elapsed().as_secs_f64();
41        let (mut d, mut n, mut se) = (0f64, 0f64, 0f64);
42        for (x, y) in got.iter().zip(img) {
43            d += (*x as f64 - *y as f64).powi(2);
44            n += (*y as f64).powi(2);
45            let q = |v: f32| ((v / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round_ties_even() as f64;
46            se += (q(*x) - q(*y)).powi(2);
47        }
48        let psnr = 10.0 * (255f64.powi(2) / (se / got.len() as f64).max(1e-12)).log10();
49        println!("rep {r}: {:.3} s   img rel {:.3e}   u8 PSNR {:.2} dB", dt, (d / n).sqrt(), psnr);
50    }
51}
zimage_vaecheck/zimage_vaecheck.rs

zimage_vaecheck/
zimage_vaecheck.rs
