Struct ZShape

Source

pub struct ZShape {
    pub h_lat: usize,
    pub w_lat: usize,
    pub grid: (usize, usize),
    pub n_img: usize,
    pub n_img_p: usize,
    pub l: usize,
    pub l_p: usize,
}

Expand description

Token geometry of one (resolution, prompt length).

Fields§

§h_lat: usize

Latent size (H/8, W/8).

§w_lat: usize§grid: (usize, usize)

Patch grid (H/16, W/16).

§n_img: usize

n_img = grid.0 · grid.1, n_img_p = ceil32(n_img).

§n_img_p: usize§l: usize

Caption tokens L and ceil32(L).

§l_p: usize

Implementations§

Source §

impl ZShape

Source

pub fn new(height: usize, width: usize, l: usize) -> Self

From the image size in pixels (multiples of 16) and the caption length.

Examples found in repository ?

examples/zimage_devcheck.rs (line 56)

44fn main() {
45    unsafe { std::env::set_var("CMF_GPU", "1") };
46    let a: Vec<String> = std::env::args().collect();
47    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
48    let dit = ZImageDit::from_cmf(&model).unwrap();
49    let od = std::path::Path::new(&a[2]);
50    let case = &a[3];
51    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
52    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
53    let t = o["t_model"][0];
54    let cap = &o["cap"];
55    let x_in = &o["x_in"];
56    let shape = ZShape::new(hh, ww, l);
57    let mods = dit.mods_for_steps(&[t]);
58    let fs = dit.final_scale_for_steps(&[t]);
59    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
60    // caption: device-refined (prepare) vs CPU-refined
61    let prep = dit.prepare(cap, shape, 1, None).unwrap();
62    println!("device prepared: {}", prep.device);
63    let mut cap_cpu = dit.embed_caption(cap, l);
64    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
65    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
66        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
67    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
68    let x_tok = zimage::pad_rows_repeat_last(&zimage::patchify(x_in, c, lh, lw), shape.n_img, shape.n_img_p, dit.geom().patch_dim);
69    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
70    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
71    // Replays must be stateless: the same inputs again, and a different
72    // step's mods in between.
73    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
74    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
75    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
76    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
77    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
78    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
79    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
80    let mut prep_cpu = prep;
81    prep_cpu.device = false;
82    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
83    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
84    for (n, v) in &taps {
85        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
86        if let Ok(b) = std::fs::read(&p) {
87            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
88            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
89            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
90        }
91    }
92}

More examples

Hide additional examples

examples/zimage_stepcheck.rs (line 48)

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        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
71        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
72        np.device = false;
73        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
74        let ok = dit.attach_device_pair(&prep, &np, 3, None);
75        println!("pair prepared: {ok}  (neg L = {})", nids.len());
76        if let Some((pp, pn)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
77            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
78            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
79                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
80        }
81        return;
82    }
83    let tok_a = dit.tokens(&xa, &shape);
84    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
85    let mut hp = zimage::ZPrepared { device: false, ..prep };
86    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
87    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
88    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
89        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
90    if let Some(b) = a.get(9) {
91        hp.device = true;
92        let xb = lat(b);
93        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
94        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
95            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
96    }
97}