[feat]: add native Z-Image transformer prototype

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

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"

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

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 = {

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: `<repo_root>/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 `<weights>/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 | `<weights>/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=<pinned-clone> 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.

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 | `<repo_root>/Z-Image/src` |
 | Official commit | `26f23eda626ffadda020b04ff79488e1d72004cd` |
 | HF weights | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` |
-| Local weights dir | `<repo_root>/official_weights/Z-Image/` (`text_encoder/`, `tokenizer/`, `vae/`, `scheduler/`) |
+| Local weights dir | `<repo_root>/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.

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)