zimage_stepcheck/
zimage_stepcheck.rs

1//! One mid-trajectory DiT step, device vs CPU, from traced latents.
2//!
3//! zimage_stepcheck <container.cmf> <prompt> <H> <W> <steps> <shift> <step i> <trace A> [trace B]
4//!
5//! Encodes the prompt (CPU text encoder), prepares the caption, then runs
6//! step i on `lat_i` of trace A (noise when i = 0 needs CMF_INIT_LATENT)
7//! on the device and on the CPU and prints rel(dev, cpu) and rel against
8//! trace A's `v_i`. With trace B it also runs the device on B's `lat_i`
9//! and prints how far the two device outputs are apart against how far
10//! the two inputs are — the trajectory's local sensitivity.
11use cortiq_engine::tokenizer::Tokenizer;
12use cortiq_engine::zimage::{self, ZImageDit, ZShape};
13use std::sync::Arc;
14
15fn rel(a: &[f32], b: &[f32]) -> f64 {
16    let (mut d, mut r) = (0f64, 0f64);
17    for (x, y) in a.iter().zip(b) {
18        d += (*x as f64 - *y as f64).powi(2);
19        r += (*y as f64).powi(2);
20    }
21    (d / r.max(1e-300)).sqrt()
22}
23
24fn read_f32(p: &str) -> Vec<f32> {
25    std::fs::read(p)
26        .unwrap_or_else(|e| panic!("{p}: {e}"))
27        .chunks_exact(4)
28        .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
29        .collect()
30}
31
32fn main() {
33    let a: Vec<String> = std::env::args().collect();
34    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
35    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
36    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
37    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
38    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
39    let cap = {
40        let _p = cortiq_engine::gpu::pause_gpu();
41        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
42    };
43    let dit = ZImageDit::from_cmf(&model).unwrap();
44    let sig = zimage::sigmas_torch_f32(steps, shift);
45    let t = zimage::t_model(sig[i]);
46    let mods = dit.mods_for_steps(&[t]);
47    let fs = dit.final_scale_for_steps(&[t]);
48    let shape = ZShape::new(hh, ww, ids.len());
49    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
50    println!("device prepared: {}", prep.device);
51    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
52    let lat = |dir: &str| -> Vec<f32> {
53        if i == 0 {
54            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
55        } else {
56            read_f32(&format!("{dir}/lat_{i}.f32"))
57        }
58    };
59    let xa = lat(&a[8]);
60    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
61    // forward against the two items stepped one by one.
62    if let Ok(neg) = std::env::var("ZC_NEG") {
63        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
64        let ncap = {
65            let _p = cortiq_engine::gpu::pause_gpu();
66            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
67        };
68        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
69        let tok_a = dit.tokens(&xa, &shape);
70        // `ZC_TAPS=<dir>`: every block's residual stream of the three
71        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
72        // then compared block by block (the pair's item rows vs the single).
73        let taps = std::env::var("ZC_TAPS").ok();
74        let set_taps = |sub: &str| {
75            if let Some(d) = &taps {
76                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
77            }
78        };
79        set_taps("pos");
80        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
81        set_taps("neg");
82        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
83        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
84        np.device = false;
85        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
86        np.device = true;
87        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
88        // (item 0 vs item 1) from a content one (the caption length).
89        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
90        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
91        let ok = dit.attach_device_pair(first, second, 3, None);
92        println!(
93            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
94            ids.len(),
95            nids.len(),
96            shape.n_img_p,
97            if swap { "neg,pos" } else { "pos,neg" }
98        );
99        set_taps("pair");
100        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
101            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
102            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
103            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
104                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
105        }
106        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
107        if let Some(d) = &taps {
108            // Row layout of the pair: image stage [item][n_img_p]; joint
109            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
110            let h = dit.cfg.dim;
111            let cp = |l: usize| l.div_ceil(32) * 32;
112            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
113            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
114            let nip = shape.n_img_p;
115            let rd = |p: String| -> Option<Vec<f32>> {
116                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
117            };
118            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
119            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
120            for n in names {
121                let (Some(sp), Some(sn), Some(pr)) = (
122                    rd(format!("{d}/pos/step{i}/{n}.f32")),
123                    rd(format!("{d}/neg/step{i}/{n}.f32")),
124                    rd(format!("{d}/pair/step{i}/{n}.f32")),
125                ) else {
126                    continue;
127                };
128                let (r0, r1) = if n.starts_with("nr") {
129                    ((0, nip), (nip, nip))
130                } else {
131                    ((0, nip + cp0), (nip + cp0, nip + cp1))
132                };
133                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
134                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
135                // image rows and caption rows separately
136                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
137                let (pi, pc) = split(ip);
138                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
139                println!(
140                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
141                    rel(&pi, &spi),
142                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
143                    rel(ineg, &sn[..ineg.len().min(sn.len())])
144                );
145            }
146        }
147        return;
148    }
149    let tok_a = dit.tokens(&xa, &shape);
150    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
151    let mut hp = zimage::ZPrepared { device: false, ..prep };
152    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
153    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
154    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
155        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
156    if let Some(b) = a.get(9) {
157        hp.device = true;
158        let xb = lat(b);
159        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
160        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
161            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
162    }
163}
zimage_stepcheck/zimage_stepcheck.rs

zimage_stepcheck/
zimage_stepcheck.rs
