"""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 torch import nn
from torch.nn import functional as F
from diffusion_cli.checkpoint import loadStateDict
from diffusion_cli.config import DIFFUSION_ROLE, ModelSource
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:
"""Load and execute the local Z-Image Turbo diffusion model."""
def __init__(
self,
model_path: ModelSource | Path | None,
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():
model_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),
)
sigma = timesteps.to(
device=model_output.device,
dtype=model_output.dtype,
)
sigma = sigma.reshape(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
model_input = latent.to(device=self.device, dtype=self.dtype)
output = model_input - model_output * sigma
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 toDevice(self, device, dtype) -> None:
"""Move the diffusion model to the active generation device."""
self.device = device
self.dtype = dtype
self.model.to(device=device, dtype=dtype)
self.model.eval()
def toCpu(self) -> None:
"""Move the diffusion model back to CPU memory."""
self.device = torch.device("cpu")
self.model.to(device=self.device)
self.model.eval()
def _loadModel(self, model_path: ModelSource | Path | None) -> nn.Module:
if model_path is None:
raise ValueError("model_path is required when model is absent")
if isinstance(model_path, ModelSource):
model_source = model_path
else:
model_source = ModelSource(model_path, DIFFUSION_ROLE)
state_dict = normalizeZImageStateDict(loadStateDict(model_source))
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)