Changes
diff --git a/diffusion_cli/zimage_model.py b/diffusion_cli/zimage_model.py
index 45dc6f2..1c9f5a4 100644
--- a/diffusion_cli/zimage_model.py
+++ b/diffusion_cli/zimage_model.py
@@ -1,12 +1,837 @@
-"""Diffusion model loading boundary for Z-Image Turbo."""
+"""Standalone Z-Image/Lumina2 diffusion model implementation."""
from __future__ import annotations
+from dataclasses import dataclass
+import math
+from pathlib import Path
+import re
+from typing import Iterable
+
+import torch
+from safetensors.torch import load_file
+from torch import nn
+from torch.nn import functional as F
+
+from diffusion_cli.errors import DiffusionCliError
+
+
+ZIMAGE_PREFIX = "model.diffusion_model."
+ZIMAGE_DIM = 3840
+ZIMAGE_CAP_FEATURE_DIM = 2560
+ZIMAGE_LAYERS = 30
+ZIMAGE_REFINER_LAYERS = 2
+ZIMAGE_HEADS = 30
+ZIMAGE_KV_HEADS = 30
+ZIMAGE_PATCH_SIZE = 2
+ZIMAGE_CHANNELS = 16
+ZIMAGE_AXES_DIMS = (32, 48, 48)
+ZIMAGE_ROPE_THETA = 256.0
+ZIMAGE_FFN_MULTIPLIER = 8.0 / 3.0
+ZIMAGE_PAD_TOKENS_MULTIPLE = 32
+ZIMAGE_TIME_SCALE = 1000.0
+TIMESTEP_EMBEDDING_SIZE = 256
+UNUSED_ZIMAGE_KEYS = frozenset({"norm_final.weight"})
+
+
+@dataclass(frozen=True)
+class ZImageConfig:
+ """Architecture settings for the Z-Image Turbo diffusion model."""
+
+ dim: int = ZIMAGE_DIM
+ cap_feat_dim: int = ZIMAGE_CAP_FEATURE_DIM
+ n_layers: int = ZIMAGE_LAYERS
+ n_refiner_layers: int = ZIMAGE_REFINER_LAYERS
+ n_heads: int = ZIMAGE_HEADS
+ n_kv_heads: int = ZIMAGE_KV_HEADS
+ patch_size: int = ZIMAGE_PATCH_SIZE
+ in_channels: int = ZIMAGE_CHANNELS
+ axes_dims: tuple[int, int, int] = ZIMAGE_AXES_DIMS
+ rope_theta: float = ZIMAGE_ROPE_THETA
+ ffn_dim_multiplier: float = ZIMAGE_FFN_MULTIPLIER
+ pad_tokens_multiple: int = ZIMAGE_PAD_TOKENS_MULTIPLE
+ time_scale: float = ZIMAGE_TIME_SCALE
+
+
+def normalizeZImageStateDict(
+ state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+ """Strip known checkpoint prefixes from Z-Image diffusion keys."""
+
+ normalized = {}
+ for key, value in state_dict.items():
+ if key.startswith(ZIMAGE_PREFIX):
+ normalized[key.removeprefix(ZIMAGE_PREFIX)] = value
+ elif key.startswith("diffusion_model."):
+ normalized[key.removeprefix("diffusion_model.")] = value
+ else:
+ normalized[key] = value
+ return normalized
+
+
+def detectZImageConfig(
+ state_dict: dict[str, torch.Tensor],
+) -> ZImageConfig:
+ """Infer the supported Z-Image architecture from state-dict shapes."""
+
+ normalized = normalizeZImageStateDict(state_dict)
+ try:
+ cap_weight = normalized["cap_embedder.1.weight"]
+ x_weight = normalized["x_embedder.weight"]
+ except KeyError as exc:
+ raise DiffusionCliError(
+ "Unsupported Z-Image checkpoint: missing Lumina2 keys"
+ ) from exc
+
+ layer_ids = _numberedChildren(normalized, "layers")
+ refiner_ids = _numberedChildren(normalized, "context_refiner")
+ dim = cap_weight.shape[0]
+ cap_feat_dim = cap_weight.shape[1]
+ patch_features = x_weight.shape[1]
+ patch_size = int(math.sqrt(patch_features // ZIMAGE_CHANNELS))
+
+ if dim != ZIMAGE_DIM:
+ raise DiffusionCliError(
+ f"Unsupported Z-Image dim: expected {ZIMAGE_DIM}, got {dim}"
+ )
+ if patch_size != ZIMAGE_PATCH_SIZE:
+ raise DiffusionCliError(
+ "Unsupported Z-Image patch size: "
+ f"expected {ZIMAGE_PATCH_SIZE}, got {patch_size}"
+ )
+
+ return ZImageConfig(
+ dim=dim,
+ cap_feat_dim=cap_feat_dim,
+ n_layers=len(layer_ids),
+ n_refiner_layers=len(refiner_ids),
+ patch_size=patch_size,
+ pad_tokens_multiple=(
+ ZIMAGE_PAD_TOKENS_MULTIPLE
+ if "cap_pad_token" in normalized
+ else 0
+ ),
+ )
+
+
+def _numberedChildren(
+ state_dict: dict[str, torch.Tensor],
+ prefix: str,
+) -> tuple[int, ...]:
+ pattern = re.compile(rf"^{re.escape(prefix)}\.(\d+)\.")
+ return tuple(sorted({int(match.group(1)) for key in state_dict
+ if (match := pattern.match(key))}))
+
+
+def timestepEmbedding(
+ timesteps: torch.Tensor,
+ dim: int,
+ *,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """Build sinusoidal timestep embeddings matching ComfyUI MMDiT."""
+
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period)
+ * torch.arange(0, half, dtype=torch.float32, device=timesteps.device)
+ / half
+ )
+ args = timesteps[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat(
+ [embedding, torch.zeros_like(embedding[:, :1])],
+ dim=-1,
+ )
+ return embedding.to(timesteps.dtype)
+
+
+def rope(pos: torch.Tensor, dim: int, theta: float) -> torch.Tensor:
+ """Create rotary-position matrices for one positional axis."""
+
+ if dim % 2 != 0:
+ raise ValueError("RoPE dimension must be even")
+ scale = torch.linspace(
+ 0,
+ (dim - 2) / dim,
+ steps=dim // 2,
+ dtype=torch.float64,
+ device=pos.device,
+ )
+ omega = 1.0 / (theta ** scale)
+ out = torch.einsum("...n,d->...nd", pos.float(), omega)
+ out = torch.stack(
+ [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)],
+ dim=-1,
+ )
+ return out.reshape(*out.shape[:-1], 2, 2).float()
+
+
+def applyRope1(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+ """Apply precomputed rotary-position matrices to a q or k tensor."""
+
+ x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
+ if (
+ x_.shape[2] != 1
+ and freqs_cis.shape[2] != 1
+ and x_.shape[2] != freqs_cis.shape[2]
+ ):
+ freqs_cis = freqs_cis[:, :, :x_.shape[2]]
+ x_out = freqs_cis[..., 0] * x_[..., 0]
+ x_out = x_out + freqs_cis[..., 1] * x_[..., 1]
+ return x_out.reshape(*x.shape).type_as(x)
+
+
+def applyRope(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ freqs_cis: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Apply RoPE to query and key tensors."""
+
+ return applyRope1(query, freqs_cis), applyRope1(key, freqs_cis)
+
+
+def posIdsX(
+ start_t: int,
+ h_tokens: int,
+ w_tokens: int,
+ batch_size: int,
+ device,
+) -> torch.Tensor:
+ """Build image-token position ids used by Lumina2 RoPE."""
+
+ pos_ids = torch.zeros(
+ (batch_size, h_tokens * w_tokens, 3),
+ dtype=torch.float32,
+ device=device,
+ )
+ pos_ids[:, :, 0] = start_t
+ pos_ids[:, :, 1] = (
+ torch.arange(h_tokens, dtype=torch.float32, device=device)
+ .view(-1, 1)
+ .repeat(1, w_tokens)
+ .flatten()
+ )
+ pos_ids[:, :, 2] = (
+ torch.arange(w_tokens, dtype=torch.float32, device=device)
+ .view(1, -1)
+ .repeat(h_tokens, 1)
+ .flatten()
+ )
+ return pos_ids
+
+
+def padToPatchSize(
+ latent: torch.Tensor,
+ patch_size: int,
+) -> torch.Tensor:
+ """Pad a latent tensor so height and width are patch-size multiples."""
+
+ pad_h = (-latent.shape[-2]) % patch_size
+ pad_w = (-latent.shape[-1]) % patch_size
+ if pad_h == 0 and pad_w == 0:
+ return latent
+ return F.pad(latent, (0, pad_w, 0, pad_h), mode="circular")
+
+
+def padZImage(
+ feats: torch.Tensor,
+ pad_token: torch.Tensor,
+ multiple: int,
+) -> tuple[torch.Tensor, int]:
+ """Pad a token sequence to the multiple used by Z-Image checkpoints."""
+
+ if multiple <= 0:
+ return feats, 0
+ pad_extra = (-feats.shape[1]) % multiple
+ if pad_extra == 0:
+ return feats, 0
+ pad = pad_token.to(device=feats.device, dtype=feats.dtype)
+ pad = pad.unsqueeze(0).repeat(feats.shape[0], pad_extra, 1)
+ return torch.cat((feats, pad), dim=1), pad_extra
+
+
+def modulate(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ """Apply Z-Image scale-only AdaLN modulation."""
+
+ return x * (1 + scale.unsqueeze(1))
+
+
+def clampFp16(x: torch.Tensor) -> torch.Tensor:
+ """Clamp fp16 overflow while leaving other dtypes untouched."""
+
+ if x.dtype == torch.float16:
+ return torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
+ return x
+
+
+class EmbedND(nn.Module):
+ """N-axis rotary embedding builder used by Lumina2."""
+
+ def __init__(self, dim: int, theta: float, axes_dim: Iterable[int]):
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ self.axes_dim = tuple(axes_dim)
+
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
+ """Return RoPE matrices for position ids."""
+
+ axes = ids.shape[-1]
+ emb = torch.cat(
+ [rope(ids[..., axis], self.axes_dim[axis], self.theta)
+ for axis in range(axes)],
+ dim=-3,
+ )
+ return emb.unsqueeze(1)
+
+
+class TimestepEmbedder(nn.Module):
+ """Embed scalar diffusion timesteps for AdaLN modulation."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ *,
+ frequency_embedding_size: int = TIMESTEP_EMBEDDING_SIZE,
+ output_size: int | None = None,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ if output_size is None:
+ output_size = hidden_size
+ self.frequency_embedding_size = frequency_embedding_size
+ self.mlp = nn.Sequential(
+ nn.Linear(
+ frequency_embedding_size,
+ hidden_size,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ ),
+ nn.SiLU(),
+ nn.Linear(
+ hidden_size,
+ output_size,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ ),
+ )
+
+ def forward(self, timesteps: torch.Tensor, dtype) -> torch.Tensor:
+ """Embed timesteps into the requested runtime dtype."""
+
+ t_freq = timestepEmbedding(
+ timesteps,
+ self.frequency_embedding_size,
+ ).to(dtype)
+ return self.mlp(t_freq)
+
+
+class JointAttention(nn.Module):
+ """Grouped-query attention used by each Lumina2 block."""
+
+ def __init__(
+ self,
+ dim: int,
+ n_heads: int,
+ n_kv_heads: int,
+ *,
+ out_bias: bool,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ self.n_heads = n_heads
+ self.n_kv_heads = n_kv_heads
+ self.head_dim = dim // n_heads
+ qkv_dim = (n_heads + n_kv_heads + n_kv_heads) * self.head_dim
+ self.qkv = nn.Linear(dim, qkv_dim, bias=False, dtype=dtype,
+ device=device)
+ self.out = nn.Linear(n_heads * self.head_dim, dim, bias=out_bias,
+ dtype=dtype, device=device)
+ self.q_norm = nn.RMSNorm(self.head_dim, dtype=dtype, device=device)
+ self.k_norm = nn.RMSNorm(self.head_dim, dtype=dtype, device=device)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ freqs_cis: torch.Tensor,
+ ) -> torch.Tensor:
+ """Run attention over the concatenated text and image sequence."""
+
+ batch_size, seq_len, _ = x.shape
+ query, key, value = torch.split(
+ self.qkv(x),
+ [
+ self.n_heads * self.head_dim,
+ self.n_kv_heads * self.head_dim,
+ self.n_kv_heads * self.head_dim,
+ ],
+ dim=-1,
+ )
+ query = query.view(batch_size, seq_len, self.n_heads, self.head_dim)
+ key = key.view(batch_size, seq_len, self.n_kv_heads, self.head_dim)
+ value = value.view(batch_size, seq_len, self.n_kv_heads,
+ self.head_dim)
+
+ query = self.q_norm(query)
+ key = self.k_norm(key)
+ query, key = applyRope(query, key, freqs_cis)
+
+ repeats = self.n_heads // self.n_kv_heads
+ if repeats > 1:
+ key = key.unsqueeze(3).repeat(1, 1, 1, repeats, 1).flatten(2, 3)
+ value = (
+ value.unsqueeze(3)
+ .repeat(1, 1, 1, repeats, 1)
+ .flatten(2, 3)
+ )
+
+ query = query.movedim(1, 2)
+ key = key.movedim(1, 2)
+ value = value.movedim(1, 2)
+ output = F.scaled_dot_product_attention(query, key, value)
+ output = output.movedim(1, 2).flatten(2)
+ return self.out(output)
+
+
+class FeedForward(nn.Module):
+ """SwiGLU feed-forward network used by Lumina2 blocks."""
+
+ def __init__(
+ self,
+ dim: int,
+ *,
+ multiple_of: int = 256,
+ ffn_dim_multiplier: float = ZIMAGE_FFN_MULTIPLIER,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ hidden_dim = int(ffn_dim_multiplier * dim)
+ hidden_dim = multiple_of * (
+ (hidden_dim + multiple_of - 1) // multiple_of
+ )
+ self.w1 = nn.Linear(dim, hidden_dim, bias=False, dtype=dtype,
+ device=device)
+ self.w2 = nn.Linear(hidden_dim, dim, bias=False, dtype=dtype,
+ device=device)
+ self.w3 = nn.Linear(dim, hidden_dim, bias=False, dtype=dtype,
+ device=device)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply the feed-forward transform."""
+
+ return self.w2(clampFp16(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class JointTransformerBlock(nn.Module):
+ """One Lumina2 transformer block with optional timestep modulation."""
+
+ def __init__(
+ self,
+ config: ZImageConfig,
+ *,
+ modulation: bool,
+ attn_out_bias: bool = False,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ self.modulation = modulation
+ self.attention = JointAttention(
+ config.dim,
+ config.n_heads,
+ config.n_kv_heads,
+ out_bias=attn_out_bias,
+ dtype=dtype,
+ device=device,
+ )
+ self.feed_forward = FeedForward(
+ config.dim,
+ ffn_dim_multiplier=config.ffn_dim_multiplier,
+ dtype=dtype,
+ device=device,
+ )
+ self.attention_norm1 = nn.RMSNorm(config.dim, eps=1e-5, dtype=dtype,
+ device=device)
+ self.attention_norm2 = nn.RMSNorm(config.dim, eps=1e-5, dtype=dtype,
+ device=device)
+ self.ffn_norm1 = nn.RMSNorm(config.dim, eps=1e-5, dtype=dtype,
+ device=device)
+ self.ffn_norm2 = nn.RMSNorm(config.dim, eps=1e-5, dtype=dtype,
+ device=device)
+ if modulation:
+ self.adaLN_modulation = nn.Sequential(
+ nn.Linear(
+ min(config.dim, 256),
+ 4 * config.dim,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ ),
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ freqs_cis: torch.Tensor,
+ adaln_input: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """Apply attention and feed-forward residual updates."""
+
+ if self.modulation:
+ if adaln_input is None:
+ raise ValueError("adaln_input is required for modulated block")
+ scale_msa, gate_msa, scale_mlp, gate_mlp = (
+ self.adaLN_modulation(adaln_input).chunk(4, dim=1)
+ )
+ attn = self.attention(
+ modulate(self.attention_norm1(x), scale_msa),
+ freqs_cis,
+ )
+ x = x + gate_msa.unsqueeze(1).tanh() * self.attention_norm2(
+ clampFp16(attn),
+ )
+ mlp = self.feed_forward(modulate(self.ffn_norm1(x), scale_mlp))
+ x = x + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(
+ clampFp16(mlp),
+ )
+ return x
+
+ attn = self.attention(self.attention_norm1(x), freqs_cis)
+ x = x + self.attention_norm2(clampFp16(attn))
+ x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
+ return x
+
+
+class FinalLayer(nn.Module):
+ """Final patch projection for Z-Image denoising output."""
+
+ def __init__(self, config: ZImageConfig, *, dtype=None, device=None):
+ super().__init__()
+ self.norm_final = nn.LayerNorm(
+ config.dim,
+ elementwise_affine=False,
+ eps=1e-6,
+ dtype=dtype,
+ device=device,
+ )
+ self.linear = nn.Linear(
+ config.dim,
+ config.patch_size * config.patch_size * config.in_channels,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ )
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(
+ 256,
+ config.dim,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ ),
+ )
+
+ def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+ """Project hidden patch tokens back to latent patch values."""
+
+ scale = self.adaLN_modulation(c)
+ return self.linear(modulate(self.norm_final(x), scale))
+
+
+class NextDiT(nn.Module):
+ """Z-Image Turbo NextDiT network for text-to-image denoising."""
+
+ def __init__(
+ self,
+ config: ZImageConfig,
+ *,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ self.config = config
+ self.x_embedder = nn.Linear(
+ config.patch_size * config.patch_size * config.in_channels,
+ config.dim,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ )
+ self.noise_refiner = nn.ModuleList([
+ JointTransformerBlock(config, modulation=True, dtype=dtype,
+ device=device)
+ for _ in range(config.n_refiner_layers)
+ ])
+ self.context_refiner = nn.ModuleList([
+ JointTransformerBlock(config, modulation=False, dtype=dtype,
+ device=device)
+ for _ in range(config.n_refiner_layers)
+ ])
+ self.t_embedder = TimestepEmbedder(
+ min(config.dim, 1024),
+ output_size=256,
+ dtype=dtype,
+ device=device,
+ )
+ self.cap_embedder = nn.Sequential(
+ nn.RMSNorm(config.cap_feat_dim, eps=1e-5, dtype=dtype,
+ device=device),
+ nn.Linear(config.cap_feat_dim, config.dim, bias=True,
+ dtype=dtype, device=device),
+ )
+ self.layers = nn.ModuleList([
+ JointTransformerBlock(
+ config,
+ modulation=True,
+ attn_out_bias=False,
+ dtype=dtype,
+ device=device,
+ )
+ for _ in range(config.n_layers)
+ ])
+ self.final_layer = FinalLayer(config, dtype=dtype, device=device)
+ self.x_pad_token = nn.Parameter(
+ torch.empty((1, config.dim), dtype=dtype, device=device),
+ )
+ self.cap_pad_token = nn.Parameter(
+ torch.empty((1, config.dim), dtype=dtype, device=device),
+ )
+ self.rope_embedder = EmbedND(
+ dim=config.dim // config.n_heads,
+ theta=config.rope_theta,
+ axes_dim=config.axes_dims,
+ )
+
+ def forward(
+ self,
+ latent: torch.Tensor,
+ timesteps: torch.Tensor,
+ context: torch.Tensor,
+ num_tokens: int,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """Run one Z-Image denoising forward pass."""
+
+ del attention_mask
+ if context.shape[1] < num_tokens:
+ raise ValueError("num_tokens exceeds context sequence length")
+ original_h, original_w = latent.shape[-2], latent.shape[-1]
+ latent = padToPatchSize(latent, self.config.patch_size)
+ timesteps = timesteps.to(device=latent.device, dtype=latent.dtype)
+ t = 1.0 - timesteps
+ adaln_input = self.t_embedder(
+ t * self.config.time_scale,
+ dtype=latent.dtype,
+ )
+
+ img, img_size, cap_size, freqs_cis = self.patchifyAndEmbed(
+ latent,
+ context,
+ adaln_input,
+ )
+ for layer in self.layers:
+ img = layer(img, freqs_cis, adaln_input)
+
+ img = self.final_layer(img, adaln_input)
+ img = self.unpatchify(img, img_size, cap_size)
+ return -img[:, :, :original_h, :original_w]
+
+ def patchifyAndEmbed(
+ self,
+ latent: torch.Tensor,
+ cap_feats: torch.Tensor,
+ adaln_input: torch.Tensor,
+ ) -> tuple[torch.Tensor, list[tuple[int, int]], list[int], torch.Tensor]:
+ """Embed text and latent patches into one transformer sequence."""
+
+ batch_size, channels, height, width = latent.shape
+ patch_size = self.config.patch_size
+ img_sizes = [(height, width)] * batch_size
+
+ cap_feats = self.cap_embedder(cap_feats)
+ cap_feats, _ = padZImage(
+ cap_feats,
+ self.cap_pad_token,
+ self.config.pad_tokens_multiple,
+ )
+ cap_pos_ids = torch.zeros(
+ batch_size,
+ cap_feats.shape[1],
+ 3,
+ dtype=torch.float32,
+ device=latent.device,
+ )
+ cap_pos_ids[:, :, 0] = (
+ torch.arange(cap_feats.shape[1], dtype=torch.float32,
+ device=latent.device)
+ + 1.0
+ )
+ cap_freqs = self.rope_embedder(cap_pos_ids).movedim(1, 2)
+
+ for layer in self.context_refiner:
+ cap_feats = layer(cap_feats, cap_freqs)
+
+ patches = (
+ latent.view(
+ batch_size,
+ channels,
+ height // patch_size,
+ patch_size,
+ width // patch_size,
+ patch_size,
+ )
+ .permute(0, 2, 4, 3, 5, 1)
+ .flatten(3)
+ .flatten(1, 2)
+ )
+ patches = self.x_embedder(patches)
+ x_pos_ids = posIdsX(
+ cap_feats.shape[1] + 1,
+ height // patch_size,
+ width // patch_size,
+ batch_size,
+ latent.device,
+ )
+ patches, _ = padZImage(
+ patches,
+ self.x_pad_token,
+ self.config.pad_tokens_multiple,
+ )
+ if patches.shape[1] != x_pos_ids.shape[1]:
+ pad = patches.shape[1] - x_pos_ids.shape[1]
+ x_pos_ids = F.pad(x_pos_ids, (0, 0, 0, pad))
+ x_freqs = self.rope_embedder(x_pos_ids).movedim(1, 2)
+
+ for layer in self.noise_refiner:
+ patches = layer(patches, x_freqs, adaln_input)
+
+ img = torch.cat((cap_feats, patches), dim=1)
+ freqs_cis = torch.cat((cap_freqs, x_freqs), dim=1).to(img.device)
+ cap_size = [cap_feats.shape[1]] * batch_size
+ return img, img_sizes, cap_size, freqs_cis
+
+ def unpatchify(
+ self,
+ x: torch.Tensor,
+ img_size: list[tuple[int, int]],
+ cap_size: list[int],
+ ) -> torch.Tensor:
+ """Convert output patch tokens back to an NCHW latent tensor."""
+
+ patch_size = self.config.patch_size
+ images = []
+ for i in range(x.size(0)):
+ height, width = img_size[i]
+ begin = cap_size[i]
+ end = begin + (height // patch_size) * (width // patch_size)
+ image = (
+ x[i][begin:end]
+ .view(
+ height // patch_size,
+ width // patch_size,
+ patch_size,
+ patch_size,
+ self.config.in_channels,
+ )
+ .permute(4, 0, 2, 1, 3)
+ .flatten(3, 4)
+ .flatten(1, 2)
+ )
+ images.append(image)
+ return torch.stack(images, dim=0)
+
class ZImageModel:
- """Placeholder boundary for the future Z-Image/Lumina2 model port."""
+ """Load and execute the local Z-Image Turbo diffusion model."""
+
+ def __init__(
+ self,
+ model_path: Path,
+ device,
+ dtype,
+ *,
+ model: nn.Module | None = None,
+ ) -> None:
+ self.device = device
+ self.dtype = dtype
+ if model is None:
+ self.model = self._loadModel(model_path)
+ else:
+ self.model = model.to(device=device, dtype=dtype)
+ self.model.eval()
+
+ def forward(
+ self,
+ latent: torch.Tensor,
+ sigma: torch.Tensor | float,
+ conditioning,
+ ) -> torch.Tensor:
+ """Run one denoising forward pass from text conditioning."""
+
+ if isinstance(sigma, torch.Tensor):
+ timesteps = sigma.to(device=self.device, dtype=self.dtype)
+ else:
+ timesteps = torch.tensor([sigma], device=self.device,
+ dtype=self.dtype)
+ if timesteps.ndim == 0:
+ timesteps = timesteps[None]
+ if timesteps.shape[0] == 1 and latent.shape[0] != 1:
+ timesteps = timesteps.expand(latent.shape[0])
+
+ with torch.inference_mode():
+ output = self.model(
+ latent.to(device=self.device, dtype=self.dtype),
+ timesteps,
+ conditioning.hidden_states.to(
+ device=self.device,
+ dtype=self.dtype,
+ ),
+ conditioning.token_count,
+ conditioning.attention_mask.to(self.device),
+ )
+
+ if output.shape != latent.shape:
+ raise DiffusionCliError(
+ f"Z-Image output shape {tuple(output.shape)} does not match "
+ f"latent shape {tuple(latent.shape)}"
+ )
+ if not torch.all(torch.isfinite(output)):
+ raise DiffusionCliError("Z-Image forward produced NaN or Inf")
+ return output
- def __init__(self, *_args, **_kwargs) -> None:
- raise NotImplementedError(
- "Z-Image diffusion model execution is not implemented yet"
+ def _loadModel(self, model_path: Path) -> nn.Module:
+ state_dict = normalizeZImageStateDict(load_file(model_path,
+ device="cpu"))
+ config = detectZImageConfig(state_dict)
+ with torch.device("meta"):
+ model = NextDiT(config, dtype=self.dtype)
+ missing_keys, unexpected_keys = model.load_state_dict(
+ state_dict,
+ strict=False,
+ assign=True,
)
+ missing = [key for key in missing_keys if key]
+ unexpected = [
+ key for key in unexpected_keys
+ if key and key not in UNUSED_ZIMAGE_KEYS
+ ]
+ if missing or unexpected:
+ details = []
+ if missing:
+ details.append(f"missing keys: {missing[:8]}")
+ if unexpected:
+ details.append(f"unexpected keys: {unexpected[:8]}")
+ raise DiffusionCliError(
+ "Could not load Z-Image diffusion model: "
+ + "; ".join(details)
+ )
+ return model.to(device=self.device, dtype=self.dtype)
diff --git a/tests/test_zimage_model.py b/tests/test_zimage_model.py
new file mode 100644
index 0000000..99bbb86
--- /dev/null
+++ b/tests/test_zimage_model.py
@@ -0,0 +1,112 @@
+import unittest
+
+import torch
+
+from diffusion_cli.text_encoder import TextConditioning
+from diffusion_cli.zimage_model import (
+ NextDiT,
+ ZImageConfig,
+ ZImageModel,
+ detectZImageConfig,
+ normalizeZImageStateDict,
+)
+
+
+class ZImageModelTest(unittest.TestCase):
+ def testStateDictNormalizationStripsModelPrefix(self):
+ state_dict = normalizeZImageStateDict(
+ {
+ "model.diffusion_model.x_embedder.weight": "a",
+ "diffusion_model.final_layer.linear.weight": "b",
+ "layers.0.ffn_norm1.weight": "c",
+ }
+ )
+
+ self.assertEqual(state_dict["x_embedder.weight"], "a")
+ self.assertEqual(state_dict["final_layer.linear.weight"], "b")
+ self.assertEqual(state_dict["layers.0.ffn_norm1.weight"], "c")
+
+ def testDetectConfigReadsZImageShapes(self):
+ state_dict = {
+ "model.diffusion_model.cap_embedder.1.weight": torch.empty(
+ 3840,
+ 2560,
+ ),
+ "model.diffusion_model.x_embedder.weight": torch.empty(3840, 64),
+ "model.diffusion_model.cap_pad_token": torch.empty(1, 3840),
+ "model.diffusion_model.layers.0.ffn_norm1.weight": torch.empty(
+ 3840,
+ ),
+ "model.diffusion_model.layers.1.ffn_norm1.weight": torch.empty(
+ 3840,
+ ),
+ "model.diffusion_model.context_refiner.0.ffn_norm1.weight":
+ torch.empty(3840),
+ }
+
+ config = detectZImageConfig(state_dict)
+
+ self.assertEqual(config.dim, 3840)
+ self.assertEqual(config.cap_feat_dim, 2560)
+ self.assertEqual(config.n_layers, 2)
+ self.assertEqual(config.n_refiner_layers, 1)
+ self.assertEqual(config.patch_size, 2)
+ self.assertEqual(config.pad_tokens_multiple, 32)
+
+ def testTinyForwardMatchesLatentShapeAndIsFinite(self):
+ torch.manual_seed(1)
+ config = ZImageConfig(
+ dim=258,
+ cap_feat_dim=8,
+ n_layers=1,
+ n_refiner_layers=1,
+ n_heads=3,
+ n_kv_heads=3,
+ patch_size=2,
+ in_channels=4,
+ axes_dims=(30, 28, 28),
+ pad_tokens_multiple=0,
+ )
+ model = NextDiT(config, dtype=torch.float32)
+ latent = torch.randn(1, 4, 4, 4)
+ context = torch.randn(1, 3, 8)
+ timesteps = torch.tensor([0.5])
+
+ output = model(latent, timesteps, context, 3)
+
+ self.assertEqual(output.shape, latent.shape)
+ self.assertTrue(torch.all(torch.isfinite(output)))
+
+ def testWrapperValidatesFiniteOutput(self):
+ class IdentityModel(torch.nn.Module):
+ def forward(
+ self,
+ latent,
+ timesteps,
+ context,
+ num_tokens,
+ attention_mask,
+ ):
+ return latent + timesteps.reshape(-1, 1, 1, 1) * 0
+
+ wrapper = ZImageModel(
+ model_path=None,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ model=IdentityModel(),
+ )
+ conditioning = TextConditioning(
+ hidden_states=torch.zeros(1, 2, 8),
+ attention_mask=torch.ones(1, 2),
+ token_count=2,
+ )
+ latent = torch.zeros(1, 4, 4, 4)
+
+ output = wrapper.forward(latent, torch.tensor([0.5]), conditioning)
+
+ self.assertEqual(output.shape, latent.shape)
+ self.assertTrue(torch.all(torch.isfinite(output)))
+
+
+if __name__ == "__main__":
+ unittest.main()