Struct ZImageDit 

pub struct ZImageDit {
    pub cfg: ZConfig,
    pub model: Option<Arc<CmfModel>>,
    /* private fields */
}

The Z-Image transformer. Tensor names are the ORIGINAL diffusers names under dit. (no rename map).

Fields

cfg: ZConfig

model: Option<Arc<CmfModel>>

The container this was loaded from (None for load_dir) — device paths need tensor indices into it.

Implementations

impl ZImageDit

pub fn load_dir(dir: &Path) -> Result<Self, String>

Load from a diffusers transformer/ directory (bf16/fp32 safetensors, read whole and widened to f32 — a dev/parity path).

pub fn from_cmf(model: &Arc<CmfModel>) -> Result<Self, String>

Load from a packaged .cmf (dit.* tensors + dit.config_json). Quantized projections stay mmap-resident.

Examples found in repository

examples/zimage_devcheck.rs (line 48)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 61)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 43)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn geom(&self) -> ZGeom

Examples found in repository

examples/zimage_devcheck.rs (line 68)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 89)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

pub fn temb(&self, t_model: f32) -> Vec<f32>

Timestep embedding for t_model = (1000 − 1000σ)/1000: sinusoid of t·1000 (cos first, f32 args) → mlp.0 → SiLU → mlp.2. Returns [256].

pub fn mods_for_steps(&self, t_models: &[f32]) -> Vec<f32>

Raw adaLN_modulation.0(temb) for every step and block: [steps][2 + n_layers][4 · dim], chunks [scale_msa, gate_msa, scale_mlp, gate_mlp] (no +1, no tanh). f64 accumulation.

Examples found in repository

examples/zimage_devcheck.rs (line 57)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 72)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 46)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn final_scale_for_steps(&self, t_models: &[f32]) -> Vec<f32>

1 + all_final_layer.2-1.adaLN_modulation.1(SiLU(temb)) per step: [steps][dim].

Examples found in repository

examples/zimage_devcheck.rs (line 58)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 73)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 47)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn embed_caption(&self, cap_feats: &[f32], l: usize) -> Vec<f32>

Caption features [l, cap_feat_dim] → [l_p, dim]: pad rows are copies of the last row, RMSNorm(w, 1e-5) → Linear + b, rows ≥ l := cap_pad_token.

Examples found in repository

examples/zimage_devcheck.rs (line 63)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 78)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

pub fn refine_caption_cpu(&self, cap: &mut [f32], rope_cap: (&[f32], &[f32]))

The two unmodulated context-refiner blocks on the host, in place on cap [l_p, dim].

Examples found in repository

examples/zimage_devcheck.rs (line 64)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 79)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

pub fn block_cpu( &self, blk: ZBlockId, x: &mut [f32], n: usize, rope: (&[f32], &[f32]), m: Option<&[f32]>, )

One block on the host, in place on x [n, dim]. m = the block’s raw modulation [4 · dim] (None = unmodulated: s = 0, gate = 1).

pub fn block_refs(&self) -> Option<ZBlockRefs<'_>>

Device views of all blocks (requires from_cmf).

pub fn prepare( &self, cap_feats: &[f32], shape: ZShape, key: u64, mods_all: Option<(&[f32], &[f32])>, ) -> Result<ZPrepared, String>

Once per (prompt, resolution): caption embed → context refiner (gpu::zimage_refine_caption, else CPU) → rope tables → gpu::zimage_prepare (sets ZPrepared::device). mods_all = (mods_for_steps, final_scale_for_steps) forwarded to the backend.

Examples found in repository

examples/zimage_devcheck.rs (line 61)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 76)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 49)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn prepare_with( &self, cap_feats: &[f32], shape: ZShape, key: u64, mods_all: Option<(&[f32], &[f32])>, device: bool, ) -> Result<ZPrepared, String>

prepare with the device use explicit: device = false builds a pure host state (CPU context refiner, no gpu::zimage_prepare) — the reference a device test diffs step against.

Examples found in repository

examples/zimage_metal_check.rs (line 109)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

pub fn prepare_host( &self, cap_feats: &[f32], shape: ZShape, key: u64, device_refine: bool, ) -> Result<ZPrepared, String>

The host half of prepare: caption embed → context refiner (gpu::zimage_refine_caption when device_refine, else CPU) → rope tables. device is false until attach_device.

pub fn preload_device(&self) -> bool

Upload every device plane now (gpu::zimage_preload); independent of the caption, so it can run beside the text encoder.

pub fn attach_device( &self, p: &mut ZPrepared, mods_all: Option<(&[f32], &[f32])>, ) -> bool

gpu::zimage_prepare for a host state (sets p.device).

pub fn attach_device_pair( &self, pos: &ZPrepared, neg: &ZPrepared, key: u64, mods_all: Option<(&[f32], &[f32])>, ) -> bool

One batch-2 device program for a CFG pair (item 0 = pos, item 1 = neg, both at the same resolution) under key. false = the backend has no batch 2 here; the caller steps the items one by one.

Examples found in repository

examples/zimage_metal_check.rs (line 127)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 91)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn step_pair_device( &self, key: u64, n_img: usize, step: usize, x_tok: &[f32], mods: &[f32], final_scale: &[f32], ) -> Option<(Vec<f32>, Vec<f32>)>

Both items of a CFG pair prepared by attach_device_pair under key, in one device forward: returns (v_pos, v_neg), or None when the backend declined (the caller steps the items separately).

Examples found in repository

examples/zimage_metal_check.rs (line 131)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 100)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn step( &self, p: &ZPrepared, step: usize, x_tok: &[f32], mods: &[f32], final_scale: &[f32], ) -> Vec<f32>

One DiT forward: gpu::zimage_step if p.device, else step_cpu. x_tok [n_img_p, 64] (pad rows = copies of the last row), mods [(2 + n_layers) · 4 · dim] of this step, final_scale [dim]. Returns v [n_img, 64] (before the pipeline’s negation).

Examples found in repository

examples/zimage_devcheck.rs (line 70)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}

examples/zimage_metal_check.rs (line 93)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

examples/zimage_stepcheck.rs (line 80)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn tokens(&self, latent: &[f32], shape: &ZShape) -> Vec<f32>

The step input for a latent [16, h_lat, w_lat]: patchify, then pad to n_img_p rows by repeating the last row.

Examples found in repository

examples/zimage_stepcheck.rs (line 69)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
    let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
    let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
    let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
    let cap = {
        let _p = cortiq_engine::gpu::pause_gpu();
        cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
    };
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let sig = zimage::sigmas_torch_f32(steps, shift);
    let t = zimage::t_model(sig[i]);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let shape = ZShape::new(hh, ww, ids.len());
    let prep = dit.prepare(&cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    let lat = |dir: &str| -> Vec<f32> {
        if i == 0 {
            read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
        } else {
            read_f32(&format!("{dir}/lat_{i}.f32"))
        }
    };
    let xa = lat(&a[8]);
    // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
    // forward against the two items stepped one by one.
    if let Ok(neg) = std::env::var("ZC_NEG") {
        let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
        let ncap = {
            let _p = cortiq_engine::gpu::pause_gpu();
            cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
        };
        let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
        let tok_a = dit.tokens(&xa, &shape);
        // `ZC_TAPS=<dir>`: every block's residual stream of the three
        // forwards (single pos, single neg, pair) into <dir>/{pos,neg,pair},
        // then compared block by block (the pair's item rows vs the single).
        let taps = std::env::var("ZC_TAPS").ok();
        let set_taps = |sub: &str| {
            if let Some(d) = &taps {
                unsafe { std::env::set_var("CMF_ZI_TAPS", format!("{d}/{sub}")) };
            }
        };
        set_taps("pos");
        let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
        set_taps("neg");
        let vn = dit.step(&np, i, &tok_a, &mods, &fs);
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        np.device = false;
        let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
        np.device = true;
        // `ZC_SWAP=1`: the pair as (neg, pos) — tells a positional cause
        // (item 0 vs item 1) from a content one (the caption length).
        let swap = std::env::var("ZC_SWAP").as_deref() == Ok("1");
        let (first, second) = if swap { (&np, &prep) } else { (&prep, &np) };
        let ok = dit.attach_device_pair(first, second, 3, None);
        println!(
            "pair prepared: {ok}  (pos L = {}, neg L = {}, n_img_p {}, order {})",
            ids.len(),
            nids.len(),
            shape.n_img_p,
            if swap { "neg,pos" } else { "pos,neg" }
        );
        set_taps("pair");
        if let Some((p0, p1)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
            let (pp, pn) = if swap { (p1, p0) } else { (p0, p1) };
            let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
            println!("pair vs singles: pos {:.3e}  neg {:.3e}  non-finite {nan}   single neg dev vs cpu {:.3e}",
                rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
        }
        unsafe { std::env::remove_var("CMF_ZI_TAPS") };
        if let Some(d) = &taps {
            // Row layout of the pair: image stage [item][n_img_p]; joint
            // stage [item0: n_img_p + cp0][item1: n_img_p + cp1].
            let h = dit.cfg.dim;
            let cp = |l: usize| l.div_ceil(32) * 32;
            let (cp_pos, cp_neg) = (cp(ids.len()), cp(nids.len()));
            let (cp0, cp1) = if swap { (cp_neg, cp_pos) } else { (cp_pos, cp_neg) };
            let nip = shape.n_img_p;
            let rd = |p: String| -> Option<Vec<f32>> {
                std::fs::read(&p).ok().map(|b| b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect())
            };
            let mut names: Vec<String> = (0..2).map(|k| format!("nr{k}_out")).collect();
            names.extend((0..dit.cfg.n_layers).map(|k| format!("l{k}_out")));
            for n in names {
                let (Some(sp), Some(sn), Some(pr)) = (
                    rd(format!("{d}/pos/step{i}/{n}.f32")),
                    rd(format!("{d}/neg/step{i}/{n}.f32")),
                    rd(format!("{d}/pair/step{i}/{n}.f32")),
                ) else {
                    continue;
                };
                let (r0, r1) = if n.starts_with("nr") {
                    ((0, nip), (nip, nip))
                } else {
                    ((0, nip + cp0), (nip + cp0, nip + cp1))
                };
                let item = |r: (usize, usize)| &pr[r.0 * h..(r.0 + r.1) * h];
                let (ip, ineg) = if swap { (item(r1), item(r0)) } else { (item(r0), item(r1)) };
                // image rows and caption rows separately
                let split = |v: &[f32]| (v[..nip * h].to_vec(), v[nip * h..].to_vec());
                let (pi, pc) = split(ip);
                let (spi, spc) = split(&sp[..ip.len().min(sp.len())]);
                println!(
                    "{n:9} pos: img {:.3e} cap {:.3e}   neg: all {:.3e}",
                    rel(&pi, &spi),
                    if pc.is_empty() { 0.0 } else { rel(&pc, &spc) },
                    rel(ineg, &sn[..ineg.len().min(sn.len())])
                );
            }
        }
        return;
    }
    let tok_a = dit.tokens(&xa, &shape);
    let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
    let mut hp = zimage::ZPrepared { device: false, ..prep };
    let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
    let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
    println!("step {i} on A's lat: dev vs cpu {:.3e}   cpu vs A's v {:.3e}   dev vs A's v {:.3e}",
        rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
    if let Some(b) = a.get(9) {
        hp.device = true;
        let xb = lat(b);
        let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
        println!("inputs A vs B {:.3e}   device outputs {:.3e}   (B's own v {:.3e})",
            rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
    }
}

pub fn embed_image( &self, x_tok: &[f32], n_img: usize, n_img_p: usize, ) -> Vec<f32>

x_tok [n_img_p, 64] → [n_img_p, dim]: Linear + b, rows ≥ n_img := x_pad_token.

pub fn final_layer(&self, u: &[f32], n: usize, final_scale: &[f32]) -> Vec<f32>

Final layer on the first n rows of u: LayerNorm(eps 1e-6, no affine) · final_scale → Linear(dim → 64) + b.

pub fn step_cpu( &self, p: &ZPrepared, x_tok: &[f32], mods: &[f32], final_scale: &[f32], ) -> Vec<f32>

The host reference forward (WP2/WP3 gate against this): embed → pad rows := x_pad_token → noise refiner (img only) → concat [img, cap] → layers → LayerNorm(1e-6)·final_scale → Linear → image rows.

Examples found in repository

examples/zimage_metal_check.rs (line 112)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let reps: usize = a.get(4).and_then(|v| v.parse().ok()).unwrap_or(2);
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let bf = od.join(format!("dit_{case}_bf16.safetensors"));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    let t0 = std::time::Instant::now();
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {} ({:.3}s incl. caption refine)", prep.device, t0.elapsed().as_secs_f64());
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    if let Some(cr) = o.get("cr1_out") {
        println!(
            "caption  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
            rel(&prep.cap, &cap_cpu),
            rel(&cap_cpu, cr),
            rel(&prep.cap, cr)
        );
    }
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let mut v_dev = Vec::new();
    for r in 0..reps {
        let ts = std::time::Instant::now();
        let v = dit.step(&prep, 0, &x_tok, &mods, &fs);
        println!("device step {r}: {:.3}s", ts.elapsed().as_secs_f64());
        if r > 0 {
            println!("replay   {:.3e}", rel(&v, &v_dev));
        }
        v_dev = v;
    }
    let orc = &o["final_out"];
    println!("v        dev vs oracle fp32 {:.3e}", rel(&v_dev, orc));
    if bf.exists() {
        let (ob, _) = read_st(&bf);
        if let Some(vb) = ob.get("final_out") {
            println!("v        oracle bf16 vs fp32 {:.3e} (the diffusers-bf16 floor)", rel(vb, orc));
        }
    }
    if std::env::var("ZC_NOCPU").as_deref() != Ok("1") {
        let mut prep_cpu = dit.prepare_with(cap, shape, 2, None, false).unwrap();
        prep_cpu.cap = prep.cap.clone();
        let ts = std::time::Instant::now();
        let v_cpu = dit.step_cpu(&prep_cpu, &x_tok, &mods, &fs);
        println!("cpu step: {:.3}s", ts.elapsed().as_secs_f64());
        println!(
            "v        dev vs cpu {:.3e}   cpu vs oracle fp32 {:.3e}",
            rel(&v_dev, &v_cpu),
            rel(&v_cpu, orc)
        );
    }
    if std::env::var("ZC_NEG").as_deref() == Ok("1") {
        // batch-2 pair: item 1 = the same caption at a different padded
        // length is not available here, so use the same caption; the pair
        // must reproduce the single forward exactly per item
        let mut p2 = dit.prepare(cap, shape, 7, None).unwrap();
        p2.device = false;
        let pair_key = 9;
        let okp = dit.attach_device_pair(&prep, &p2, pair_key, None);
        println!("pair prepared: {okp}");
        if okp {
            let ts = std::time::Instant::now();
            let r = dit.step_pair_device(pair_key, shape.n_img, 0, &x_tok, &mods, &fs).unwrap();
            println!("pair step: {:.3}s", ts.elapsed().as_secs_f64());
            println!("pair     item0 vs single {:.3e}   item1 vs single {:.3e}", rel(&r.0, &v_dev), rel(&r.1, &v_dev));
        }
    }
}

pub fn step_cpu_taps( &self, p: &ZPrepared, x_tok: &[f32], mods: &[f32], final_scale: &[f32], tap: &mut dyn FnMut(&str, &[f32]), ) -> Vec<f32>

step_cpu with a tap callback: (name, tensor) at the oracle’s tap points (x_seq, nr{i}_out, u_in, l{i}_out, final_out).

Examples found in repository

examples/zimage_devcheck.rs (line 82)

fn main() {
    unsafe { std::env::set_var("CMF_GPU", "1") };
    let a: Vec<String> = std::env::args().collect();
    let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
    let dit = ZImageDit::from_cmf(&model).unwrap();
    let od = std::path::Path::new(&a[2]);
    let case = &a[3];
    let (o, meta) = read_st(&od.join(format!("dit_{case}_fp32.safetensors")));
    let (hh, ww, l) = (meta_usize(&meta, "H"), meta_usize(&meta, "W"), meta_usize(&meta, "L"));
    let t = o["t_model"][0];
    let cap = &o["cap"];
    let x_in = &o["x_in"];
    let shape = ZShape::new(hh, ww, l);
    let mods = dit.mods_for_steps(&[t]);
    let fs = dit.final_scale_for_steps(&[t]);
    let rope = zimage::ids_and_rope(shape.grid, l, dit.cfg.rope_theta, dit.cfg.axes_dims);
    // caption: device-refined (prepare) vs CPU-refined
    let prep = dit.prepare(cap, shape, 1, None).unwrap();
    println!("device prepared: {}", prep.device);
    let mut cap_cpu = dit.embed_caption(cap, l);
    dit.refine_caption_cpu(&mut cap_cpu, (&rope.cap.0, &rope.cap.1));
    println!("cr1_out  dev vs cpu {:.3e}   cpu vs oracle {:.3e}   dev vs oracle {:.3e}",
        rel(&prep.cap, &cap_cpu), rel(&cap_cpu, &o["cr1_out"]), rel(&prep.cap, &o["cr1_out"]));
    let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
    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);
    let tapdir = std::env::var("CMF_ZI_TAPS").unwrap_or("/root/zb/taps".into());
    let v_dev = dit.step(&prep, 0, &x_tok, &mods, &fs);
    // Replays must be stateless: the same inputs again, and a different
    // step's mods in between.
    let mods2 = dit.mods_for_steps(&[t * 0.5 + 0.25]);
    let fs2 = dit.final_scale_for_steps(&[t * 0.5 + 0.25]);
    let _ = dit.step(&prep, 1, &x_tok, &mods2, &fs2);
    let v_dev2 = dit.step(&prep, 2, &x_tok, &mods, &fs);
    println!("replay   dev2 vs dev1 {:.3e}", rel(&v_dev2, &v_dev));
    let mut taps: Vec<(String, Vec<f32>)> = Vec::new();
    // the CPU path must see the SAME (device-refined) caption to isolate the DiT
    let mut prep_cpu = prep;
    prep_cpu.device = false;
    let v_cpu = dit.step_cpu_taps(&prep_cpu, &x_tok, &mods, &fs, &mut |n, v| taps.push((n.to_string(), v.to_vec())));
    println!("v        dev vs cpu {:.3e}   cpu vs oracle(final_out) {:.3e}", rel(&v_dev, &v_cpu), rel(&v_cpu, &o["final_out"]));
    for (n, v) in &taps {
        let p = std::path::Path::new(&tapdir).join("step0").join(format!("{n}.f32"));
        if let Ok(b) = std::fs::read(&p) {
            let d: Vec<f32> = b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
            let or = o.get(n.as_str()).map(|w| format!("{:.3e}", rel(v, w))).unwrap_or_default();
            println!("{n:10} dev vs cpu {:.3e}   (cpu vs oracle {or})", rel(&d, v));
        }
    }
}