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.

pub fn geom(&self) -> ZGeom

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.

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].

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.

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].

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.

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.

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.

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).

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).

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.

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.

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).