From fac19aa6fa9bfdf6d914b7b372de406070aeef27 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Mon, 13 Jul 2026 17:45:41 -0700 Subject: [PATCH] [feat]: add native Z-Image transformer prototype --- fastvideo/configs/models/dits/__init__.py | 3 +- fastvideo/configs/models/dits/zimage.py | 60 ++ fastvideo/models/dits/zimage.py | 570 ++++++++++++++++++ fastvideo/models/registry.py | 1 + tests/local_tests/zimage/PORT_STATUS.md | 29 +- tests/local_tests/zimage/README.md | 39 +- .../zimage/test_zimage_transformer_parity.py | 339 +++++++++++ 7 files changed, 1018 insertions(+), 23 deletions(-) create mode 100644 fastvideo/configs/models/dits/zimage.py create mode 100644 fastvideo/models/dits/zimage.py create mode 100644 tests/local_tests/zimage/test_zimage_transformer_parity.py diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 64a5c118d..91cd4b3cf 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -12,6 +12,7 @@ from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig from fastvideo.configs.models.dits.stable_audio import StableAudioConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig +from fastvideo.configs.models.dits.zimage import ZImageDiTConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig @@ -19,5 +20,5 @@ __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", - "StableAudioConfig", "GlmImageDiTConfig" + "StableAudioConfig", "GlmImageDiTConfig", "ZImageDiTConfig" ] diff --git a/fastvideo/configs/models/dits/zimage.py b/fastvideo/configs/models/dits/zimage.py new file mode 100644 index 000000000..ce9bf7848 --- /dev/null +++ b/fastvideo/configs/models/dits/zimage.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def is_zimage_block(name: str, module) -> bool: + parts = name.split(".") + return len(parts) >= 2 and parts[-2] in {"noise_refiner", "context_refiner", "layers"} and parts[-1].isdigit() + + +@dataclass +class ZImageDiTArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_block]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.TORCH_SDPA, ) + + all_patch_size: tuple[int, ...] = (2, ) + all_f_patch_size: tuple[int, ...] = (1, ) + in_channels: int = 16 + dim: int = 3840 + n_layers: int = 30 + n_refiner_layers: int = 2 + n_heads: int = 30 + n_kv_heads: int = 30 + norm_eps: float = 1e-5 + qk_norm: bool = True + cap_feat_dim: int = 2560 + rope_theta: float = 256.0 + t_scale: float = 1000.0 + axes_dims: tuple[int, ...] = (32, 48, 48) + axes_lens: tuple[int, ...] = (1536, 512, 512) + + adaln_embed_dim: int = 256 + frequency_embedding_size: int = 256 + timestep_mid_size: int = 1024 + max_period: int = 10000 + seq_multi_of: int = 32 + + def __post_init__(self) -> None: + super().__post_init__() + if len(self.all_patch_size) != len(self.all_f_patch_size): + raise ValueError("all_patch_size and all_f_patch_size must have equal length") + if self.dim % self.n_heads: + raise ValueError("dim must be divisible by n_heads") + if self.dim // self.n_heads != sum(self.axes_dims): + raise ValueError("attention head dimension must equal sum(axes_dims)") + if len(self.axes_dims) != len(self.axes_lens) or any(dim % 2 for dim in self.axes_dims): + raise ValueError("RoPE axes require matching lengths and even dimensions") + + self.hidden_size = self.dim + self.num_attention_heads = self.n_heads + self.num_channels_latents = self.in_channels + self.out_channels = self.in_channels + + +@dataclass +class ZImageDiTConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=ZImageDiTArchConfig) + prefix: str = "ZImage" diff --git a/fastvideo/models/dits/zimage.py b/fastvideo/models/dits/zimage.py new file mode 100644 index 000000000..969bf8be9 --- /dev/null +++ b/fastvideo/models/dits/zimage.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo-native Z-Image transformer. + +Z-Image attends over padded variable-length image/text streams and requires a +key-padding mask. FastVideo's distributed attention wrappers do not expose that +mask contract yet, so this implementation uses torch SDPA and is SP=1 only. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence + +from fastvideo.configs.models.dits.zimage import ZImageDiTConfig +from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + + +def _linear(layer: ReplicatedLinear, x: torch.Tensor) -> torch.Tensor: + return layer(x)[0] + + +def _prepare_attention_mask(attention_mask: torch.Tensor | None, dtype: torch.dtype) -> torch.Tensor | None: + if attention_mask is None: + return None + if attention_mask.ndim == 2: + attention_mask = attention_mask[:, None, None, :] + if attention_mask.dtype == torch.bool: + additive_mask = torch.zeros_like(attention_mask, dtype=dtype) + additive_mask.masked_fill_(~attention_mask, float("-inf")) + return additive_mask + return attention_mask + + +class TimestepEmbedder(nn.Module): + + def __init__( + self, + out_size: int, + mid_size: int, + frequency_embedding_size: int, + max_period: int, + ) -> None: + super().__init__() + self.mlp = nn.ModuleList([ + ReplicatedLinear(frequency_embedding_size, mid_size, bias=True), + nn.SiLU(), + ReplicatedLinear(mid_size, out_size, bias=True), + ]) + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half) + args = t[:, 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 + + def forward(self, t: torch.Tensor) -> torch.Tensor: + t_freq = self.timestep_embedding(t, self.frequency_embedding_size, self.max_period) + weight_dtype = self.mlp[0].weight.dtype + if weight_dtype.is_floating_point: + t_freq = t_freq.to(weight_dtype) + return _linear(self.mlp[2], self.mlp[1](_linear(self.mlp[0], t_freq))) + + +class RMSNorm(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-5) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return output * self.weight + + +class FeedForward(nn.Module): + + def __init__(self, dim: int, hidden_dim: int) -> None: + super().__init__() + self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False) + self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False) + self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _linear(self.w2, F.silu(_linear(self.w1, x)) * _linear(self.w3, x)) + + +def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + x_out = torch.view_as_real(x * freqs_cis.unsqueeze(2)).flatten(3) + return x_out.type_as(x_in) + + +class ZImageAttention(nn.Module): + + def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool = True, eps: float = 1e-5) -> None: + super().__init__() + self.n_heads = n_heads + self.n_kv_heads = n_kv_heads + self.head_dim = dim // n_heads + + self.to_q = ReplicatedLinear(dim, n_heads * self.head_dim, bias=False) + self.to_k = ReplicatedLinear(dim, n_kv_heads * self.head_dim, bias=False) + self.to_v = ReplicatedLinear(dim, n_kv_heads * self.head_dim, bias=False) + self.to_out = nn.ModuleList([ReplicatedLinear(n_heads * self.head_dim, dim, bias=False)]) + + self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else None + self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, + ) -> torch.Tensor: + query = _linear(self.to_q, hidden_states).unflatten(-1, (self.n_heads, -1)) + key = _linear(self.to_k, hidden_states).unflatten(-1, (self.n_kv_heads, -1)) + value = _linear(self.to_v, hidden_states).unflatten(-1, (self.n_kv_heads, -1)) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + if freqs_cis is not None: + query = apply_rotary_emb(query, freqs_cis) + key = apply_rotary_emb(key, freqs_cis) + + mask = _prepare_attention_mask(attention_mask, query.dtype) + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).contiguous() + return _linear(self.to_out[0], hidden_states.flatten(2, 3).to(query.dtype)) + + +class ZImageTransformerBlock(nn.Module): + + def __init__( + self, + layer_id: int, + dim: int, + n_heads: int, + n_kv_heads: int, + norm_eps: float, + qk_norm: bool, + adaln_embed_dim: int, + modulation: bool = True, + ) -> None: + super().__init__() + self.dim = dim + self.head_dim = dim // n_heads + self.layer_id = layer_id + self.modulation = modulation + + self.attention = ZImageAttention(dim, n_heads, n_kv_heads, qk_norm, norm_eps) + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + self.attention_norm1 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.attention_norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + if modulation: + self.adaLN_modulation = nn.ModuleList( + [ReplicatedLinear(min(dim, adaln_embed_dim), 4 * dim, bias=True)]) + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + freqs_cis: torch.Tensor, + adaln_input: torch.Tensor | None = None, + ) -> torch.Tensor: + if self.modulation: + assert adaln_input is not None + scale_msa, gate_msa, scale_mlp, gate_mlp = _linear( + self.adaLN_modulation[0], adaln_input).unsqueeze(1).chunk(4, dim=2) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, + attention_mask=attn_mask, + freqs_cis=freqs_cis, + ) + x = x + gate_msa * self.attention_norm2(attn_out) + x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + else: + attn_out = self.attention( + self.attention_norm1(x), + attention_mask=attn_mask, + freqs_cis=freqs_cis, + ) + x = x + self.attention_norm2(attn_out) + x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x))) + return x + + +class FinalLayer(nn.Module): + + def __init__(self, hidden_size: int, out_channels: int, adaln_embed_dim: int) -> None: + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = ReplicatedLinear(hidden_size, out_channels, bias=True) + self.adaLN_modulation = nn.ModuleList([ + nn.SiLU(), + ReplicatedLinear(min(hidden_size, adaln_embed_dim), hidden_size, bias=True), + ]) + + def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + scale = 1.0 + _linear(self.adaLN_modulation[1], self.adaLN_modulation[0](c)) + return _linear(self.linear, self.norm_final(x) * scale.unsqueeze(1)) + + +class RopeEmbedder: + + def __init__(self, theta: float, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...]) -> None: + if len(axes_dims) != len(axes_lens): + raise ValueError("RoPE axes require matching dimensions and lengths") + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + self.freqs_cis: list[torch.Tensor] | None = None + + @staticmethod + def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float) -> list[torch.Tensor]: + with torch.device("cpu"): + freqs_cis = [] + for axis_dim, axis_end in zip(dim, end): + freqs = 1.0 / (theta**(torch.arange(0, axis_dim, 2, dtype=torch.float64) / axis_dim)) + timestep = torch.arange(axis_end, dtype=torch.float64) + angles = torch.outer(timestep, freqs).float() + freqs_cis.append(torch.polar(torch.ones_like(angles), angles).to(torch.complex64)) + return freqs_cis + + def __call__(self, ids: torch.Tensor) -> torch.Tensor: + if ids.ndim != 2 or ids.shape[-1] != len(self.axes_dims): + raise ValueError("RoPE ids must have shape [sequence, number_of_axes]") + if self.freqs_cis is None: + self.freqs_cis = [ + freqs.to(ids.device) + for freqs in self.precompute_freqs_cis(self.axes_dims, self.axes_lens, self.theta) + ] + elif self.freqs_cis[0].device != ids.device: + self.freqs_cis = [freqs.to(ids.device) for freqs in self.freqs_cis] + return torch.cat([self.freqs_cis[i][ids[:, i]] for i in range(len(self.axes_dims))], dim=-1) + + +class ZImageTransformer2DModel(BaseDiT): + _default_config = ZImageDiTConfig() + _fsdp_shard_conditions = _default_config.arch_config._fsdp_shard_conditions + _compile_conditions = _default_config.arch_config._compile_conditions + _supported_attention_backends = (AttentionBackendEnum.TORCH_SDPA, ) + param_names_mapping = _default_config.arch_config.param_names_mapping + reverse_param_names_mapping = _default_config.arch_config.reverse_param_names_mapping + + def __init__(self, config: ZImageDiTConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + arch = config.arch_config + + self.in_channels = arch.in_channels + self.out_channels = arch.in_channels + self.all_patch_size = tuple(arch.all_patch_size) + self.all_f_patch_size = tuple(arch.all_f_patch_size) + self.dim = arch.dim + self.n_heads = arch.n_heads + self.rope_theta = arch.rope_theta + self.t_scale = arch.t_scale + self.seq_multi_of = arch.seq_multi_of + + self.hidden_size = arch.dim + self.num_attention_heads = arch.n_heads + self.num_channels_latents = arch.in_channels + + self.all_x_embedder = nn.ModuleDict({ + f"{patch_size}-{f_patch_size}": ReplicatedLinear( + f_patch_size * patch_size * patch_size * arch.in_channels, arch.dim, bias=True) + for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size) + }) + self.all_final_layer = nn.ModuleDict({ + f"{patch_size}-{f_patch_size}": FinalLayer( + arch.dim, + patch_size * patch_size * f_patch_size * self.out_channels, + arch.adaln_embed_dim, + ) + for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size) + }) + + block_kwargs = { + "dim": arch.dim, + "n_heads": arch.n_heads, + "n_kv_heads": arch.n_kv_heads, + "norm_eps": arch.norm_eps, + "qk_norm": arch.qk_norm, + "adaln_embed_dim": arch.adaln_embed_dim, + } + self.noise_refiner = nn.ModuleList([ + ZImageTransformerBlock(1000 + layer_id, modulation=True, **block_kwargs) + for layer_id in range(arch.n_refiner_layers) + ]) + self.context_refiner = nn.ModuleList([ + ZImageTransformerBlock(layer_id, modulation=False, **block_kwargs) + for layer_id in range(arch.n_refiner_layers) + ]) + self.t_embedder = TimestepEmbedder( + min(arch.dim, arch.adaln_embed_dim), + mid_size=arch.timestep_mid_size, + frequency_embedding_size=arch.frequency_embedding_size, + max_period=arch.max_period, + ) + self.cap_embedder = nn.ModuleList([ + RMSNorm(arch.cap_feat_dim, eps=arch.norm_eps), + ReplicatedLinear(arch.cap_feat_dim, arch.dim, bias=True), + ]) + self.x_pad_token = nn.Parameter(torch.empty((1, arch.dim))) + self.cap_pad_token = nn.Parameter(torch.empty((1, arch.dim))) + self.layers = nn.ModuleList([ + ZImageTransformerBlock(layer_id, modulation=True, **block_kwargs) for layer_id in range(arch.n_layers) + ]) + self.axes_dims = tuple(arch.axes_dims) + self.axes_lens = tuple(arch.axes_lens) + self.rope_embedder = RopeEmbedder(arch.rope_theta, self.axes_dims, self.axes_lens) + self.__post_init__() + + def unpatchify( + self, + x: list[torch.Tensor], + size: list[tuple[int, int, int]], + patch_size: int, + f_patch_size: int, + ) -> list[torch.Tensor]: + patch_height = patch_width = patch_size + patch_frames = f_patch_size + if len(x) != len(size): + raise ValueError("output batch and original sizes must have equal length") + for i, (frames, height, width) in enumerate(size): + original_length = (frames // patch_frames) * (height // patch_height) * (width // patch_width) + x[i] = (x[i][:original_length].view( + frames // patch_frames, + height // patch_height, + width // patch_width, + patch_frames, + patch_height, + patch_width, + self.out_channels, + ).permute(6, 0, 3, 1, 4, 2, 5).reshape(self.out_channels, frames, height, width)) + return x + + @staticmethod + def create_coordinate_grid( + size: tuple[int, int, int], + start: tuple[int, int, int] | None = None, + device: torch.device | None = None, + ) -> torch.Tensor: + start = start or (0, ) * len(size) + axes = [ + torch.arange(axis_start, axis_start + span, dtype=torch.int32, device=device) + for axis_start, span in zip(start, size) + ] + return torch.stack(torch.meshgrid(axes, indexing="ij"), dim=-1) + + def patchify_and_embed( + self, + all_image: list[torch.Tensor], + all_cap_feats: list[torch.Tensor], + patch_size: int, + f_patch_size: int, + ) -> tuple[ + list[torch.Tensor], + list[torch.Tensor], + list[tuple[int, int, int]], + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], + ]: + patch_height = patch_width = patch_size + patch_frames = f_patch_size + device = all_image[0].device + + image_out = [] + image_sizes = [] + image_pos_ids = [] + image_pad_masks = [] + cap_pos_ids = [] + cap_pad_masks = [] + cap_feats_out = [] + + for image, cap_feat in zip(all_image, all_cap_feats): + cap_length = len(cap_feat) + cap_padding = (-cap_length) % self.seq_multi_of + cap_pos_ids.append( + self.create_coordinate_grid( + (cap_length + cap_padding, 1, 1), + start=(1, 0, 0), + device=device, + ).flatten(0, 2)) + cap_pad_masks.append( + torch.cat([ + torch.zeros(cap_length, dtype=torch.bool, device=device), + torch.ones(cap_padding, dtype=torch.bool, device=device), + ]) if cap_padding else torch.zeros(cap_length, dtype=torch.bool, device=device)) + cap_feats_out.append( + torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding, 1)]) if cap_padding else cap_feat) + + channels, frames, height, width = image.size() + image_sizes.append((frames, height, width)) + frame_tokens, height_tokens, width_tokens = ( + frames // patch_frames, + height // patch_height, + width // patch_width, + ) + image = image.view( + channels, + frame_tokens, + patch_frames, + height_tokens, + patch_height, + width_tokens, + patch_width, + ) + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape( + frame_tokens * height_tokens * width_tokens, + patch_frames * patch_height * patch_width * channels, + ) + + image_length = len(image) + image_padding = (-image_length) % self.seq_multi_of + original_pos_ids = self.create_coordinate_grid( + (frame_tokens, height_tokens, width_tokens), + start=(cap_length + cap_padding + 1, 0, 0), + device=device, + ).flatten(0, 2) + if image_padding: + padding_pos_ids = self.create_coordinate_grid((1, 1, 1), device=device).flatten(0, 2).repeat( + image_padding, 1) + image_pos_ids.append(torch.cat([original_pos_ids, padding_pos_ids])) + else: + image_pos_ids.append(original_pos_ids) + image_pad_masks.append( + torch.cat([ + torch.zeros(image_length, dtype=torch.bool, device=device), + torch.ones(image_padding, dtype=torch.bool, device=device), + ]) if image_padding else torch.zeros(image_length, dtype=torch.bool, device=device)) + image_out.append( + torch.cat([image, image[-1:].repeat(image_padding, 1)]) if image_padding else image) + + return ( + image_out, + cap_feats_out, + image_sizes, + image_pos_ids, + cap_pos_ids, + image_pad_masks, + cap_pad_masks, + ) + + @staticmethod + def _attention_mask(lengths: list[int], device: torch.device) -> torch.Tensor: + mask = torch.zeros((len(lengths), max(lengths)), dtype=torch.bool, device=device) + for i, length in enumerate(lengths): + mask[i, :length] = True + return mask + + def forward( + self, + hidden_states: torch.Tensor | list[torch.Tensor], + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + patch_size: int = 2, + f_patch_size: int = 1, + **kwargs, + ) -> tuple[list[torch.Tensor], dict]: + del encoder_hidden_states_image, guidance, kwargs + if model_parallel_is_initialized() and get_sp_world_size() != 1: + raise NotImplementedError( + "Z-Image masked SDPA does not support sequence parallelism; run with sp_size=1") + if patch_size not in self.all_patch_size or f_patch_size not in self.all_f_patch_size: + raise ValueError(f"unsupported patch sizes: spatial={patch_size}, temporal={f_patch_size}") + if isinstance(hidden_states, torch.Tensor): + hidden_states = list(hidden_states.unbind(0)) + if isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = list(encoder_hidden_states.unbind(0)) + + device = hidden_states[0].device + timestep_embedding = self.t_embedder(timestep * self.t_scale) + ( + hidden_states, + encoder_hidden_states, + image_sizes, + image_pos_ids, + cap_pos_ids, + image_inner_pad_masks, + cap_inner_pad_masks, + ) = self.patchify_and_embed(hidden_states, encoder_hidden_states, patch_size, f_patch_size) + + image_lengths = [len(item) for item in hidden_states] + if not all(length % self.seq_multi_of == 0 for length in image_lengths): + raise ValueError("padded image sequence lengths must be aligned") + hidden_states = torch.cat(hidden_states) + hidden_states = _linear(self.all_x_embedder[f"{patch_size}-{f_patch_size}"], hidden_states) + adaln_input = timestep_embedding.type_as(hidden_states) + hidden_states[torch.cat(image_inner_pad_masks)] = self.x_pad_token + hidden_states = list(hidden_states.split(image_lengths)) + image_freqs_cis = list( + self.rope_embedder(torch.cat(image_pos_ids)).split([len(item) for item in image_pos_ids])) + hidden_states = pad_sequence(hidden_states, batch_first=True, padding_value=0.0) + image_freqs_cis = pad_sequence(image_freqs_cis, batch_first=True, padding_value=0.0) + image_freqs_cis = image_freqs_cis[:, :hidden_states.shape[1]] + image_attn_mask = self._attention_mask(image_lengths, device) + for layer in self.noise_refiner: + hidden_states = layer(hidden_states, image_attn_mask, image_freqs_cis, adaln_input) + + cap_lengths = [len(item) for item in encoder_hidden_states] + if not all(length % self.seq_multi_of == 0 for length in cap_lengths): + raise ValueError("padded caption sequence lengths must be aligned") + encoder_hidden_states = torch.cat(encoder_hidden_states) + encoder_hidden_states = _linear(self.cap_embedder[1], self.cap_embedder[0](encoder_hidden_states)) + encoder_hidden_states[torch.cat(cap_inner_pad_masks)] = self.cap_pad_token + encoder_hidden_states = list(encoder_hidden_states.split(cap_lengths)) + cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids)).split([len(item) for item in cap_pos_ids])) + encoder_hidden_states = pad_sequence(encoder_hidden_states, batch_first=True, padding_value=0.0) + cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) + cap_freqs_cis = cap_freqs_cis[:, :encoder_hidden_states.shape[1]] + cap_attn_mask = self._attention_mask(cap_lengths, device) + for layer in self.context_refiner: + encoder_hidden_states = layer(encoder_hidden_states, cap_attn_mask, cap_freqs_cis) + + unified = [] + unified_freqs_cis = [] + for i, (image_length, cap_length) in enumerate(zip(image_lengths, cap_lengths)): + unified.append( + torch.cat([hidden_states[i][:image_length], encoder_hidden_states[i][:cap_length]])) + unified_freqs_cis.append( + torch.cat([image_freqs_cis[i][:image_length], cap_freqs_cis[i][:cap_length]])) + unified_lengths = [image_length + cap_length for image_length, cap_length in zip(image_lengths, cap_lengths)] + unified = pad_sequence(unified, batch_first=True, padding_value=0.0) + unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) + unified_attn_mask = self._attention_mask(unified_lengths, device) + for layer in self.layers: + unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input) + + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input) + outputs = self.unpatchify(list(unified.unbind(0)), image_sizes, patch_size, f_patch_size) + return outputs, {} + + +EntryClass = ZImageTransformer2DModel diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index d88d91fef..74aba3435 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -64,6 +64,7 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { # Text-to-image DiT models (2D image generation) _TEXT_TO_IMAGE_DIT_MODELS = { "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"), + "ZImageTransformer2DModel": ("dits", "zimage", "ZImageTransformer2DModel"), } _TEXT_ENCODER_MODELS = { diff --git a/tests/local_tests/zimage/PORT_STATUS.md b/tests/local_tests/zimage/PORT_STATUS.md index 1dc19b2d4..156bbb4b8 100644 --- a/tests/local_tests/zimage/PORT_STATUS.md +++ b/tests/local_tests/zimage/PORT_STATUS.md @@ -5,7 +5,7 @@ - model_family: `zimage` - variant: `Z-Image-Turbo` - workload_types: `T2I` -- PR scope: component-only (`text_encoder`, `tokenizer`, `vae`, `scheduler`) +- PR scope: components plus native transformer prototype/parity - official_ref: `https://github.com/Tongyi-MAI/Z-Image` - official_ref_dir: `/Z-Image/src` - official_ref_commit: `26f23eda626ffadda020b04ff79488e1d72004cd` @@ -20,10 +20,10 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension ## Current Phase -- phase: Phase 6 (component parity revalidation) +- phase: Phase 6 (transformer parity activation and component revalidation) - status: in_progress -- owner: parity -- last_updated: 2026-07-04 +- owner: dit/parity +- last_updated: 2026-07-13 ## Component Matrix @@ -33,7 +33,7 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension | Tokenizer | tokenizer | reused | `AutoTokenizer` from `/tokenizer/`; official thinking template and length 512 | `TokenizerLoader` | DONE | not_needed | REVALIDATION REQUIRED (historical unpinned PASS) | I003, I010 | | VAE | vae | reused with accepted wrapper exception | Pinned `zimage.AutoencoderKL`; official decode uses `(latents / scaling_factor) + shift_factor` | Existing Diffusers-backed `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` through production `VAELoader`, plus direct implementation coverage | DONE | not_needed | REVALIDATION REQUIRED (historical unpinned direct-decode PASS only) | I008 | | Scheduler | scheduler | reused with Z-Image extension | Pinned `zimage.FlowMatchEulerDiscreteScheduler`; official pipeline sets `sigma_min=0.0` | `FlowMatchEulerDiscreteScheduler(use_reference_discrete_timesteps=True, sigma_min=0.0)` | DONE | not_needed | REVALIDATION REQUIRED; asset-free regressions remain separately runnable | I002, I009 | -| Transformer (`ZImageTransformer2DModel`) | dit | planned native port | `/transformer/` and pinned official source | future PR | NOT_STARTED | unknown_until_prototype | NOT_STARTED | I004 | +| Transformer (`ZImageTransformer2DModel`) | dit | native port | Pinned `src/zimage/transformer.py`; published Diffusers transformer subfolder | `fastvideo.models.dits.zimage.ZImageTransformer2DModel` through `TransformerLoader` | DONE (production meta + tiny CPU forward) | not_needed; exact 521-key surface | TINY NON-SKIP PASS; production real-weight forward PENDING | I011 | | Pipeline | pipeline | planned | Pinned `ZImagePipeline` | future PR | NOT_STARTED | depends_on_transformer | NOT_STARTED | I005 | ## Conversion State @@ -41,8 +41,9 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension - conversion_script: none for the reused components in this component-only PR - converted_weights_dir: none - source_layout: diffusers -- future_transformer_conversion: unknown until native key/shape prototype exists -- strict_load_status: revalidation_required +- transformer_conversion: not_needed; official, HF index, and native names align exactly +- transformer_key_surface: 521 tensors; sorted-key SHA256 `3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354` +- strict_load_status: transformer real-weight strict load pending; other components revalidation_required - production_passthrough: body-only Qwen3 `AutoModel`, tokenizer assets - direct_load_components: native Qwen3 parity target, VAE, scheduler - native_qwen_exclusion: `lm_head.weight` only; the native encoder owns the body, not an LM head @@ -60,6 +61,7 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension | Text encoder fp32 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | RERUN REQUIRED; historical unpinned native PASS on L40S, 2026-06-21 | New run must cover independent reference, body-only production path, and native implementation. | | Text encoder bf16 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16] -v -s` | RERUN REQUIRED; historical unpinned native PASS on L40S, 2026-06-21 | Historical thresholds retained pending pinned-snapshot validation. | | Per-layer bf16 diagnostic | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | RERUN REQUIRED; historical unpinned informational run on L40S, 2026-06-21 | Pinned architecture has 36 blocks and 37 hidden-state entries under the tested contract. | +| Transformer meta/tiny parity | `ZIMAGE_OFFICIAL_REF_DIR= pytest tests/local_tests/zimage/test_zimage_transformer_parity.py -v -s` | 4 PASSED, 1 SKIPPED on CPU, 2026-07-13 | Production 521-key/shape surface and pinned HF key digest match; additive-mask attention and tiny variable-length full-forward parity are exact; SP>1 guard passes. Real-weight test skipped because the 24.6 GB transformer subfolder/CUDA were absent. | ## Open Questions @@ -68,7 +70,7 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension | Q001 | Which exact Z-Image source revision is the parity oracle? | prep | Phase 1 | resolved | Pinned `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`; scheduler/VAE tests fail on a wrong HEAD or import origin. | | Q002 | Which immutable HF snapshot supplies component weights? | prep | Phase 1 | resolved | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6`. | | Q003 | What prompt/tokenization/hidden-state contract does the official pipeline use? | pipeline | Phase 6 | resolved | `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)`, tokenization length 512, and valid-token `hidden_states[-2]`. | -| Q004 | Does the future native transformer require conversion? | conversion | Phase 5 of future transformer PR | open | Decide from official/native key-shape dumps; Diffusers layout alone is not proof. | +| Q004 | Does the native transformer require conversion? | conversion | Phase 5 | resolved | No. Pinned HF index, official production meta model, and FastVideo production meta model expose the same 521 keys; official/native key shapes match exactly. | ## Issues And Blockers @@ -77,13 +79,14 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension | I001 | Phase 6 | text_encoder | medium | The full checkpoint includes `lm_head.weight`, while the native FastVideo encoder owns only the body. | Native parity allowlists exactly `lm_head.weight`; production uses body-only Transformers `AutoModel`. | parity | in_progress | Implementation is complete; resolve only after the pinned fp32/bf16 rerun. | | I002 | Phase 7 | scheduler/pipeline | high | Published `scheduler_config.json` omits the reference-schedule flag and official zero endpoint. | Pinned pipeline mutates `sigma_min=0.0`; parity requires both scheduler values. | pipeline | open | Future pipeline config must serialize `use_reference_discrete_timesteps=True` and `sigma_min=0.0`. | | I003 | Phase 6 | tokenizer/text_encoder | medium | Prompt formatting and target length previously differed from the official call. | Pinned source uses thinking chat template, 512 tokens, and `hidden_states[-2]`. | parity | in_progress | Test/config behavior corrected; pinned asset rerun required before resolving. | -| I004 | Phase 4 | transformer | high | `ZImageTransformer2DModel` is not ported. | No FastVideo target or parity test exists. | port | open | Separate future PR. | +| I004 | Phase 4 | transformer | high | `ZImageTransformer2DModel` was not ported. | Native config/model/registry and deterministic parity test now exist. | port | resolved | Production meta surface matches all 521 pinned HF keys; tiny full-forward parity passes exactly. | | I005 | Phase 7 | pipeline | high | No FastVideo pipeline/config/preset/registry/example or quality regression exists. | Component-only PR scope. | pipeline | open | Separate future PR after component gates pass. | | I006 | Phase 6 | text_encoder | high | Production loading previously used a causal-LM object and could override CPU placement, while parity exercised only direct native construction. | Production path is now body-only `AutoModel`; loader records `_fastvideo_input_device`; stage moves tokens to that device; parity uses a distinct reference instance. | encoder | in_progress | Implementation complete; pinned production-loader rerun required. | | I007 | Phase 6 | text_encoder | high | Lenient native loading could miss one fused source shard or an auxiliary quantization scale. | Loader completeness now covers every destination, required split shard, exact fused destination, and mapped scale parameter. | encoder | in_progress | Asset-free strictness regressions and pinned checkpoint load must pass before resolving. | | I008 | Phase 6 | vae | medium | Direct decode alone did not cover production registry/config/device/strict-load behavior; the shared target is a runtime Diffusers wrapper. | VAE parity scope is now `both` and includes `VAELoader` plus official normalization. | parity | in_progress | Existing Diffusers `AutoencoderKL` wrapper exception is accepted only for this component-only PR; pinned rerun required. | | I009 | Phase 2 | reference assets | high | Adding a nonexistent clone path to `sys.path` could import another installed `zimage` package and yield false parity. | Tests now verify exact clone HEAD and module paths under `Z-Image/src`; `/Z-Image/` is ignored. | prep | resolved | Wrong SHA/path/import is a failure; absent local assets may skip. | | I010 | Phase 6 | tokenizer | medium | Missing `apply_chat_template` previously skipped and could hide an incompatible tokenizer. | Tokenizer parity now asserts the API and exact official kwargs. | parity | resolved | Missing chat-template behavior fails. Asset-backed numerical rerun remains I003. | +| I011 | Phase 6 | transformer | high | Production real-weight strict load and full-forward parity are not verified. | The pinned transformer is 24.6 GB and is not present locally; the CUDA test scaffold therefore skipped. | parity | open | Run `test_zimage_transformer_production_loader_forward_parity` non-skip with the pinned HF transformer subfolder on CUDA and record drift. | ## Escape Hatches @@ -94,6 +97,7 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension | Date | Decision | Rationale | Impact | |---|---|---|---| +| 2026-07-13 | Preserve the published transformer key surface and use raw masked SDPA at SP=1. | The checkpoint already matches the official class, while padded variable-length attention needs a key mask unavailable in current distributed wrappers. | No converter is needed; runtime rejects SP>1 until a mask-aware distributed path has parity evidence. | | 2026-07-04 | Pin both external sources and treat all earlier GPU results as historical until rerun. | Earlier evidence did not record an immutable HF revision, and source imports were not origin-checked. | README has executable clone/download commands; tests enforce the source SHA; all asset-backed parity rows are revalidation-required. | | 2026-07-04 | Use a body-only Transformers `AutoModel` for production and a separate `AutoModel` object as the parity oracle. | Every FastVideo consumer needs hidden states, not full-vocabulary logits or an LM head; comparing an object with itself is not independent evidence. | Production avoids the LM head while parity covers production loading and native implementation separately. | | 2026-07-04 | Make fused-source and quantization-scale completeness part of native Qwen strict loading. | A loaded destination name alone did not prove that every required Q/K/V or gate/up shard and auxiliary scale arrived. | Missing split shards or mapped scale parameters fail; an exact fused destination remains valid. | @@ -115,6 +119,7 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension fused/split source completeness, and auxiliary quantization scales. - VAE revalidation must cover both direct implementation decode and production `VAELoader` behavior with official scaling/shift normalization. -- Next separate-PR work: transformer prototype/parity, conversion decision, - pipeline/config/preset/registry/example, pipeline parity, and image-quality - regression. +- Transformer production key/shape and tiny output parity pass. The real-weight + production-loader/full-forward CUDA test remains I011 and must pass non-skip. +- Next work after I011: pipeline/config/preset/registry/example, pipeline parity, + and image-quality regression. diff --git a/tests/local_tests/zimage/README.md b/tests/local_tests/zimage/README.md index b1c4e3c38..584d38535 100644 --- a/tests/local_tests/zimage/README.md +++ b/tests/local_tests/zimage/README.md @@ -1,10 +1,10 @@ # Z-Image Local Tests Local-only component coverage for the `zimage` FastVideo port (`T2I`, -Z-Image-Turbo only). This component-only PR covers the Qwen3 text encoder, -tokenizer, AutoencoderKL VAE, and FlowMatchEulerDiscreteScheduler. The native -`ZImageTransformer2DModel`, pipeline, pipeline parity, example, and image-quality -regression remain out of scope. +Z-Image-Turbo only). Coverage includes the Qwen3 text encoder, tokenizer, +AutoencoderKL VAE, FlowMatchEulerDiscreteScheduler, and a native +`ZImageTransformer2DModel`. Pipeline wiring, pipeline parity, examples, and +image-quality regression remain out of scope. > **Status:** DRAFT port, in progress. Every test that consumes the reference > clone or HF assets requires a fresh non-skip rerun against the immutable pins @@ -18,14 +18,14 @@ regression remain out of scope. |---|---| | Model family / variant | `zimage` / `Z-Image-Turbo` | | Workload | `T2I` | -| Component-only PR scope | text encoder, tokenizer, VAE, scheduler | +| Component scope | text encoder, tokenizer, VAE, scheduler, native transformer | | Official reference | `https://github.com/Tongyi-MAI/Z-Image` | | Local reference dir | `/Z-Image/src` | | Official commit | `26f23eda626ffadda020b04ff79488e1d72004cd` | | HF weights | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` | -| Local weights dir | `/official_weights/Z-Image/` (`text_encoder/`, `tokenizer/`, `vae/`, `scheduler/`) | +| Local weights dir | `/official_weights/Z-Image/` (component subfolders; transformer is 24.6 GB) | | Source layout | Diffusers-style per-component subfolders | -| Conversion | not needed for the reused components in this PR; the future native transformer/full-pipeline decision will come from native and official key/shape prototypes | +| Conversion | not needed; the native and official production transformer expose the same 521 checkpoint keys/shapes | | HF token env | `HF_TOKEN` (the pinned repository is public; never record a token value) | The pinned text-encoder config declares `Qwen3ForCausalLM` with 36 decoder @@ -69,14 +69,18 @@ python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py \ --allow-pattern 'scheduler/*' ``` +The real-weight transformer test additionally needs `transformer/*` from the +same pinned revision (24.6 GB). Do not download it without approving that cost; +point `ZIMAGE_TRANSFORMER_DIR` at an existing pinned transformer subfolder. + Do not change core dependency versions (`torch`, `diffusers`, `transformers`, `flash-attn`, `triton`, or CUDA packages) without explicit approval. ```text dependency_changes: none -official_env_status: not_verified +official_env_status: transformer source import and tiny CPU parity verified; asset-backed components pending private_dep_stubs: none -blocked_on: pinned-snapshot non-skip component reruns +blocked_on: pinned-snapshot component reruns and transformer real-weight CUDA parity ``` ## Run the tests @@ -91,6 +95,7 @@ pytest tests/local_tests/zimage/ -v -s | Tokenizer ([`test_zimage_tokenizer_parity.py`](./test_zimage_tokenizer_parity.py)) | `production_loader`; exact `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)` and `max_length=512`; absence of `apply_chat_template` is a failure, not a skip | REVALIDATION REQUIRED (historical unpinned PASS) | | VAE ([`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py)) | `both`; direct implementation decode plus production `VAELoader`; production check applies `(latents / scaling_factor) + shift_factor` before decode | REVALIDATION REQUIRED (historical unpinned direct-decode PASS only) | | Text encoder ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `both`; independent Transformers `AutoModel` reference, body-only production loader, and FastVideo-native implementation; fp32/bf16 output checks plus fused/split/quant-scale strictness | REVALIDATION REQUIRED (historical unpinned native PASS only) | +| Transformer ([`test_zimage_transformer_parity.py`](./test_zimage_transformer_parity.py)) | `both`; pinned production meta key/shape surface, additive-mask attention parity, deterministic tiny CPU full-forward parity, explicit `sp_size=1` guard, and real-weight production-loader/full-forward scaffold | TINY PARITY PASS (4 passed); REAL-WEIGHT PARITY PENDING (24.6 GB weights absent) | ## Component contracts @@ -133,10 +138,24 @@ its `num_steps + 1` schedule. A future Z-Image pipeline config must set both `use_reference_discrete_timesteps=True` and `sigma_min=0.0`; the published `scheduler_config.json` alone does not encode the complete runtime contract. +### Transformer + +- The FastVideo-native transformer preserves all 521 published checkpoint keys. + The sorted key digest from the pinned HF safetensors index is + `3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354`. +- Tiny CPU parity uses the pinned official class, identical deterministic state, + variable-length image/text batches, and concrete full-forward output checks. +- Attention uses raw torch SDPA because the padded variable-length stream needs + a key mask not exposed by FastVideo's distributed wrappers. Runtime rejects + initialized sequence-parallel world sizes above one. +- The production-loader/full-forward test runs when the pinned transformer + subfolder and CUDA are available; its non-skip result is still required. + ## Remaining work - Run every asset-backed component test non-skip against the pinned clone and HF snapshot before claiming component parity. -- Port and validate `ZImageTransformer2DModel` in a separate PR. +- Run the transformer production-loader/full-forward parity test non-skip with + the pinned 24.6 GB transformer weights on CUDA. - Add the Z-Image pipeline/config/preset/registry/example, pipeline parity, and image-quality regression after all component gates pass. diff --git a/tests/local_tests/zimage/test_zimage_transformer_parity.py b/tests/local_tests/zimage/test_zimage_transformer_parity.py new file mode 100644 index 000000000..471b84d18 --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_transformer_parity.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Z-Image transformer parity. + +Coverage scope: both. The tiny CPU tests compare the pinned implementation and +FastVideo model directly. The production-loader test activates when the real +transformer subfolder is available. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys +from types import ModuleType + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.configs.models.dits import ZImageDiTConfig +from fastvideo.models.dits.zimage import ZImageAttention, ZImageTransformer2DModel, _prepare_attention_mask + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_REF_DIR = Path(os.getenv("ZIMAGE_OFFICIAL_REF_DIR", REPO_ROOT / "Z-Image")) +OFFICIAL_SRC = OFFICIAL_REF_DIR / "src" +TRANSFORMER_DIR = Path( + os.getenv("ZIMAGE_TRANSFORMER_DIR", REPO_ROOT / "official_weights" / "Z-Image" / "transformer")) +REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" +HF_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6" +HF_KEY_COUNT = 521 +HF_KEY_SHA256 = "3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354" +PARITY_SCOPE = "both" + +PRODUCTION_CONFIG = { + "all_patch_size": (2, ), + "all_f_patch_size": (1, ), + "in_channels": 16, + "dim": 3840, + "n_layers": 30, + "n_refiner_layers": 2, + "n_heads": 30, + "n_kv_heads": 30, + "norm_eps": 1e-5, + "qk_norm": True, + "cap_feat_dim": 2560, + "rope_theta": 256.0, + "t_scale": 1000.0, + "axes_dims": (32, 48, 48), + "axes_lens": (1536, 512, 512), +} + +TINY_CONFIG = { + "all_patch_size": (2, ), + "all_f_patch_size": (1, ), + "in_channels": 2, + "dim": 12, + "n_layers": 1, + "n_refiner_layers": 1, + "n_heads": 2, + "n_kv_heads": 2, + "norm_eps": 1e-5, + "qk_norm": True, + "cap_feat_dim": 8, + "rope_theta": 256.0, + "t_scale": 1000.0, + "axes_dims": (2, 2, 2), + "axes_lens": (192, 16, 16), +} + + +def _load_source_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + pytest.fail(f"Cannot build import spec for {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def reference_transformer_cls(): + source_file = OFFICIAL_SRC / "zimage" / "transformer.py" + if not OFFICIAL_REF_DIR.exists(): + pytest.skip(f"Pinned Z-Image reference clone not found: {OFFICIAL_REF_DIR}") + if not source_file.is_file(): + pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}") + + try: + revision = subprocess.run( + ["git", "-C", str(OFFICIAL_REF_DIR), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + pytest.fail(f"Cannot verify Z-Image reference revision: {exc}") + assert revision == REFERENCE_REVISION, ( + f"expected Z-Image reference {REFERENCE_REVISION}, got {revision}") + + module_names = ( + "config", + "config.inference", + "config.model", + "utils", + "utils.import_utils", + "utils.attention", + "_zimage_reference_transformer", + ) + saved_modules = {name: sys.modules.get(name) for name in module_names} + added_path = str(OFFICIAL_SRC) not in sys.path + if added_path: + sys.path.insert(0, str(OFFICIAL_SRC)) + + try: + for name in module_names: + sys.modules.pop(name, None) + utils_package = ModuleType("utils") + utils_package.__path__ = [str(OFFICIAL_SRC / "utils")] + sys.modules["utils"] = utils_package + _load_source_module("utils.import_utils", OFFICIAL_SRC / "utils" / "import_utils.py") + _load_source_module("utils.attention", OFFICIAL_SRC / "utils" / "attention.py") + module = _load_source_module("_zimage_reference_transformer", source_file) + assert Path(module.__file__).resolve() == source_file.resolve() + yield module.ZImageTransformer2DModel + finally: + for name in module_names: + sys.modules.pop(name, None) + if saved_modules[name] is not None: + sys.modules[name] = saved_modules[name] + if added_path: + sys.path.remove(str(OFFICIAL_SRC)) + + +def _build_fastvideo(config_values: dict) -> ZImageTransformer2DModel: + config = ZImageDiTConfig() + config.update_model_arch(config_values) + return ZImageTransformer2DModel( + config=config, + hf_config={"_class_name": "ZImageTransformer2DModel", **config_values}, + ) + + +def _state_shapes(model: torch.nn.Module) -> dict[str, tuple[int, ...]]: + return {name: tuple(tensor.shape) for name, tensor in model.state_dict().items()} + + +def test_zimage_transformer_production_meta_key_surface(reference_transformer_cls): + with torch.device("meta"): + reference = reference_transformer_cls(**PRODUCTION_CONFIG) + fastvideo = _build_fastvideo(PRODUCTION_CONFIG) + + reference_shapes = _state_shapes(reference) + fastvideo_shapes = _state_shapes(fastvideo) + assert fastvideo_shapes == reference_shapes + keys = sorted(reference_shapes) + key_digest = hashlib.sha256(("\n".join(keys) + "\n").encode()).hexdigest() + assert len(keys) == HF_KEY_COUNT + assert key_digest == HF_KEY_SHA256 + print(f"Z-Image production transformer key/shape surface: {len(reference_shapes)} tensors") + + +def test_zimage_transformer_tiny_cpu_parity(reference_transformer_cls): + reference = reference_transformer_cls(**TINY_CONFIG).eval() + generator = torch.Generator(device="cpu").manual_seed(123) + with torch.no_grad(): + for parameter in reference.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) + + fastvideo = _build_fastvideo(TINY_CONFIG).eval() + assert _state_shapes(fastvideo) == _state_shapes(reference) + fastvideo.load_state_dict(reference.state_dict(), strict=True) + + generator.manual_seed(456) + images = [ + torch.randn((2, 1, 4, 4), generator=generator), + torch.randn((2, 1, 8, 10), generator=generator), + ] + caption_features = [ + torch.randn((3, 8), generator=generator), + torch.randn((35, 8), generator=generator), + ] + timestep = torch.tensor([0.25, 0.75], dtype=torch.float32) + + with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + reference_outputs = reference(images, timestep, caption_features)[0] + positional_outputs = fastvideo(images, caption_features, timestep)[0] + keyword_outputs = fastvideo( + hidden_states=images, + encoder_hidden_states=caption_features, + timestep=timestep, + )[0] + + assert len(positional_outputs) == len(keyword_outputs) == len(reference_outputs) == 2 + for positional, actual, expected in zip(positional_outputs, keyword_outputs, reference_outputs): + assert actual.shape == expected.shape + assert expected.abs().max() > 0 + assert_close(positional, expected, atol=1e-6, rtol=1e-6) + assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + +def test_zimage_attention_additive_mask_parity(reference_transformer_cls): + reference_module = sys.modules[reference_transformer_cls.__module__] + reference = reference_module.ZImageAttention(12, 2, 2).eval() + generator = torch.Generator(device="cpu").manual_seed(321) + with torch.no_grad(): + for parameter in reference.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) + + fastvideo = ZImageAttention(12, 2, 2).eval() + fastvideo.load_state_dict(reference.state_dict(), strict=True) + hidden_states = torch.randn((2, 5, 12), generator=generator) + attention_mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, True]]) + additive_mask = _prepare_attention_mask(attention_mask, hidden_states.dtype) + assert additive_mask is not None + assert additive_mask.dtype == hidden_states.dtype + assert additive_mask.shape == (2, 1, 1, 5) + assert additive_mask[0, 0, 0, 0] == 0 + assert torch.isneginf(additive_mask[0, 0, 0, -1]) + + with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + expected = reference(hidden_states, attention_mask=attention_mask) + actual = fastvideo(hidden_states, attention_mask=attention_mask) + assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + +def test_zimage_transformer_rejects_sequence_parallelism(monkeypatch: pytest.MonkeyPatch): + model = _build_fastvideo(TINY_CONFIG) + monkeypatch.setattr("fastvideo.models.dits.zimage.model_parallel_is_initialized", lambda: True) + monkeypatch.setattr("fastvideo.models.dits.zimage.get_sp_world_size", lambda: 2) + + with pytest.raises(NotImplementedError, match="sp_size=1"): + model( + hidden_states=[torch.zeros(2, 1, 4, 4)], + encoder_hidden_states=[torch.zeros(3, 8)], + timestep=torch.tensor([0.5]), + ) + + +def _has_real_weights() -> bool: + return any(TRANSFORMER_DIR.glob("*.safetensors")) + + +def _load_real_config() -> dict: + with (TRANSFORMER_DIR / "config.json").open(encoding="utf-8") as file: + config = json.load(file) + assert config.pop("_class_name") == "ZImageTransformer2DModel" + config.pop("_diffusers_version", None) + config.pop("_name_or_path", None) + for key, expected in PRODUCTION_CONFIG.items(): + actual = config[key] + if isinstance(expected, tuple): + actual = tuple(actual) + assert actual == expected, f"unexpected production config {key}={actual!r}" + assert set(config) == set(PRODUCTION_CONFIG) + return config + + +def _run_fastvideo_production(inputs: tuple, dtype: torch.dtype) -> list[torch.Tensor]: + from fastvideo.configs.pipelines.base import PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import TransformerLoader + + precision = "bf16" if dtype == torch.bfloat16 else "fp16" + args = FastVideoArgs( + model_path=str(TRANSFORMER_DIR), + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, + pipeline_config=PipelineConfig(dit_config=ZImageDiTConfig(), dit_precision=precision), + ) + model = TransformerLoader().load(str(TRANSFORMER_DIR), args).eval() + assert isinstance(model, ZImageTransformer2DModel) + assert args.model_paths["transformer"] == str(TRANSFORMER_DIR) + assert next(model.parameters()).device.type == "cuda" + with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + outputs = model( + hidden_states=inputs[0], + encoder_hidden_states=inputs[1], + timestep=inputs[2], + )[0] + result = [output.detach().float().cpu() for output in outputs] + del model + torch.cuda.empty_cache() + return result + + +def _run_reference_production(reference_transformer_cls, config: dict, inputs: tuple, + dtype: torch.dtype) -> list[torch.Tensor]: + from safetensors.torch import load_file + + state_dict = {} + for shard in sorted(TRANSFORMER_DIR.glob("*.safetensors")): + state_dict.update(load_file(str(shard), device="cpu")) + with torch.device("meta"): + model = reference_transformer_cls(**config) + model.load_state_dict(state_dict, strict=True, assign=True) + del state_dict + model = model.to(device="cuda", dtype=dtype).eval() + with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + outputs = model(inputs[0], inputs[2], inputs[1])[0] + result = [output.detach().float().cpu() for output in outputs] + del model + torch.cuda.empty_cache() + return result + + +@pytest.mark.skipif( + not _has_real_weights(), + reason=( + "Pinned Z-Image transformer weights are absent; set ZIMAGE_TRANSFORMER_DIR " + f"to Tongyi-MAI/Z-Image-Turbo@{HF_REVISION}/transformer"), +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Real-weight Z-Image transformer parity requires CUDA") +def test_zimage_transformer_production_loader_forward_parity(reference_transformer_cls): + config = _load_real_config() + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + generator = torch.Generator(device="cuda").manual_seed(789) + inputs = ( + [torch.randn((16, 1, 4, 4), device="cuda", dtype=dtype, generator=generator)], + [torch.randn((3, 2560), device="cuda", dtype=dtype, generator=generator)], + torch.tensor([0.25], device="cuda", dtype=dtype), + ) + + fastvideo_outputs = _run_fastvideo_production(inputs, dtype) + reference_outputs = _run_reference_production(reference_transformer_cls, config, inputs, dtype) + assert len(fastvideo_outputs) == len(reference_outputs) == 1 + + actual, expected = fastvideo_outputs[0], reference_outputs[0] + diff = (actual - expected).abs() + mean_scale = expected.abs().mean().clamp_min(1e-6) + print(f"Z-Image real-weight parity: max={diff.max():.6f}, mean={diff.mean():.6f}") + assert diff.mean() / mean_scale < 0.01 + assert_close(actual, expected, atol=1e-2, rtol=1e-2)