diff --git a/.gitignore b/.gitignore index 6fbe89528..0fb04721d 100644 --- a/.gitignore +++ b/.gitignore @@ -34,6 +34,7 @@ env *.log weights/ logs/ +/Z-Image/ official_weights/ converted_weights/ diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index cf4eb7386..00863afb5 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -333,15 +333,25 @@ surfaces: use_distill: sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig] scheduler_arch: - sources: [fastvideo.configs.pipelines.sd35.SD35Config] + sources: + - fastvideo.configs.pipelines.sd35.SD35Config + - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig text_encoder_archs: - sources: [fastvideo.configs.pipelines.sd35.SD35Config] + sources: + - fastvideo.configs.pipelines.sd35.SD35Config + - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig tokenizer_archs: - sources: [fastvideo.configs.pipelines.sd35.SD35Config] + sources: + - fastvideo.configs.pipelines.sd35.SD35Config + - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig transformer_arch: - sources: [fastvideo.configs.pipelines.sd35.SD35Config] + sources: + - fastvideo.configs.pipelines.sd35.SD35Config + - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig vae_arch: - sources: [fastvideo.configs.pipelines.sd35.SD35Config] + sources: + - fastvideo.configs.pipelines.sd35.SD35Config + - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig expand_timesteps: sources: - fastvideo.configs.pipelines.wan.FastWan2_2_TI2V_5B_Config @@ -421,6 +431,8 @@ surfaces: frame_receptive_field: "MagiHuman internal data-proxy receptive-field setting." image_conditioning: "MagiHuman preset variant marker for reference-image conditioning." ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset." + scheduler_sigma_min: "Z-Image scheduler parity invariant; not part of the public typed inference API." + scheduler_use_reference_discrete_timesteps: "Z-Image scheduler parity invariant; not part of the public typed inference API." sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection." text_offset: "MagiHuman internal data-proxy text alignment offset." vae_stride: "MagiHuman internal VAE/data-proxy stride setting." @@ -447,6 +459,7 @@ surfaces: output_video_name: request.output.output_video_name num_videos_per_prompt: request.sampling.num_videos_per_prompt seed: request.sampling.seed + max_sequence_length: request.sampling.max_sequence_length num_frames: request.sampling.num_frames height: request.sampling.height width: request.sampling.width @@ -457,6 +470,8 @@ surfaces: num_inference_steps_sr: request.sampling.num_inference_steps_sr guidance_scale: request.sampling.guidance_scale guidance_scale_2: request.sampling.guidance_scale_2 + cfg_normalization: request.sampling.cfg_normalization + cfg_truncation: request.sampling.cfg_truncation guidance_rescale: request.sampling.guidance_rescale use_embedded_guidance: request.sampling.use_embedded_guidance true_cfg_scale: request.sampling.true_cfg_scale @@ -508,7 +523,6 @@ surfaces: internal_only: data_type: "Derived from the request shape and not a public input." latents: "Pre-generated diffusion latents supplied by parity/debug harnesses; not a public input." - max_sequence_length: "Model-specific text-encoder sequence cap; not part of the public typed inference API." sampling_param_extensions: {} diff --git a/examples/inference/basic/basic_zimage.py b/examples/inference/basic/basic_zimage.py new file mode 100644 index 000000000..0ecbf5d87 --- /dev/null +++ b/examples/inference/basic/basic_zimage.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run Z-Image-Turbo text-to-image generation through FastVideo. + +User story: + "I want the official Z-Image-Turbo defaults and a deterministic PNG from + a local or Hugging Face checkpoint." +""" + +import argparse +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api import ( + EngineConfig, + GenerationRequest, + GeneratorConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) + + +DEFAULT_PROMPT = ( + "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. " + "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. " + "Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, " + "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights." +) +DEFAULT_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run Z-Image-Turbo text-to-image generation.") + parser.add_argument("--model-path", default="Tongyi-MAI/Z-Image-Turbo") + parser.add_argument("--revision", default=DEFAULT_REVISION) + parser.add_argument("--output", default="outputs/zimage/zimage_turbo.png") + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--steps", type=int, default=8) + parser.add_argument("--guidance-scale", type=float, default=0.0) + parser.add_argument("--max-sequence-length", type=int, default=512) + parser.add_argument("--cfg-normalization", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--cfg-truncation", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=args.model_path, + revision=args.revision, + engine=EngineConfig( + num_gpus=1, + parallelism=ParallelismConfig(tp_size=1, sp_size=1), + use_fsdp_inference=False, + ), + # The model registry selects the native zimage_turbo preset. + pipeline=PipelineSelection(workload_type="t2i"), + )) + try: + generator.generate( + GenerationRequest( + prompt=args.prompt, + negative_prompt=args.negative_prompt, + sampling=SamplingConfig( + height=args.height, + width=args.width, + num_frames=1, + fps=1, + num_inference_steps=args.steps, + guidance_scale=args.guidance_scale, + max_sequence_length=args.max_sequence_length, + cfg_normalization=args.cfg_normalization, + cfg_truncation=args.cfg_truncation, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output), + save_video=True, + return_frames=False, + ), + )) + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 89d437cc7..77a9bd8a3 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -90,6 +90,11 @@ class SamplingParam: num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 guidance_scale_2: float | None = None + # Z-Image CFG controls. ``cfg_normalization=True`` caps the guided + # prediction norm at the positive-prediction norm; ``cfg_truncation`` + # disables CFG above the normalized-noise threshold. + cfg_normalization: bool = False + cfg_truncation: float | None = 1.0 # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG. use_embedded_guidance: bool = False # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding). @@ -323,6 +328,18 @@ class SamplingParam: default=SamplingParam.guidance_scale, help="Classifier-free guidance scale", ) + parser.add_argument( + "--cfg-normalization", + action=StoreBoolean, + default=SamplingParam.cfg_normalization, + help="Cap Z-Image CFG prediction norm to the positive-prediction norm", + ) + parser.add_argument( + "--cfg-truncation", + type=float, + default=SamplingParam.cfg_truncation, + help="Disable Z-Image CFG above this normalized-noise threshold", + ) parser.add_argument( "--guidance-rescale", type=float, diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index d3dbb0516..8c5745247 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -138,6 +138,7 @@ class InputConfig: class SamplingConfig: num_videos_per_prompt: int = 1 seed: int = 1024 + max_sequence_length: int | None = None num_frames: int = 125 height: int = 720 width: int = 1280 @@ -148,6 +149,8 @@ class SamplingConfig: num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 guidance_scale_2: float | None = None + cfg_normalization: bool = False + cfg_truncation: float | None = 1.0 guidance_rescale: float = 0.0 true_cfg_scale: float | None = None use_embedded_guidance: bool | None = None 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/configs/models/encoders/base.py b/fastvideo/configs/models/encoders/base.py index 1ce6167f2..09ee57067 100644 --- a/fastvideo/configs/models/encoders/base.py +++ b/fastvideo/configs/models/encoders/base.py @@ -83,6 +83,7 @@ class TextEncoderConfig(EncoderConfig): arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig) is_chat_model: bool = False treat_empty_as_dot: bool = False + chat_template_enable_thinking: bool = field(default=False, kw_only=True) @dataclass diff --git a/fastvideo/configs/pipelines/zimage.py b/fastvideo/configs/pipelines/zimage.py new file mode 100644 index 000000000..c4395fac8 --- /dev/null +++ b/fastvideo/configs/pipelines/zimage.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import EncoderConfig +from fastvideo.configs.models.dits.zimage import ZImageDiTConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig +from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text + + +def _zimage_text_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor: + if outputs.hidden_states is None: + raise RuntimeError("Z-Image requires Qwen3 hidden states") + return outputs.hidden_states[-2] + + +@dataclass +class ZImagePipelineConfig(PipelineConfig): + """Configuration for the native Z-Image text-to-image pipeline.""" + + scheduler_arch: str = "FlowMatchEulerDiscreteScheduler" + transformer_arch: str = "ZImageTransformer2DModel" + vae_arch: str = "AutoencoderKL" + text_encoder_archs: tuple[str, ...] = ("Qwen3Model", ) + tokenizer_archs: tuple[str, ...] = ("Qwen2Tokenizer", ) + + dit_config: ZImageDiTConfig = field(default_factory=ZImageDiTConfig) + vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (Qwen3TextConfig(chat_template_enable_thinking=True), )) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (preprocess_text, )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: (_zimage_text_postprocess, )) + + dit_precision: str = "bf16" + vae_precision: str = "fp32" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) + vae_tiling: bool = False + vae_sp: bool = False + + embedded_cfg_scale: float = 0.0 + flow_shift: float | None = 3.0 + scheduler_step_in_fp32: bool = True + scheduler_sigma_min: float = 0.0 + scheduler_use_reference_discrete_timesteps: bool = True 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/encoders/qwen3.py b/fastvideo/models/encoders/qwen3.py index 59461bf56..e664463ad 100644 --- a/fastvideo/models/encoders/qwen3.py +++ b/fastvideo/models/encoders/qwen3.py @@ -327,17 +327,15 @@ class Qwen3ForCausalLM(TextEncoder): dtype: torch.dtype, device: torch.device, ) -> nn.Module: - from transformers import AutoModelForCausalLM + from transformers import AutoModel - if device.type == "cpu" and torch.cuda.is_available(): - from fastvideo.distributed import get_local_torch_device - - device = get_local_torch_device() - - return AutoModelForCausalLM.from_pretrained( + # FastVideo uses Qwen3 only as a text encoder. Loading the body avoids + # materializing an unused LM head and full-vocabulary logits, including + # for checkpoints whose metadata names Qwen3ForCausalLM. + return AutoModel.from_pretrained( model_path, local_files_only=True, - torch_dtype=dtype, + dtype=dtype, low_cpu_mem_usage=True, ).eval().to(device) @@ -368,9 +366,18 @@ class Qwen3ForCausalLM(TextEncoder): residual = None if position_ids is None: + # Expand to [batch_size, seq_len]: the rotary layer flattens + # positions to ``num_tokens`` and reshapes q/k to + # ``(num_tokens, -1, head_dim)``. A bare [1, seq_len] only matches + # ``num_tokens`` when batch_size == 1; for batched inputs it folds + # the batch dim into the head dim and misaligns RoPE. Expanding to + # ``batch_size * seq_len`` tokens keeps the layout correct. position_ids = torch.arange( - 0, hidden_states.shape[1], device=hidden_states.device - ).unsqueeze(0) + 0, + hidden_states.shape[1], + device=hidden_states.device, + dtype=torch.long, + ).unsqueeze(0).expand(hidden_states.shape[0], -1) all_hidden_states: tuple[Any, ...] | None = ( () if output_hidden_states else None @@ -405,6 +412,20 @@ class Qwen3ForCausalLM(TextEncoder): ) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + stacked_params_mapping = self.config.arch_config.stacked_params_mapping + # A fused destination is initialized either by one already-fused tensor + # or after every split source projection has been loaded. Include + # auxiliary quantization parameters (for example scale_weight) rather + # than limiting completeness checks to weight/bias tensors. + expected_stacked_shards = { + (name, shard_id) + for name in params_dict + for param_name, _, shard_id in stacked_params_mapping + if param_name in name + } + fused_param_names = {name for name, _ in expected_stacked_shards} + loaded_stacked_shards: set[tuple[str, str | int]] = set() + loaded_fused_params: set[str] = set() for name, loaded_weight in weights: if name.startswith("model."): @@ -423,37 +444,66 @@ class Qwen3ForCausalLM(TextEncoder): continue name = kv_scale_name - for ( - param_name, - weight_name, - shard_id, - ) in self.config.arch_config.stacked_params_mapping: + matched_stacked_param = False + for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue - name = name.replace(weight_name, param_name) + matched_stacked_param = True + target_name = name.replace(weight_name, param_name) - if name.endswith(".bias") and name not in params_dict: - continue + if target_name.endswith(".bias") and target_name not in params_dict: + break - if name not in params_dict: - continue + if target_name not in params_dict: + break - param = params_dict[name] + param = params_dict[target_name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) + loaded_params.add(target_name) + loaded_stacked_shards.add((target_name, shard_id)) break + + if matched_stacked_param: + continue + + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + if name in fused_param_names and name.endswith(".scale_weight"): + # Merged scale loaders interpret a missing shard id as shard 0. + # An exact fused key is already a complete vector, so copy it + # atomically and retain the default loader's shape validation. + default_weight_loader(param, loaded_weight) else: - if name.endswith(".bias") and name not in params_dict: - continue - - if name not in params_dict: - continue - - param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) - loaded_params.add(name) + if name in fused_param_names: + loaded_fused_params.add(name) + + required_split_shards = { + (name, shard_id) + for name, shard_id in expected_stacked_shards + if name not in loaded_fused_params + } + missing_stacked_shards = required_split_shards - loaded_stacked_shards + if missing_stacked_shards: + formatted_missing = ", ".join( + f"{name}[{shard_id}]" + for name, shard_id in sorted( + missing_stacked_shards, + key=lambda item: (item[0], str(item[1])), + ) + ) + raise ValueError( + "Missing required stacked checkpoint shards: " + f"{formatted_missing}" + ) return loaded_params diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 6b097c317..a6cbd9293 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -391,18 +391,22 @@ class TextEncoderLoader(ComponentLoader): model_config.quant_config = quant_cls() with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): - with target_device: - architectures = getattr(model_config, "architectures", []) - model_cls, _ = ModelRegistry.resolve_model_cls(architectures) - if getattr(model_cls, "supports_hf_from_pretrained", False): - model = model_cls.from_pretrained_local( # type: ignore[attr-defined] - model_path, - model_config, # type: ignore[arg-type] - dtype=PRECISION_TO_TYPE[dtype], - device=target_device, - ) - return model.eval() + architectures = getattr(model_config, "architectures", []) + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + if getattr(model_cls, "supports_hf_from_pretrained", False): + model = model_cls.from_pretrained_local( # type: ignore[attr-defined] + model_path, + model_config, # type: ignore[arg-type] + dtype=PRECISION_TO_TYPE[dtype], + device=target_device, + ) + # HF passthrough encoders return before FastVideo's FSDP + # wrapping path, so the text stage needs their placement to + # put token tensors on the same device. + model._fastvideo_input_device = target_device + return model.eval() + with target_device: model = model_cls(model_config) # type: ignore weights_to_load = {name for name, _ in model.named_parameters()} diff --git a/fastvideo/models/loader/utils.py b/fastvideo/models/loader/utils.py index 033cd0058..d3bbc3ce1 100644 --- a/fastvideo/models/loader/utils.py +++ b/fastvideo/models/loader/utils.py @@ -18,8 +18,10 @@ def set_default_torch_dtype(dtype: torch.dtype): """Sets the default torch dtype to the given dtype.""" old_dtype = torch.get_default_dtype() torch.set_default_dtype(dtype) - yield - torch.set_default_dtype(old_dtype) + try: + yield + finally: + torch.set_default_dtype(old_dtype) def get_param_names_mapping( diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index fdbc4fc2e..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 = { @@ -78,6 +79,9 @@ _TEXT_ENCODER_MODELS = { "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"), "Qwen2_5_VLForConditionalGeneration": ("encoders", "reason1", "Reason1TextEncoder"), + # Z-Image-Turbo's text_encoder/config.json declares architecture + # "Qwen3Model"; route it to the shared Qwen3 encoder (added for Flux2 Klein). + "Qwen3Model": ("encoders", "qwen3", "Qwen3ForCausalLM"), "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"), "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"), "Mistral3ForConditionalGeneration": diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 5e0d65853..b4f1935e1 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -96,6 +96,14 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, The minimum sigma value for the noise schedule. sigma_data (`float`, *optional*): The sigma data value for scaling. + use_reference_discrete_timesteps (`bool`, defaults to False): + Some reference schedulers (e.g. Z-Image) construct the timestep + schedule by linspacing `num_inference_steps + 1` points from + `t_max` to `t_min` and dropping the terminal point. Default + (`False`) preserves the original `np.linspace(t_max, t_min, + num_inference_steps)` (float64) behaviour used by every existing + model. Enable this flag only when matching a reference scheduler + that expects the +1 + drop-terminal construction. """ _compatibles: list[Any] = [] @@ -122,6 +130,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, sigma_max: float | None = None, sigma_min: float | None = None, sigma_data: float | None = None, + use_reference_discrete_timesteps: bool = False, ): if sum([ self.config.use_beta_sigmas, self.config.use_exponential_sigmas, @@ -155,7 +164,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, self.sigmas = sigmas.to( "cpu") # to avoid too much CPU/GPU communication - self.sigma_min = self.sigmas[-1].item() + self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() BaseScheduler.__init__(self) @@ -350,7 +359,19 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, if timesteps_array is None: t_max = self._sigma_to_t(self.sigma_max) t_min = self._sigma_to_t(self.sigma_min) - timesteps_array = np.linspace(t_max, t_min, num_inference_steps) + if self.config.use_reference_discrete_timesteps: + # Some reference schedulers (for example Z-Image) build a + # float64 num_steps+1 linspace and drop the terminal point. + timesteps_array = np.linspace( + t_max, + t_min, + num_inference_steps + 1, + )[:-1] + else: + # Preserve the original numpy default (float64) here — + # casting to float32 silently shifts rounded timestep + # values for every existing model that uses this branch. + timesteps_array = np.linspace(t_max, t_min, num_inference_steps) sigmas_array = timesteps_array / self.config.num_train_timesteps else: sigmas_array = np.array(sigmas).astype(np.float32) diff --git a/fastvideo/pipelines/__init__.py b/fastvideo/pipelines/__init__.py index abafbe48c..ec6d5708b 100644 --- a/fastvideo/pipelines/__init__.py +++ b/fastvideo/pipelines/__init__.py @@ -35,7 +35,7 @@ def build_pipeline(fastvideo_args: FastVideoArgs, """ # Get pipeline type model_path = fastvideo_args.model_path - model_path = maybe_download_model(model_path) + model_path = maybe_download_model(model_path, revision=fastvideo_args.revision) # fastvideo_args.downloaded_model_path = model_path logger.info("Model path: %s", model_path) diff --git a/fastvideo/pipelines/basic/zimage/__init__.py b/fastvideo/pipelines/basic/zimage/__init__.py new file mode 100644 index 000000000..9428b7d51 --- /dev/null +++ b/fastvideo/pipelines/basic/zimage/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Z-Image pipeline.""" + +from fastvideo.pipelines.basic.zimage.zimage_pipeline import ZImagePipeline + +__all__ = ["ZImagePipeline"] diff --git a/fastvideo/pipelines/basic/zimage/presets.py b/fastvideo/pipelines/basic/zimage/presets.py new file mode 100644 index 000000000..59e2e554c --- /dev/null +++ b/fastvideo/pipelines/basic/zimage/presets.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Z-Image inference presets.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Z-Image denoising pass", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + "cfg_normalization", + "cfg_truncation", + }), +) + +ZIMAGE_TURBO = InferencePreset( + name="zimage_turbo", + version=1, + model_family="zimage", + description="Z-Image-Turbo text-to-image generation", + workload_type="t2i", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 1024, + "width": 1024, + "num_frames": 1, + "fps": 1, + "seed": 42, + "guidance_scale": 0.0, + "num_inference_steps": 8, + "negative_prompt": "", + "max_sequence_length": 512, + "cfg_normalization": False, + "cfg_truncation": 1.0, + }, +) + +ALL_PRESETS = (ZIMAGE_TURBO, ) diff --git a/fastvideo/pipelines/basic/zimage/stages.py b/fastvideo/pipelines/basic/zimage/stages.py new file mode 100644 index 000000000..0c33fe437 --- /dev/null +++ b/fastvideo/pipelines/basic/zimage/stages.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline stages for the native Z-Image text-to-image path.""" + +from __future__ import annotations + +import inspect + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.hooks.activation_trace import trace_step +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.input_validation import InputValidationStage +from fastvideo.utils import PRECISION_TO_TYPE + + +class ZImageInputValidationStage(InputValidationStage): + """Validate the image geometry and reproduce the official device RNG.""" + + def _generate_seeds(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None: + del fastvideo_args + assert batch.seed is not None + batch.seeds = [batch.seed] + device = get_local_torch_device() + batch.generator = torch.Generator(device=device).manual_seed(batch.seed) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.do_classifier_free_guidance and batch.negative_prompt is None and not batch.negative_prompt_embeds: + batch.negative_prompt = "" + batch = super().forward(batch, fastvideo_args) + if batch.num_frames != 1: + raise ValueError(f"Z-Image is text-to-image and requires num_frames=1, got {batch.num_frames}") + if batch.height is None or batch.width is None: + raise ValueError("Z-Image requires height and width") + if batch.height % 16 or batch.width % 16: + raise ValueError("Z-Image height and width must be divisible by 16; " + f"got {batch.height}x{batch.width}") + return batch + + +class ZImageConditioningStage(PipelineStage): + """Trim padded Qwen states and materialize variable-length CFG streams.""" + + @staticmethod + def _trim_embeddings( + embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + ) -> list[torch.Tensor]: + if attention_mask is None: + return list(embeds.unbind(0)) + return [ + sample[mask.to(device=sample.device, dtype=torch.bool)] + for sample, mask in zip(embeds, attention_mask, strict=True) + ] + + @staticmethod + def _repeat(items: list[torch.Tensor], count: int) -> list[torch.Tensor]: + return [item for item in items for _ in range(count)] + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + del fastvideo_args + if len(batch.prompt_embeds) != 1: + raise ValueError(f"Z-Image expects one text encoder, got {len(batch.prompt_embeds)}") + + prompt_mask = batch.prompt_attention_mask[0] if batch.prompt_attention_mask else None + prompt_embeds = self._trim_embeddings(batch.prompt_embeds[0], prompt_mask) + batch.extra["zimage_prompt_embeds"] = self._repeat(prompt_embeds, batch.num_videos_per_prompt) + + if batch.do_classifier_free_guidance: + if not batch.negative_prompt_embeds: + raise ValueError("Z-Image CFG requires negative prompt embeddings") + negative_mask = batch.negative_attention_mask[0] if batch.negative_attention_mask else None + negative_embeds = self._trim_embeddings(batch.negative_prompt_embeds[0], negative_mask) + batch.extra["zimage_negative_prompt_embeds"] = self._repeat( + negative_embeds, + batch.num_videos_per_prompt, + ) + else: + batch.extra["zimage_negative_prompt_embeds"] = [] + return batch + + +class ZImageLatentPreparationStage(PipelineStage): + """Create the official fp32 image latents on the transformer device.""" + + def __init__(self, transformer) -> None: + self.transformer = transformer + + @staticmethod + def _randn( + shape: tuple[int, ...], + generators: torch.Generator | list[torch.Generator] | None, + device: torch.device, + ) -> torch.Tensor: + if isinstance(generators, list): + if len(generators) != shape[0]: + raise ValueError(f"generator list length {len(generators)} does not match batch size {shape[0]}") + sample_shape = (1, *shape[1:]) + return torch.cat([ + torch.randn(sample_shape, generator=generator, device=device, dtype=torch.float32) + for generator in generators + ]) + return torch.randn(shape, generator=generators, device=device, dtype=torch.float32) + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.height is None or batch.width is None: + raise ValueError("Z-Image requires height and width before latent preparation") + + prompt_embeds = batch.extra.get("zimage_prompt_embeds") + if not isinstance(prompt_embeds, list) or not prompt_embeds: + raise ValueError("Z-Image conditioning must run before latent preparation") + + channels = int(getattr(self.transformer, "in_channels", 16)) + spatial_ratio = int(fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio) + shape = ( + len(prompt_embeds), + channels, + 1, + batch.height // spatial_ratio, + batch.width // spatial_ratio, + ) + device = get_local_torch_device() + if batch.latents is None: + latents = self._randn(shape, batch.generator, device) + else: + latents = batch.latents + if latents.ndim == 4: + latents = latents.unsqueeze(2) + if tuple(latents.shape) != shape: + raise ValueError(f"Expected Z-Image latents with shape {shape}, got {tuple(latents.shape)}") + latents = latents.to(device=device, dtype=torch.float32) + + batch.latents = latents + batch.raw_latent_shape = shape + return batch + + +class ZImageTimestepPreparationStage(PipelineStage): + """Apply the native scheduler's zero endpoint and discrete schedule.""" + + def __init__(self, scheduler) -> None: + self.scheduler = scheduler + + @staticmethod + def _calculate_shift( + image_seq_len: int, + base_seq_len: int, + max_seq_len: int, + base_shift: float, + max_shift: float, + ) -> float: + slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) + return image_seq_len * slope + base_shift - slope * base_seq_len + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.latents is None: + raise ValueError("Z-Image latents must be prepared before timesteps") + if batch.timesteps is not None and batch.sigmas is not None: + raise ValueError("Only one of timesteps or sigmas may be supplied") + + scheduler = self.scheduler + sigma_min = float(fastvideo_args.pipeline_config.scheduler_sigma_min) + use_reference_timesteps = bool(fastvideo_args.pipeline_config.scheduler_use_reference_discrete_timesteps) + scheduler.sigma_min = sigma_min + scheduler.register_to_config( + sigma_min=sigma_min, + use_reference_discrete_timesteps=use_reference_timesteps, + ) + config = scheduler.config + image_seq_len = (batch.latents.shape[-2] // 2) * (batch.latents.shape[-1] // 2) + mu = self._calculate_shift( + image_seq_len, + int(config.get("base_image_seq_len", 256)), + int(config.get("max_image_seq_len", 4096)), + float(config.get("base_shift", 0.5)), + float(config.get("max_shift", 1.15)), + ) + device = get_local_torch_device() + + if batch.timesteps is not None: + if "timesteps" not in inspect.signature(scheduler.set_timesteps).parameters: + raise ValueError(f"{type(scheduler).__name__} does not accept custom timesteps") + scheduler.set_timesteps(timesteps=batch.timesteps, device=device, mu=mu) + elif batch.sigmas is not None: + if "sigmas" not in inspect.signature(scheduler.set_timesteps).parameters: + raise ValueError(f"{type(scheduler).__name__} does not accept custom sigmas") + scheduler.set_timesteps(sigmas=batch.sigmas, device=device, mu=mu) + else: + scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu) + + batch.timesteps = scheduler.timesteps + batch.num_inference_steps = len(batch.timesteps) + return batch + + +class ZImageDenoisingStage(PipelineStage): + """Run the native Z-Image flow-matching loop.""" + + performance_component_metric = "transformer_time_s" + + def __init__(self, transformer, scheduler) -> None: + self.transformer = transformer + self.scheduler = scheduler + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.latents is None or batch.timesteps is None: + raise ValueError("Z-Image denoising requires latents and timesteps") + + latents = batch.latents.float() + positive = batch.extra.get("zimage_prompt_embeds") + negative = batch.extra.get("zimage_negative_prompt_embeds", []) + if not isinstance(positive, list) or not positive: + raise ValueError("Z-Image denoising requires prompt embeddings") + + device = get_local_torch_device() + target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] + batch_size = latents.shape[0] + + for index, timestep_value in enumerate(batch.timesteps): + if timestep_value.item() == 0 and index == len(batch.timesteps) - 1: + continue + + timestep = timestep_value.expand(batch_size).to(device=device, dtype=torch.float32) + timestep = (1000.0 - timestep) / 1000.0 + current_guidance_scale = float(batch.guidance_scale) + if (batch.do_classifier_free_guidance and batch.cfg_truncation is not None and batch.cfg_truncation <= 1.0 + and timestep[0].item() > batch.cfg_truncation): + current_guidance_scale = 0.0 + apply_cfg = batch.do_classifier_free_guidance and current_guidance_scale > 0.0 + + if apply_cfg: + if not isinstance(negative, list) or len(negative) != batch_size: + raise ValueError("Z-Image CFG requires one negative embedding per image") + model_latents = latents.to(target_dtype).repeat(2, 1, 1, 1, 1) + model_embeddings = positive + negative + model_timestep = timestep.repeat(2) + else: + model_latents = latents.to(target_dtype) + model_embeddings = positive + model_timestep = timestep + + with ( + torch.autocast( + device_type=device.type, + enabled=False, + ), + trace_step(index), + set_forward_context( + current_timestep=index, + attn_metadata=None, + forward_batch=batch, + ), + ): + model_outputs = self.transformer( + hidden_states=model_latents, + encoder_hidden_states=model_embeddings, + timestep=model_timestep, + )[0] + + if apply_cfg: + positive_outputs = model_outputs[:batch_size] + negative_outputs = model_outputs[batch_size:] + guided_outputs = [] + for positive_output, negative_output in zip( + positive_outputs, + negative_outputs, + strict=True, + ): + positive_fp32 = positive_output.float() + prediction = positive_fp32 + current_guidance_scale * (positive_fp32 - negative_output.float()) + if batch.cfg_normalization: + positive_norm = torch.linalg.vector_norm(positive_fp32) + prediction_norm = torch.linalg.vector_norm(prediction) + if prediction_norm > positive_norm: + prediction = prediction * (positive_norm / prediction_norm) + guided_outputs.append(prediction) + noise_pred = torch.stack(guided_outputs) + else: + noise_pred = torch.stack([output.float() for output in model_outputs]) + + noise_pred = -noise_pred + latents = self.scheduler.step( + noise_pred, + timestep_value, + latents, + return_dict=False, + )[0].float() + + batch.latents = latents + return batch + + +class ZImageDecodingStage(PipelineStage): + """Apply the official latent transform and decode one image frame.""" + + performance_component_metric = "vae_decode_time_s" + + def __init__(self, vae) -> None: + self.vae = vae + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.latents is None: + raise ValueError("Z-Image decoding requires latents") + if fastvideo_args.output_type == "latent": + # FastVideo standardizes image and video latents as [B,C,T,H,W]. + # Tongyi's image-only API returns the equivalent tensor with T squeezed. + batch.output = batch.latents + return batch + + latents = batch.latents + if latents.ndim != 5 or latents.shape[2] != 1: + raise ValueError(f"Expected Z-Image latents [B,C,1,H,W], got {tuple(latents.shape)}") + + device = get_local_torch_device() + self.vae = self.vae.to(device) + vae_dtype = getattr(self.vae, "dtype", None) + if vae_dtype is None: + vae_dtype = next(self.vae.parameters()).dtype + config = self.vae.config + scaling_factor = float(config.scaling_factor) + shift_factor = float(config.shift_factor or 0.0) + latents_2d = latents.squeeze(2).to(device=device, dtype=vae_dtype) + latents_2d = latents_2d / scaling_factor + shift_factor + decoded = self.vae.decode(latents_2d, return_dict=False)[0] + decoded = (decoded / 2 + 0.5).clamp(0, 1) + batch.output = decoded.unsqueeze(2).float() + + if fastvideo_args.vae_cpu_offload: + self.vae.to("cpu") + return batch diff --git a/fastvideo/pipelines/basic/zimage/zimage_pipeline.py b/fastvideo/pipelines/basic/zimage/zimage_pipeline.py new file mode 100644 index 000000000..ecf897c78 --- /dev/null +++ b/fastvideo/pipelines/basic/zimage/zimage_pipeline.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages.text_encoding import TextEncodingStage + +from .stages import ( + ZImageConditioningStage, + ZImageDecodingStage, + ZImageDenoisingStage, + ZImageInputValidationStage, + ZImageLatentPreparationStage, + ZImageTimestepPreparationStage, +) + + +class ZImagePipeline(ComposedPipelineBase): + """Native Z-Image text-to-image pipeline.""" + + pipeline_config_cls: type[ZImagePipelineConfig] = ZImagePipelineConfig + _required_config_modules = [ + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + scheduler = self.get_module("scheduler") + transformer = self.get_module("transformer") + + self.add_stage("input_validation_stage", ZImageInputValidationStage()) + self.add_stage( + "text_encoding_stage", + TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + self.add_stage("zimage_conditioning_stage", ZImageConditioningStage()) + self.add_stage( + "latent_preparation_stage", + ZImageLatentPreparationStage(transformer=transformer), + ) + self.add_stage( + "timestep_preparation_stage", + ZImageTimestepPreparationStage(scheduler=scheduler), + ) + self.add_stage( + "denoising_stage", + ZImageDenoisingStage(transformer=transformer, scheduler=scheduler), + ) + self.add_stage( + "decoding_stage", + ZImageDecodingStage(vae=self.get_module("vae")), + ) + + +EntryClass = ZImagePipeline diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 962a0a525..d6076668f 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -303,7 +303,8 @@ class ComposedPipelineBase(ABC): self.modules[module_name] = module def _load_config(self, model_path: str) -> dict[str, Any]: - model_path = maybe_download_model(self.model_path) + revision = getattr(self.fastvideo_args, "revision", None) + model_path = maybe_download_model(self.model_path, revision=revision) self.model_path = model_path # fastvideo_args.downloaded_model_path = model_path logger.info("Model path: %s", model_path) diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 243bd2ae0..d9a74e07b 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -178,6 +178,8 @@ class ForwardBatch: num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 guidance_scale_2: float | None = None + cfg_normalization: bool = False + cfg_truncation: float | None = 1.0 guidance_rescale: float = 0.0 eta: float = 0.0 sigmas: list[float] | None = None diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 82c5dfb22..15a8e27a4 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -73,25 +73,36 @@ class DecodingStage(PipelineStage): cfg = getattr(self.vae, "config", None) - # Matrix-Game 2.0-style: z = z * std + mean - if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")): - latents_mean = torch.tensor(cfg.latents_mean, device=latents.device, + # Matrix-Game 2.0-style: z = z * std + mean. Some configs declare these + # fields as None, so test the values rather than mere attribute presence + # -- otherwise torch.tensor(None) raises instead of falling through. + latents_mean_value = getattr(cfg, "latents_mean", None) if cfg is not None else None + latents_std_value = getattr(cfg, "latents_std", None) if cfg is not None else None + if latents_mean_value is not None and latents_std_value is not None: + latents_mean = torch.tensor(latents_mean_value, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1) - latents_std = torch.tensor(cfg.latents_std, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(latents_std_value, device=latents.device, + dtype=latents.dtype).view(1, -1, 1, 1, 1) return latents * latents_std + latents_mean - # Diffusers-style: scaling_factor (+ optional shift_factor) - if hasattr(self.vae, "scaling_factor"): - if isinstance(self.vae.scaling_factor, torch.Tensor): - latents = latents / self.vae.scaling_factor.to(latents.device, latents.dtype) + # Diffusers-style: scaling_factor (+ optional shift_factor), read from + # the module only. Nearly every VAE *config* also carries a + # scaling_factor whose decode() already accounts for it, so sourcing it + # from cfg here would newly divide latents for VAEs that have always + # passed through untouched (flux2, hunyuanvideo, hunyuanvideo15, ltx2). + scaling_factor = getattr(self.vae, "scaling_factor", None) + if scaling_factor is not None: + if isinstance(scaling_factor, torch.Tensor): + latents = latents / scaling_factor.to(latents.device, latents.dtype) else: - latents = latents / self.vae.scaling_factor + latents = latents / scaling_factor - if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: - if isinstance(self.vae.shift_factor, torch.Tensor): - latents = latents + self.vae.shift_factor.to(latents.device, latents.dtype) + shift_factor = getattr(self.vae, "shift_factor", None) + if shift_factor is not None: + if isinstance(shift_factor, torch.Tensor): + latents = latents + shift_factor.to(latents.device, latents.dtype) else: - latents = latents + self.vae.shift_factor + latents = latents + shift_factor return latents diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 0a5427ee8..264506902 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -73,6 +73,7 @@ class TextEncodingStage(PipelineStage): fastvideo_args, encoder_index=all_indices, return_attention_mask=True, + max_length=batch.max_sequence_length, ) if self._last_audio_embeds is not None: batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds @@ -91,6 +92,7 @@ class TextEncodingStage(PipelineStage): fastvideo_args, encoder_index=all_indices, return_attention_mask=True, + max_length=batch.max_sequence_length, ) if self._last_audio_embeds is not None: batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds @@ -221,6 +223,20 @@ class TextEncodingStage(PipelineStage): encoder_device = torch.device(target_device) moved_for_forward = True + # An explicit `device=` wins. Otherwise follow the encoder's real + # param device. Once it has been moved for the forward that is the + # target device, and an HF-passthrough encoder's + # _fastvideo_input_device (stamped at load, e.g. "cpu" under + # text_encoder_cpu_offload) is stale -- honouring it would feed cpu + # token ids to cuda weights. The marker only speaks when nothing + # moved and the module has no parameters to speak for it. + if device is not None: + input_device = torch.device(target_device) + elif moved_for_forward: + input_device = encoder_device + else: + input_device = getattr(text_encoder, "_fastvideo_input_device", encoder_device) + tok_kwargs = dict(encoder_config.tokenizer_kwargs) if max_length is not None: tok_kwargs["max_length"] = max_length @@ -264,7 +280,7 @@ class TextEncodingStage(PipelineStage): # pre-format prompts into message lists upstream and rely on # the inner tokenizer + full tokenizer_kwargs (which include # add_generation_prompt). Preserve that original path exactly. - text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(encoder_device) + text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(input_device) else: # Two-step approach matching Diffusers: format with chat # template first, then tokenize the resulting strings. @@ -275,12 +291,12 @@ class TextEncodingStage(PipelineStage): messages, tokenize=False, add_generation_prompt=True, - enable_thinking=False, + enable_thinking=encoder_config.chat_template_enable_thinking, ) formatted_texts.append(formatted) - text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(encoder_device) + text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(input_device) else: - text_inputs = tok(processed_texts, **tok_kwargs).to(encoder_device) + text_inputs = tok(processed_texts, **tok_kwargs).to(input_device) input_ids = text_inputs["input_ids"] attention_mask = text_inputs["attention_mask"] diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 30d885937..fc46d4b28 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -62,6 +62,7 @@ from fastvideo.configs.pipelines.glm_image import GlmImageConfig from fastvideo.configs.pipelines.flux import FluxPipelineConfig from fastvideo.configs.pipelines.sd35 import SD35Config from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig) +from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig from fastvideo.api.sampling_param import SamplingParam from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam @@ -1104,6 +1105,19 @@ def _register_configs() -> None: ], ) + # Z-Image-Turbo + register_configs( + sampling_param_cls=None, + pipeline_config_cls=ZImagePipelineConfig, + workload_types=(WorkloadType.T2I, ), + hf_model_paths=["Tongyi-MAI/Z-Image-Turbo"], + model_detectors=[ + lambda path: "zimagepipeline" in path or "z-image" in path or "z_image" in path, + ], + model_family="zimage", + default_preset="zimage_turbo", + ) + # --- Part 3: Main Resolver --- @@ -1226,6 +1240,8 @@ def _register_presets() -> None: ALL_PRESETS as TURBODIFFUSION_PRESETS, ) from fastvideo.pipelines.basic.wan.presets import ( ALL_PRESETS as WAN_PRESETS, ) + from fastvideo.pipelines.basic.zimage.presets import ( + ALL_PRESETS as ZIMAGE_PRESETS, ) from fastvideo.pipelines.basic.flux_2.presets import ( ALL_PRESETS as FLUX2_PRESETS, ) @@ -1248,6 +1264,7 @@ def _register_presets() -> None: STABLE_AUDIO_PRESETS, TURBODIFFUSION_PRESETS, WAN_PRESETS, + ZIMAGE_PRESETS, ) for group in all_preset_groups: for preset in group: diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index 64da28dc8..1db450963 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -171,6 +171,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "sampling": { "num_videos_per_prompt": 1, "seed": 1024, + "max_sequence_length": None, "num_frames": 16, "height": 720, "width": 1280, @@ -181,6 +182,8 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "num_inference_steps_sr": 50, "guidance_scale": 1.0, "guidance_scale_2": None, + "cfg_normalization": False, + "cfg_truncation": 1.0, "guidance_rescale": 0.0, "true_cfg_scale": None, "use_embedded_guidance": None, diff --git a/fastvideo/tests/ssim/reference_videos_cli.py b/fastvideo/tests/ssim/reference_videos_cli.py index 3c6c8c419..8b9c99ec6 100644 --- a/fastvideo/tests/ssim/reference_videos_cli.py +++ b/fastvideo/tests/ssim/reference_videos_cli.py @@ -18,7 +18,7 @@ VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv") # pixel-space models (e.g. LTX-2 distilled). Keeping them in the same upload # flow means seeding a new test only requires one HF round-trip. LATENT_EXTENSIONS = (".pt",) -REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS +REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS + (".png",) HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN") HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO" HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE" @@ -51,10 +51,10 @@ def _default_repo_type() -> str: def _iter_reference_files(root: Path) -> Iterable[Path]: - """Yield video and latent (.pt) references under `root`. + """Yield supported references under `root`. Used by copy-local and the "has local references" marker probe so that - latent-only tests (no mp4) still satisfy readiness checks. + image- and latent-only tests (no mp4) still satisfy readiness checks. """ for path in root.rglob("*"): if path.is_file() and path.suffix.lower() in REFERENCE_EXTENSIONS: diff --git a/fastvideo/tests/ssim/test_reference_videos_cli.py b/fastvideo/tests/ssim/test_reference_videos_cli.py index 8c740d9da..20a51d4e6 100644 --- a/fastvideo/tests/ssim/test_reference_videos_cli.py +++ b/fastvideo/tests/ssim/test_reference_videos_cli.py @@ -44,6 +44,22 @@ def _fake_snapshot_download(*, allow_patterns, cache_dir, **_kwargs): return str(snapshot_root) +def test_png_references_are_copied_and_mark_local_cache_ready(tmp_path): + generated_dir = tmp_path / "generated" + generated = generated_dir / "model-a" / "flash" / "sample.png" + generated.parent.mkdir(parents=True) + generated.write_bytes(b"image") + reference_dir = tmp_path / "reference_videos" / "default" / "L40S_reference_videos" + + assert reference_videos_cli.copy_generated_to_reference( + generated_dir, + reference_dir, + ) == 1 + marker = reference_dir.parent / ".download_complete_default" + marker.touch() + assert reference_videos_cli._has_local_reference_videos(tmp_path, "default") + + @pytest.fixture(autouse=True) def _fake_hf(monkeypatch): _FakeHfApi.instances = [] diff --git a/fastvideo/tests/ssim/test_zimage_similarity.py b/fastvideo/tests/ssim/test_zimage_similarity.py new file mode 100644 index 000000000..2a96d7efe --- /dev/null +++ b/fastvideo/tests/ssim/test_zimage_similarity.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PNG regression for Tongyi-MAI/Z-Image@26f23eda's native Turbo recipe.""" + +from __future__ import annotations + +import os + +import pytest +import torch + +from fastvideo.logger import init_logger +from fastvideo.tests.ssim.inference_similarity_utils import ( + run_text_to_video_similarity_test, +) +from fastvideo.tests.ssim.reference_utils import ( + get_cuda_device_name, + resolve_device_reference_folder, +) + +logger = init_logger(__name__) + +REQUIRED_GPUS = 1 +ZIMAGE_MIN_SSIM = 0.98 + +ZIMAGE_MODEL_PATH = os.getenv( + "ZIMAGE_MODEL_DIR", + "Tongyi-MAI/Z-Image-Turbo", +) +ZIMAGE_MODEL_REVISION = os.getenv( + "ZIMAGE_MODEL_REVISION", + "f332072aa78be7aecdf3ee76d5c247082da564a6", +) + +TEST_PROMPTS = [ + "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. " + "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. " + "Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, " + "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.", +] + +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H100", "H100"), + ("H200", "H200"), + ("B200", "B200"), + ), + device_name=get_cuda_device_name(), + fallback_device_prefix="L40S", + logger=logger, +) + +ZIMAGE_MODEL_TO_PARAMS: dict[str, dict[str, object]] = { + "Tongyi-MAI__Z-Image-Turbo": { + "num_gpus": 1, + "model_path": ZIMAGE_MODEL_PATH, + "sp_size": 1, + "tp_size": 1, + "height": 1024, + "width": 1024, + "num_frames": 1, + "fps": 1, + "num_inference_steps": 8, + "guidance_scale": 0.0, + "seed": 42, + }, +} + +ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS = { + model_id: dict(params) + for model_id, params in ZIMAGE_MODEL_TO_PARAMS.items() +} + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Z-Image SSIM test requires CUDA", +) +@pytest.mark.parametrize("prompt", TEST_PROMPTS) +@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) +@pytest.mark.parametrize("model_id", list(ZIMAGE_MODEL_TO_PARAMS)) +def test_zimage_similarity( + prompt: str, + attention_backend_name: str, + model_id: str, +) -> None: + is_hf_repo = "/" in ZIMAGE_MODEL_PATH and not ZIMAGE_MODEL_PATH.startswith("/") + if not is_hf_repo and not os.path.isdir(ZIMAGE_MODEL_PATH): + pytest.skip( + f"Z-Image weights not found at {ZIMAGE_MODEL_PATH} " + "(set ZIMAGE_MODEL_DIR to override)" + ) + + run_text_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + attention_backend_name=attention_backend_name, + model_id=model_id, + default_params_map=ZIMAGE_MODEL_TO_PARAMS, + full_quality_params_map=ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS, + min_acceptable_ssim=ZIMAGE_MIN_SSIM, + media_extension=".png", + init_kwargs_override={ + "revision": ZIMAGE_MODEL_REVISION, + "workload_type": "t2i", + "trust_remote_code": True, + "use_fsdp_inference": False, + "dit_cpu_offload": False, + "text_encoder_cpu_offload": False, + "vae_cpu_offload": False, + "image_encoder_cpu_offload": False, + "pin_cpu_memory": False, + }, + generation_kwargs_override={ + "save_video": True, + "negative_prompt": "", + "max_sequence_length": 512, + "cfg_normalization": False, + "cfg_truncation": 1.0, + }, + ) diff --git a/fastvideo/tests/stages/test_text_encoding.py b/fastvideo/tests/stages/test_text_encoding.py index fa7383368..25213d0ab 100644 --- a/fastvideo/tests/stages/test_text_encoding.py +++ b/fastvideo/tests/stages/test_text_encoding.py @@ -2,6 +2,7 @@ import torch import types import pytest +from fastvideo.configs.models.encoders.clip import CLIPTextArchConfig, CLIPTextConfig from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.configs.pipelines.base import PipelineConfig @@ -13,6 +14,9 @@ class TensorDict(dict): return TensorDict({k: v.to(device) for k, v in self.items()}) class FakeTokenizer: + def __init__(self): + self.last_chat_template_kwargs = None + def __call__(self, texts, **kwargs): B = len(texts) seq_len = int(kwargs.get("max_length", 4)) @@ -21,6 +25,10 @@ class FakeTokenizer: "attention_mask": torch.ones(B, seq_len, dtype=torch.long), }) + def apply_chat_template(self, messages, **kwargs): + self.last_chat_template_kwargs = kwargs + return "formatted prompt" + class FakeChatTokenizer: def __init__(self): @@ -40,17 +48,22 @@ class FakeChatTokenizer: "attention_mask": torch.ones(B, seq_len, dtype=torch.long), }) - class FakeTextEncoder(torch.nn.Module): def __init__(self, hidden_size=8): super().__init__() self.hidden_size = hidden_size + self.last_input_device = None self.last_output_hidden_states = None def forward(self, input_ids, attention_mask, output_hidden_states=False): + self.last_input_device = input_ids.device self.last_output_hidden_states = bool(output_hidden_states) B, T = input_ids.shape - last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size) + last_hidden_state = torch.arange( + B * T * self.hidden_size, + dtype=torch.float32, + device=input_ids.device, + ).view(B, T, self.hidden_size) hidden_states = (last_hidden_state, ) if output_hidden_states else None return types.SimpleNamespace(last_hidden_state=last_hidden_state, hidden_states=hidden_states) @@ -213,3 +226,104 @@ def test_chat_list_preprocess_output_is_not_stripped(): {"role": "user", "content": "a robotic arm welding a metal structure"}, ]] assert tokenizer.last_kwargs["return_tensors"] == "pt" + + +def test_chat_template_thinking_is_keyword_only_and_preserves_subclass_positions(): + arch = CLIPTextArchConfig() + config = CLIPTextConfig( + arch, + "legacy-prefix", + None, + None, + False, + False, + 7, + False, + True, + False, + ) + + assert config.num_hidden_layers_override == 7 + assert config.require_post_norm is False + assert config.enable_scale is True + assert config.is_causal is False + assert config.chat_template_enable_thinking is False + + configured = CLIPTextConfig(chat_template_enable_thinking=True) + assert configured.chat_template_enable_thinking is True + + with pytest.raises(TypeError): + CLIPTextConfig( + arch, + "legacy-prefix", + None, + None, + False, + False, + 7, + False, + True, + False, + True, + ) + + +@pytest.mark.parametrize("enable_thinking", [False, True]) +def test_encode_text_forwards_chat_template_thinking_config(enable_thinking): + fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8) + encoder_config = fastvideo_args.pipeline_config.text_encoder_configs[0] + encoder_config.is_chat_model = True + encoder_config.chat_template_enable_thinking = enable_thinking + + stage = make_stage(num_encoders=1, hidden_size=hidden) + stage.encode_text("a", fastvideo_args, encoder_index=[0]) + + assert stage.tokenizers[0].last_chat_template_kwargs == { + "tokenize": False, + "add_generation_prompt": True, + "enable_thinking": enable_thinking, + } + + +def test_encode_text_uses_hf_passthrough_input_device(monkeypatch): + fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8) + stage = make_stage(num_encoders=1, hidden_size=hidden) + stage.text_encoders[0]._fastvideo_input_device = torch.device("cpu") + + monkeypatch.setattr( + "fastvideo.pipelines.stages.text_encoding.get_local_torch_device", + lambda: torch.device("meta"), + ) + + batch = ForwardBatch( + data_type="video", + prompt="a", + do_classifier_free_guidance=False, + prompt_embeds=[], + negative_prompt_embeds=None, + prompt_attention_mask=[], + negative_attention_mask=None, + ) + output = stage.forward(batch, fastvideo_args) + + # The marker governs where the encoder *receives* its tokens... + assert stage.text_encoders[0].last_input_device == torch.device("cpu") + # ...while the stage still normalizes the embeds it returns onto the + # caller's target device. + assert output.prompt_embeds[0].device.type == "meta" + + +def test_encode_text_explicit_device_overrides_hf_passthrough_marker(): + fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8) + stage = make_stage(num_encoders=1, hidden_size=hidden) + stage.text_encoders[0]._fastvideo_input_device = torch.device("cpu") + + output = stage.encode_text( + "a", + fastvideo_args, + encoder_index=[0], + device=torch.device("meta"), + ) + + assert stage.text_encoders[0].last_input_device == torch.device("meta") + assert output[0].device.type == "meta" diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 37ebc0140..c0f66cd17 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -493,7 +493,12 @@ def import_pynvml(): return pynvml -def maybe_download_model(model_name_or_path: str, local_dir: str | None = None, download: bool = True) -> str: +def maybe_download_model( + model_name_or_path: str, + local_dir: str | None = None, + download: bool = True, + revision: str | None = None, +) -> str: """ Check if the model path is a Hugging Face Hub model ID and download it if needed. @@ -510,6 +515,7 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None, ``org/repo/subfolder`` umbrella-repo reference. local_dir: Local directory to save the model download: Whether to download the model from Hugging Face Hub + revision: Optional immutable Hub revision. Returns: Local path to the model (or to the subfolder inside the snapshot). @@ -550,6 +556,7 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None, repo_id=repo_id, allow_patterns=[f"{subfolder}/**"], local_dir=local_dir, + revision=revision, ) # Defense-in-depth: ensure the resolved subfolder path stays # inside the snapshot root and that snapshot_download actually @@ -569,7 +576,8 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None, with get_lock(model_name_or_path): local_path = snapshot_download(repo_id=model_name_or_path, ignore_patterns=["*.onnx", "*.msgpack"], - local_dir=local_dir) + local_dir=local_dir, + revision=revision) logger.info("Downloaded model to %s", local_path) return str(local_path) except Exception as e: diff --git a/tests/local_tests/pipelines/test_zimage_pipeline_parity.py b/tests/local_tests/pipelines/test_zimage_pipeline_parity.py new file mode 100644 index 000000000..50fef280c --- /dev/null +++ b/tests/local_tests/pipelines/test_zimage_pipeline_parity.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native Z-Image stage contract and pinned-repository pipeline parity.""" + +from __future__ import annotations + +import gc +import importlib +import os +import subprocess +import sys +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from fastvideo.pipelines.basic.zimage.stages import ( + ZImageConditioningStage, + ZImageDecodingStage, + ZImageDenoisingStage, + ZImageInputValidationStage, + ZImageLatentPreparationStage, + ZImageTimestepPreparationStage, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" +REFERENCE_REPO = Path(os.getenv("ZIMAGE_REFERENCE_REPO", REPO_ROOT / "Z-Image")) +MODEL_DIR = Path(os.getenv("ZIMAGE_MODEL_DIR", REPO_ROOT / "official_weights" / "Z-Image")) + + +class _ConstantTransformer: + in_channels = 16 + + def __init__(self) -> None: + self.calls: list[tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]] = [] + + def __call__( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: list[torch.Tensor], + timestep: torch.Tensor, + ) -> tuple[list[torch.Tensor], dict]: + self.calls.append((hidden_states, encoder_hidden_states, timestep)) + return [torch.ones_like(sample) for sample in hidden_states.unbind(0)], {} + + +class _RecordingVAE(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace(scaling_factor=0.3611, shift_factor=0.1159) + self.decode_input: torch.Tensor | None = None + + def decode(self, latents: torch.Tensor, return_dict: bool = False) -> tuple[torch.Tensor]: + assert return_dict is False + self.decode_input = latents.detach().clone() + batch, _, height, width = latents.shape + return (torch.zeros(batch, 3, height * 8, width * 8, device=latents.device),) + + +def _stage_args(*, output_type: str = "pil") -> SimpleNamespace: + return SimpleNamespace( + pipeline_config=ZImagePipelineConfig(), + output_type=output_type, + vae_cpu_offload=False, + disable_autocast=False, + ) + + +def test_zimage_native_default_stage_math(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise the pinned native 1-step formulas without model assets.""" + autocast_calls: list[dict[str, object]] = [] + + @contextmanager + def record_autocast(**kwargs): + autocast_calls.append(kwargs) + yield + + monkeypatch.setattr( + "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device", + lambda: torch.device("cpu"), + ) + monkeypatch.setattr("fastvideo.pipelines.basic.zimage.stages.torch.autocast", record_autocast) + args = _stage_args() + transformer = _ConstantTransformer() + scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) + vae = _RecordingVAE() + batch = ForwardBatch( + data_type="image", + prompt="test prompt", + negative_prompt="", + prompt_embeds=[torch.arange(32, dtype=torch.float32).reshape(1, 4, 8)], + prompt_attention_mask=[torch.tensor([[True, True, False, False]])], + height=64, + width=64, + num_frames=1, + num_inference_steps=1, + guidance_scale=0.0, + seed=42, + latents=torch.zeros(1, 16, 1, 8, 8), + ) + + ZImageInputValidationStage().forward(batch, args) + ZImageConditioningStage().forward(batch, args) + ZImageLatentPreparationStage(transformer).forward(batch, args) + ZImageTimestepPreparationStage(scheduler).forward(batch, args) + + assert len(batch.extra["zimage_prompt_embeds"]) == 1 + assert batch.extra["zimage_prompt_embeds"][0].shape == (2, 8) + assert_close(batch.timesteps, torch.tensor([1000.0])) + assert scheduler.sigma_min == 0.0 + assert scheduler.config.use_reference_discrete_timesteps is True + + ZImageDenoisingStage(transformer, scheduler).forward(batch, args) + assert autocast_calls == [{"device_type": "cpu", "enabled": False}] + assert_close(batch.latents, torch.ones_like(batch.latents)) + _, encoded, normalized_timestep = transformer.calls[0] + assert encoded[0].shape == (2, 8) + assert_close(normalized_timestep, torch.tensor([0.0])) + + ZImageDecodingStage(vae).forward(batch, args) + assert vae.decode_input is not None + expected_vae_input = torch.full((1, 16, 8, 8), 1 / 0.3611 + 0.1159) + assert_close(vae.decode_input, expected_vae_input) + assert batch.output is not None + assert batch.output.shape == (1, 3, 1, 64, 64) + assert_close(batch.output, torch.full_like(batch.output, 0.5)) + + +def test_zimage_batch_rng_matches_one_native_generator(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device", + lambda: torch.device("cpu"), + ) + args = _stage_args(output_type="latent") + transformer = _ConstantTransformer() + batch = ForwardBatch( + data_type="image", + prompt=["one", "two"], + negative_prompt="", + height=64, + width=64, + num_frames=1, + num_inference_steps=1, + guidance_scale=0.0, + seed=42, + ) + ZImageInputValidationStage().forward(batch, args) + batch.extra["zimage_prompt_embeds"] = [torch.zeros(1, 8), torch.zeros(1, 8)] + ZImageLatentPreparationStage(transformer).forward(batch, args) + + expected = torch.randn( + (2, 16, 1, 8, 8), + generator=torch.Generator("cpu").manual_seed(42), + ) + assert_close(batch.latents, expected) + + +def test_zimage_text_stage_honors_request_max_length(monkeypatch: pytest.MonkeyPatch) -> None: + from fastvideo.pipelines.stages.text_encoding import TextEncodingStage + + args = _stage_args() + stage = TextEncodingStage(text_encoders=[object()], tokenizers=[object()]) + observed: list[int | None] = [] + + def fake_encode_text( + text, + fastvideo_args, + encoder_index, + return_attention_mask, + max_length=None, + ): + del text, fastvideo_args, encoder_index, return_attention_mask + observed.append(max_length) + return [torch.zeros(1, 4, 8)], [torch.ones(1, 4, dtype=torch.long)] + + monkeypatch.setattr(stage, "encode_text", fake_encode_text) + batch = ForwardBatch( + data_type="image", + prompt="test", + negative_prompt="", + max_sequence_length=64, + guidance_scale=0.0, + ) + stage.forward(batch, args) + assert observed == [64] + + +def test_zimage_cfg_none_uses_native_empty_negative_prompt(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device", + lambda: torch.device("cpu"), + ) + batch = ForwardBatch( + data_type="image", + prompt="test", + negative_prompt=None, + height=64, + width=64, + num_frames=1, + num_inference_steps=1, + guidance_scale=2.0, + seed=42, + ) + ZImageInputValidationStage().forward(batch, _stage_args()) + assert batch.negative_prompt == "" + + +def _require_pinned_reference() -> Path: + if not REFERENCE_REPO.is_dir(): + pytest.skip(f"Pinned Tongyi Z-Image clone not found: {REFERENCE_REPO}") + result = subprocess.run( + ["git", "-C", str(REFERENCE_REPO), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + actual = result.stdout.strip() + assert actual == REFERENCE_REVISION, ( + f"Tongyi Z-Image oracle must be pinned at {REFERENCE_REVISION}, got {actual}" + ) + return REFERENCE_REPO / "src" + + +def _import_reference_modules() -> tuple[object, object]: + source_root = _require_pinned_reference().resolve() + if str(source_root) not in sys.path: + sys.path.insert(0, str(source_root)) + pipeline_module = importlib.import_module("zimage.pipeline") + utils_module = importlib.import_module("utils") + for module in (pipeline_module, utils_module): + module_path = Path(module.__file__ or "").resolve() + assert module_path.is_relative_to(source_root), ( + f"Z-Image oracle import escaped the pinned repository: {module_path}" + ) + return pipeline_module, utils_module + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Full Z-Image pipeline parity requires CUDA") +def test_zimage_pipeline_latents_match_pinned_native_repo(tmp_path: Path) -> None: + """Compare one complete native step; the oracle never imports Diffusers' pipeline.""" + if not (MODEL_DIR / "model_index.json").is_file(): + pytest.skip(f"Pinned Z-Image weights not found: {MODEL_DIR}") + pipeline_module, utils_module = _import_reference_modules() + prompt = "Young Chinese woman in red Hanfu, intricate embroidery." + + previous_fastvideo_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") + try: + # Exercise the shared PyTorch SDPA path with each process's native + # kernel selection; the official and FastVideo implementations both + # pass the same padding mask into scaled_dot_product_attention. + utils_module.set_attention_backend(None) + components = utils_module.load_from_local_dir( + MODEL_DIR, + device="cuda", + dtype=torch.bfloat16, + compile=False, + ) + reference = pipeline_module.generate( + prompt=prompt, + **components, + height=64, + width=64, + num_inference_steps=1, + guidance_scale=0.0, + negative_prompt="", + max_sequence_length=512, + generator=torch.Generator("cuda").manual_seed(42), + output_type="latent", + ).detach().float().cpu() + del components + gc.collect() + torch.cuda.empty_cache() + + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_pretrained( + str(MODEL_DIR), + num_gpus=1, + tp_size=1, + sp_size=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=False, + vae_cpu_offload=False, + pin_cpu_memory=False, + output_type="latent", + ) + try: + result = generator.generate_video( + prompt, + negative_prompt="", + output_path=str(tmp_path), + save_video=False, + return_frames=True, + height=64, + width=64, + num_frames=1, + fps=1, + num_inference_steps=1, + guidance_scale=0.0, + max_sequence_length=512, + seed=42, + ) + finally: + generator.shutdown() + actual = result["samples"].squeeze(2).float() + assert actual.shape == reference.shape + diff = (actual - reference).abs() + print(f"Z-Image pipeline latent parity: max={diff.max():.6f}, mean={diff.mean():.6f}") + assert_close(actual, reference, atol=1e-2, rtol=1e-2) + finally: + if previous_fastvideo_backend is None: + os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) + else: + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous_fastvideo_backend diff --git a/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py b/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py new file mode 100644 index 000000000..4b7dc9ea3 --- /dev/null +++ b/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Import/registry preflight and optional real Z-Image-Turbo smoke.""" + +from __future__ import annotations + +import os +from contextlib import nullcontext +from pathlib import Path +from typing import Any, cast + +import pytest +import torch + + +MODEL_DIR_ENV = os.getenv("ZIMAGE_MODEL_DIR") + + +def test_zimage_typed_surface_preflight() -> None: + from fastvideo.api import GeneratorConfig, PipelineSelection + from fastvideo.api.compat import generator_config_to_fastvideo_args + import fastvideo.registry as registry + from fastvideo.api.presets import get_preset, get_presets_for_family + from fastvideo.configs.models.dits.zimage import ZImageDiTConfig + from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig + from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig + from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig + from fastvideo.fastvideo_args import WorkloadType + from fastvideo.pipelines.basic.zimage.zimage_pipeline import EntryClass, ZImagePipeline + + assert ZImagePipeline.__name__ == "ZImagePipeline" + assert EntryClass is ZImagePipeline + assert ZImagePipeline.pipeline_config_cls is ZImagePipelineConfig + assert ZImagePipeline._required_config_modules == [ + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ] + + default_preset, model_family = registry.get_preset_selection("Tongyi-MAI/Z-Image-Turbo") + assert (default_preset, model_family) == ("zimage_turbo", "zimage") + info = registry.get_model_info( + "Tongyi-MAI/Z-Image-Turbo", + workload_type=WorkloadType.T2I, + override_pipeline_cls_name="ZImagePipeline", + ) + assert info.pipeline_cls is ZImagePipeline + assert info.pipeline_config_cls is ZImagePipelineConfig + + assert {preset.name for preset in get_presets_for_family("zimage")} == {"zimage_turbo"} + preset = get_preset("zimage_turbo", "zimage") + assert preset.defaults == { + "height": 1024, + "width": 1024, + "num_frames": 1, + "fps": 1, + "seed": 42, + "guidance_scale": 0.0, + "num_inference_steps": 8, + "negative_prompt": "", + "max_sequence_length": 512, + "cfg_normalization": False, + "cfg_truncation": 1.0, + } + + config = ZImagePipelineConfig() + assert isinstance(config.dit_config, ZImageDiTConfig) + assert isinstance(config.vae_config, AutoencoderKLVAEConfig) + assert isinstance(config.text_encoder_configs[0], Qwen3TextConfig) + assert config.text_encoder_configs[0].chat_template_enable_thinking is True + assert config.text_encoder_configs[0].output_hidden_states is True + assert config.text_encoder_configs[0].tokenizer_kwargs["max_length"] == 512 + assert config.scheduler_sigma_min == 0.0 + assert config.scheduler_use_reference_discrete_timesteps is True + + args = generator_config_to_fastvideo_args( + GeneratorConfig( + model_path="Tongyi-MAI/Z-Image-Turbo", + pipeline=PipelineSelection(workload_type="t2i"), + )) + assert isinstance(args.pipeline_config, ZImagePipelineConfig) + assert args.workload_type is WorkloadType.T2I + + +def test_model_download_honors_pinned_revision(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + import fastvideo.utils as utils + + observed: dict[str, object] = {} + + def fake_snapshot_download(**kwargs): + observed.update(kwargs) + return str(tmp_path) + + monkeypatch.setattr(utils, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(utils, "get_lock", lambda _key: nullcontext()) + resolved = utils.maybe_download_model( + "Tongyi-MAI/Z-Image-Turbo", + revision="f332072aa78be7aecdf3ee76d5c247082da564a6", + ) + + assert resolved == str(tmp_path) + assert observed["repo_id"] == "Tongyi-MAI/Z-Image-Turbo" + assert observed["revision"] == "f332072aa78be7aecdf3ee76d5c247082da564a6" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Z-Image load/generate smoke requires CUDA") +def test_zimage_pipeline_load_generate_smoke() -> None: + if MODEL_DIR_ENV is None: + pytest.skip("Set ZIMAGE_MODEL_DIR to activate the real load/generate smoke") + model_dir = Path(MODEL_DIR_ENV) + if not (model_dir / "model_index.json").is_file(): + pytest.skip(f"Z-Image model_index.json not found under {model_dir}") + + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_pretrained( + str(model_dir), + num_gpus=1, + tp_size=1, + sp_size=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=False, + vae_cpu_offload=False, + pin_cpu_memory=False, + output_type="latent", + ) + try: + result = generator.generate_video( + prompt="a red panda reading a book", + negative_prompt="", + output_path="outputs/zimage/smoke", + save_video=False, + return_frames=True, + height=64, + width=64, + num_frames=1, + fps=1, + num_inference_steps=1, + guidance_scale=0.0, + max_sequence_length=64, + cfg_normalization=False, + cfg_truncation=1.0, + seed=42, + ) + finally: + generator.shutdown() + + result_dict = cast(dict[str, Any], result) + samples = result_dict["samples"] + assert torch.is_tensor(samples) + assert samples.ndim in (4, 5) + assert torch.isfinite(samples).all() diff --git a/tests/local_tests/zimage/PORT_STATUS.md b/tests/local_tests/zimage/PORT_STATUS.md new file mode 100644 index 000000000..b66f631ac --- /dev/null +++ b/tests/local_tests/zimage/PORT_STATUS.md @@ -0,0 +1,162 @@ +# Z-Image Port Status + +## Summary + +- model_family: `zimage` +- variant: `Z-Image-Turbo` +- workload_types: `T2I` +- PR scope: required components plus same-PR FastVideo-native pipeline, example/parity, SSIM test, and L40S seed workflow +- official_ref: `https://github.com/Tongyi-MAI/Z-Image` +- official_ref_dir: `/Z-Image/src` +- official_ref_commit: `26f23eda626ffadda020b04ff79488e1d72004cd` +- hf_weights_path: `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` +- local_weights_dir: `/official_weights/Z-Image/` +- source_layout: official repo-native `src/zimage/`; HF component subfolders provide assets only +- weight_layout: HF per-component subfolders; these are assets, not the reference implementation +- local_tests_readme: `tests/local_tests/zimage/README.md` + +The pinned text encoder has 36 decoder blocks, hidden size 2560, +intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension +128. FastVideo obtains those values from the pinned `config.json`. + +## Current Phase + +- phase: Phase 11 (post-parity review and handoff); quality reference published +- status: external_ci_gate +- owner: infra/reviewer +- last_updated: 2026-07-15 + +## Pinned Official Pipeline Context + +- official_pipeline: `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`, + `src/zimage/pipeline.py::generate` and `src/config/inference.py` +- oracle_policy: official Z-Image source only; Diffusers is not the implementation oracle +- workload_and_output: `T2I`; prompt plus optional negative prompt to PIL images, or latents when requested +- model_index_contract: `_class_name=ZImagePipeline`; required modules are `transformer`, `vae`, + `text_encoder`, `tokenizer`, and `scheduler` +- official_defaults: `height=1024`, `width=1024`, `num_inference_steps=8`, + `guidance_scale=0.0`, `cfg_truncation=1.0`, `max_sequence_length=512` +- reproducible_rng: `torch.Generator("cuda").manual_seed(42)`; the API accepts a generator and has no fixed seed default +- scheduler_contract: `use_reference_discrete_timesteps=True`, runtime `sigma_min=0.0`, endpoint-preserving + `num_steps + 1` schedule, and no DiT call for the final zero timestep +- decode_contract: apply `(latents / 0.3611) + 0.1159` before official VAE decode +- implementation_status: native pipeline/config/preset/registry/example and the 8-step 1024x1024 PNG SSIM test are + implemented; one-step native GPU parity is bit-exact; the L40S candidate generated at inference head `9a6dbcadc` + was author-approved and published at dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c` + +## Component Matrix + +| Component | Type | Reuse/Port | Official Definition / Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +|---|---|---|---|---|---|---|---|---| +| Production text encoder (Qwen3 body) | text_encoder | reused production passthrough | Official `src/utils/loader.py` uses Transformers with typed load then device-only move; official pipeline consumes `hidden_states[-2]` | `TextEncoderLoader` -> `Qwen3ForCausalLM.from_pretrained_local` -> body-only Transformers `AutoModel` | DONE | not_needed for current HF subfolder | PINNED GB200 PASS; reference/production outputs exact | none | +| Native Qwen quality | implementation_subcomponent | shared native implementation | Same pinned text-encoder weights and official `hidden_states[-2]` boundary; not loaded by the pipeline | direct `fastvideo.models.encoders.qwen3.Qwen3ForCausalLM` construction | DONE | not_needed | PINNED GB200 PASS; fp32 `atol=2e-4`, `rtol=1e-4`; bf16 thresholds unchanged | none | +| Tokenizer | tokenizer | reused | Official loader's `AutoTokenizer` from `/tokenizer/`; official thinking template and length 512 | `TokenizerLoader` | DONE | not_needed | PINNED GB200 PASS (2 passed) | none | +| VAE | vae | reused with accepted wrapper exception | Pinned official `src/zimage/autoencoder.py`; raw decode is the component boundary | Existing Diffusers-backed `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` through production `VAELoader`, plus direct raw-decode coverage | DONE | not_needed | PINNED GB200 PASS (2 passed) | none | +| Scheduler | scheduler | reused with Z-Image extension | Pinned official `src/zimage/scheduler.py`; official pipeline sets `sigma_min=0.0` | `FlowMatchEulerDiscreteScheduler(use_reference_discrete_timesteps=True, sigma_min=0.0)` | DONE | not_needed | PINNED GB200 PASS (5 passed) | none | +| Transformer (`ZImageTransformer2DModel`) | dit | native port | Pinned official `src/zimage/transformer.py`; HF `transformer/` supplies weights only | `fastvideo.models.dits.zimage.ZImageTransformer2DModel` through `TransformerLoader` | DONE | not_needed; exact 521-key surface | PINNED GB200 PASS (5 passed; real-weight full forward exact) | none | +| Pipeline | pipeline | native same-PR implementation | Pinned official `src/zimage/pipeline.py::generate` | FastVideo `ZImagePipeline` plus config/preset/registry/example/parity | DONE | not_needed | RTX 5090 ONE-STEP LATENTS BIT-EXACT; L40S REFERENCE PUBLISHED | I005 (CI verification only) | + +## Conversion State + +- conversion_script: none for the reused components or native transformer +- converted_weights_dir: none +- source_layout: official repo-native Python (`Tongyi-MAI/Z-Image`) +- hf_weight_layout: per-component subfolders; not a Diffusers implementation oracle +- 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 passed; native Qwen strictness regressions passed +- passthrough_components: 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 +- native_qwen_fused_contract: every Q/K/V and gate/up shard is required unless the exact fused destination is present +- native_qwen_auxiliary_contract: mapped quantization auxiliaries, including scale parameters, must load and may not be silently discarded +- retry_history: none + +## Parity Commands + +| Scope | Command | Last Result | Notes | +|---|---|---|---| +| Scheduler | `pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v -s` | 5 PASSED, no skips, GB200, 2026-07-15 | Enforces official-source SHA/import origin, fp32 scalar contract, exact zero endpoint, and positional/default regressions. | +| Tokenizer | `pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -s` | 2 PASSED as part of the pinned full GB200 run, 2026-07-15 | Uses `enable_thinking=True`, `max_length=512`, and fails if chat templating is unavailable. | +| VAE | `pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v -s` | 2 PASSED, no skips, GB200, 2026-07-15 | Direct raw decode plus production `VAELoader`; config values are checked, while latent scale/shift application is pipeline behavior. | +| Production text encoder | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | REFERENCE/PRODUCTION CHECKS EXACT, GB200, 2026-07-15 | The required pipeline path matches an independent official-loader materialization exactly. | +| Native Qwen fp32 quality | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | PASS, GB200, 2026-07-15 | Z-Image-only `atol=2e-4`, `rtol=1e-4`; production equality and shared Qwen math are unchanged. | +| Native Qwen bf16 quality | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16] -v -s` | PASS, GB200, 2026-07-15 | Existing mean/median distribution thresholds pass unchanged against an independent materialization of the official loader contract. | +| Native Qwen per-layer bf16 diagnostic | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | PASS, GB200, 2026-07-15 | Pinned architecture has 36 blocks and 37 hidden-state entries under the tested contract. | +| Transformer parity | `ZIMAGE_OFFICIAL_REF_DIR= pytest tests/local_tests/zimage/test_zimage_transformer_parity.py -v -s` | 5 PASSED, no skips, GB200, 2026-07-15 | Production 521-key/shape surface, tiny CPU parity, mask behavior, SP guard, real-weight strict load, and 6.15B full forward pass; real-weight output was exact (`max=0`, `mean=0`). | +| Complete component closure | `ZIMAGE_OFFICIAL_REF_DIR= ZIMAGE_TRANSFORMER_DIR= pytest tests/local_tests/zimage/ -v -s` | 34 PASSED, no skips, GB200, 2026-07-15, `f310aacb1` | Official source `26f23eda...`, HF snapshot `f332072a...`; production/native encoder, scheduler, tokenizer, transformer, and VAE all passed. | +| Pipeline typed surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS, hardware-free, 2026-07-14 | Covers exact `EntryClass`, config, preset, registry, model detection, and official defaults. | +| Native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS, hardware-free, 2026-07-14 | Covers conditioning trim, endpoint-preserving timestep setup, denoising sign/timestep math, and VAE scale/shift. | +| Full pipeline parity | `ZIMAGE_REFERENCE_REPO= ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS, RTX 5090, `fec2baa6c` | Real FastVideo one-step latents match pinned native Tongyi exactly (`max=0`, `mean=0`). Later commits through `f310aacb1` change only API inventory/tests, not inference math. | +| PNG SSIM quality | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | L40S REFERENCE UPLOADED; END-TO-END SSIM CI PENDING | Dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`; remote SHA-256 `6c2be648...f7a6`; upload scoped to `Tongyi-MAI__Z-Image-Turbo/TORCH_SDPA`. | + +## Open Questions + +| ID | Question | Owner | Needed By Phase | Status | Resolution | +|---|---|---|---|---|---| +| 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 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 + +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +|---|---|---|---|---|---|---|---|---| +| I001 | Phase 6 | native_qwen_quality | medium | The optional native fp32 `hidden_states[-2]` tail narrowly exceeded `atol=1e-4` after smooth cross-kernel accumulation. | Exact layer trace: embeddings exact, layer-0 max `3.815e-6`, first misses at layer 33, largest failing boundary difference `1.686811e-4`; final RMSNorm passes. The required production `AutoModel` path is exact. | parity | resolved | Author approved Z-Image-only `_FP32_ATOL=2e-4` with `_FP32_RTOL=1e-4`; complete pinned GB200 suite passed 34/34 with bf16 and all other gates unchanged. | +| I002 | Phase 7 | scheduler/pipeline | high | Published `scheduler_config.json` omits the reference-schedule flag and official zero endpoint. | The native stage applies both values; hardware-free stage coverage and pinned native GPU parity pass. | pipeline | resolved | The Z-Image pipeline applies `use_reference_discrete_timesteps=True` and `sigma_min=0.0`; one-step latents match official Tongyi exactly. | +| 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 | resolved | Corrected behavior passed the pinned GB200 run. | +| 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 10 | pipeline/quality | high | Verify the published L40S PNG through end-to-end SSIM CI. | Author approved the reviewed candidate; the guarded upload published exactly one file at dataset revision `7b27d000...f01c`, and remote bytes match SHA-256 `6c2be648...f7a6`. | infra | external_ci_gate | Reference publication is resolved. The full-suite launcher remains blocked by the external-fork documentation workflow's `action_required` state; merge is separate. | +| I006 | Phase 6 | text_encoder | high | HF passthrough loading inherited an ambient target-device context, and the earlier test harness incorrectly cast float32 RoPE buffers during bf16 placement. | HF loading now occurs outside the ambient context; typed load is followed by device-only movement, matching official `src/utils/loader.py`; production/independent outputs match exactly. | encoder | resolved | Hardware-free placement/default-dtype tests and the pinned GB200 production path pass. | +| I007 | Phase 6 | text_encoder | high | Lenient native loading could miss one fused source shard or an auxiliary quantization scale. | Loader completeness covers every destination, required split shard, exact fused destination, and mapped scale parameter. | encoder | resolved | Strictness regressions and pinned checkpoint loading pass. | +| 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 now covers direct raw decode and production `VAELoader` against official `src/zimage/autoencoder.py`. | parity | resolved | Pinned GB200 subset passed 2/2. The wrapper is an implementation exception, not the oracle; official latent normalization remains Phase 7 pipeline scope. | +| 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; pinned asset-backed parity passed. | +| I011 | Phase 6 | transformer | high | Production real-weight strict load and full-forward parity were not verified. | Pinned 24.6 GB weights and official source were exercised on GB200. | parity | resolved | Strict load and full forward passed exactly (`max=0`, `mean=0`). | +| I012 | Phase 6 | required_components | high | Record a complete non-skip run after the loader/oracle corrections and approved native-Qwen tolerance. | Exact PR head `f310aacb1`, pinned official source `26f23eda...`, immutable HF snapshot `f332072a...`, GB200. | parity | resolved | Complete suite passed 34/34 in 94.21s with no skips; real-weight transformer output was exact (`max=0`, `mean=0`). | + +## Escape Hatches + +| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | +|---|---|---|---|---|---|---| +| E001 | Phase 6 | cost | Approve one standard 8-hour 4×GB200 hold to run the complete pinned component and real-weight transformer parity suite at PR #1339? | Approve `model/r9`; reuse a suitable existing allocation if available. | resolved | Author approved. Existing job `1545932` was reused, so no new hold was created. Excludes Modal, SSIM seeding/upload, and merge. | +| E002 | Phase 6 | numerical | Accept the traced Z-Image-only native encoder fp32 tail without weakening shared Qwen or bf16 gates? | Approve `model/r10`: change only `_FP32_ATOL` from `1e-4` to `2e-4`, retain `_FP32_RTOL=1e-4`, then rerun native parity at the exact resulting head. | resolved | Author approved. The one-line test threshold changed at `f310aacb1`; complete pinned GB200 suite passed 34/34 with production equality, bf16, strict-load, and official-oracle checks unchanged. | + +## Decisions + +| Date | Decision | Rationale | Impact | +|---|---|---|---| +| 2026-07-15 | Publish the visually approved L40S PNG reference to `FastVideo/ssim-reference-videos`. | The author confirmed the staged image looks good and explicitly requested the HF upload. | Guarded no-force upload added one Z-Image PNG at dataset revision `7b27d000...f01c`; remote bytes match the reviewed artifact. | +| 2026-07-15 | Accept the traced Z-Image-only native encoder fp32 absolute tolerance of `2e-4` while retaining `rtol=1e-4`. | Exact tracing showed gradual fused/split-kernel accumulation rather than a semantic break; the largest failing absolute tail was `1.686811e-4`. | One test constant changed. The full pinned suite passed 34/34; production equality, bf16 thresholds, strict loading, and official Tongyi oracle are unchanged. | +| 2026-07-14 | Carry the native FastVideo pipeline, example/parity, SSIM test, and L40S seed workflow on PR #1339. | The author explicitly requested the pipeline and quality regression on the same PR rather than a follow-up. | Implementation, component closure, native GPU parity, and reference publication are complete. Later commits do not change inference math. | +| 2026-07-14 | Treat the exact production `AutoModel` loader as the required pipeline text encoder and direct native-Qwen parity as a separate quality row. | `TextEncoderLoader` calls `Qwen3ForCausalLM.from_pretrained_local`, which returns body-only Transformers `AutoModel`; the pipeline does not direct-construct FastVideo-native Qwen. | The production component is exact. The separately approved native-Qwen tolerance changed no production math. | +| 2026-07-13 | Use pinned `Tongyi-MAI/Z-Image@26f23eda...` as the implementation oracle; treat HF subfolders only as weight/config assets. | The official repo defines transformer, VAE, scheduler, pipeline, and loader semantics. A runtime component may inherit Diffusers without making Diffusers the reference. | Tests verify official clone SHA/import origin; docs distinguish the oracle from storage/runtime implementation details. | +| 2026-07-13 | Keep official latent scale/shift application out of the VAE component assertion. | `(latents / 0.3611) + 0.1159` is pipeline-stage behavior, while VAE component parity compares raw decode and production loading. | The Z-Image stage implements the exact formula; pinned native GPU parity passed. | +| 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 were revalidated. | +| 2026-07-04 | Use a body-only Transformers `AutoModel` for production and a separate `AutoModel` object to materialize the official loader contract independently. | The pinned official Z-Image source is the oracle. Every FastVideo consumer needs hidden states, not full-vocabulary logits or an LM head, and 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. | +| 2026-07-04 | Reuse the existing Diffusers-backed AutoencoderKL wrapper for the Z-Image VAE component. | The target is established shared code with production-loader parity against pinned official `src/zimage/autoencoder.py`. | The wrapper remains an implementation detail for this VAE only; official Z-Image source defines pipeline behavior and numerical parity. | +| 2026-07-04 | Treat absent chat-template behavior as incompatibility, not an optional test environment. | Thinking-format prompt construction is part of the official numerical contract. | Tokenizer parity fails when `apply_chat_template` is unavailable. | +| 2026-07-04 | Record the pinned Qwen architecture as 36 blocks, hidden size 2560, intermediate size 9728, 32 attention heads, and 8 KV heads. | Earlier notes described a different architecture and could mis-size the native target. | Config/review evidence now matches `f332072aa78be7aecdf3ee76d5c247082da564a6`; expected hidden-state length is 37. | +| 2026-06-21 | Expand generated Qwen `position_ids` to `[batch_size, seq_len]`. | The rotary layer flattens positions; `[1, seq_len]` misaligned RoPE for batch sizes greater than one. | Batch-one behavior is unchanged; pinned batch-two parity passed. | +| 2026-06-21 | Reuse the shared native `Qwen3ForCausalLM` implementation rather than maintain a second Z-Image encoder. | The shared implementation supports the same config-driven GQA, fused QKV/gate-up, and hidden-state output contract. | Registry/native parity uses the shared class; architecture values come from the pinned config. | +| 2026-05-12 | Use mean and median drift checks for deep bf16 encoder parity. | Cross-kernel fused/unfused bf16 arithmetic produces a long max-error tail; fp32 and per-layer diagnostics distinguish this from structural divergence. | Existing thresholds passed unchanged on the pinned 36-block snapshot. | +| 2026-05-12 | Forward the complete scheduler config except loader metadata. | Hand-picking three fields could silently omit future scheduler behavior. | Scheduler parity consumes all on-disk constructor fields and explicitly supplies Z-Image runtime overrides. | + +## Handoff Notes + +- The implementation oracle is the pinned official `Tongyi-MAI/Z-Image` clone, + never Diffusers. Transformers is used for text components only because the + official loader does so; HF component subfolders provide immutable assets. +- The required production text encoder is the exact body-only `AutoModel` path. + `model/r10` is resolved: only the Z-Image test's fp32 absolute tolerance is + `2e-4`; relative tolerance, shared Qwen math, bf16, and production loading are unchanged. +- I012 is closed by the exact-head 34/34 non-skip GB200 run at `f310aacb1`. +- Native pipeline/config/preset/registry/example, parity tests, and the 8-step + 1024x1024 PNG SSIM test are implemented. Native RTX 5090 one-step parity is + bit-exact. The reviewed L40S PNG is live at dataset revision + `7b27d00004b723ef0cf4923e77ae57128573f01c` with matching remote bytes. +- End-to-end SSIM/full-suite execution remains an infra/reviewer gate; reference + publication does not authorize PR approval or merge. diff --git a/tests/local_tests/zimage/README.md b/tests/local_tests/zimage/README.md new file mode 100644 index 000000000..223b2b0f9 --- /dev/null +++ b/tests/local_tests/zimage/README.md @@ -0,0 +1,200 @@ +# Z-Image Local Tests + +Local component coverage and quality handoff for the `zimage` FastVideo port +(`T2I`, Z-Image-Turbo only). The same PR is authorized to carry the components, +FastVideo-native `ZImagePipeline`, example, pipeline parity, SSIM test, and L40S +reference seed. The native pipeline surface and 8-step 1024x1024 PNG SSIM test +are implemented. Native one-step GPU parity is bit-exact, and the L40S candidate +has been published to `FastVideo/ssim-reference-videos`. + +> **Status:** Implementation and quality-reference seeding complete; Phase 11 +> final handoff. The +> Phase 6 exact-head closure is complete: the full pinned component suite passed +> 34/34 with no skips on GB200 at `f310aacb1`. The required production text +> encoder is the official-loader-compatible Transformers `AutoModel` path, which +> matches an independently materialized reference exactly. Scheduler, tokenizer, +> VAE, native Qwen fp32/bf16, and the 6.15B real-weight transformer all pass. +> Native one-step pipeline parity is bit-exact on RTX 5090. The reviewed +> 1024x1024 L40S PNG generated at inference head `9a6dbcadc` is published at +> dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`. See +> [`PORT_STATUS.md`](./PORT_STATUS.md) for the live state. + +## Reference assets and scope + +| Field | Value | +|---|---| +| Model family / variant | `zimage` / `Z-Image-Turbo` | +| Workload | `T2I` | +| PR scope | required components, native pipeline, example/parity, SSIM test, and L40S seed | +| Required component scope | production `AutoModel` text encoder, tokenizer, VAE, scheduler, native transformer | +| Additional quality scope | direct FastVideo-native Qwen parity (`model/r10`; not used by the pipeline loader) | +| 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/` (component subfolders; transformer is 24.6 GB) | +| Reference implementation layout | Official repo-native `src/zimage/` Python; Diffusers is not the parity oracle | +| Weight layout | HF per-component subfolders (`text_encoder/`, `vae/`, `scheduler/`, `transformer/`) | +| 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 +blocks, hidden size 2560, intermediate size 9728, 32 attention heads, 8 KV +heads, and head dimension 128. FastVideo loads those values from +`text_encoder/config.json`; they are not hard-coded in the Z-Image integration. + +## Shared environment setup + +Run from the FastVideo repository root in the same environment used for +FastVideo. Keep the reference clone and weights inside the ignored paths shown +below. + +Clone and pin the official reference: + +```bash +python .agents/skills/add-model-01-prep/scripts/clone_reference_repo.py \ + https://github.com/Tongyi-MAI/Z-Image.git \ + Z-Image \ + --commit 26f23eda626ffadda020b04ff79488e1d72004cd \ + --update-gitignore +git -C Z-Image rev-parse HEAD +``` + +The final command must print +`26f23eda626ffadda020b04ff79488e1d72004cd`. The scheduler and VAE tests also +enforce that exact HEAD and verify that imported `zimage` modules resolve under +`/Z-Image/src`. A missing clone may skip local parity; a wrong SHA or +wrong import origin fails rather than silently testing another installation. + +Download the immutable HF snapshot: + +```bash +python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py \ + Tongyi-MAI/Z-Image-Turbo \ + official_weights/Z-Image \ + --revision f332072aa78be7aecdf3ee76d5c247082da564a6 \ + --allow-pattern 'text_encoder/*' \ + --allow-pattern 'tokenizer/*' \ + --allow-pattern 'vae/*' \ + --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: official source pinned and imported; GB200 component validation run completed +private_dep_stubs: none +blocked_on: repo-wide full-suite/SSIM launch gate and PR review; model implementation work is complete +``` + +## Run the tests + +```bash +pytest tests/local_tests/zimage/ -v -s +``` + +| Component | Coverage scope and contract | Status | +|---|---|---| +| Scheduler ([`test_zimage_scheduler_parity.py`](./test_zimage_scheduler_parity.py)) | `implementation_subcomponent`; exact official `sigma_min=0.0` plus `use_reference_discrete_timesteps=True`; positional/default-path regressions remain asset-free | PINNED GB200 PASS (5 passed) | +| 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 | PINNED GB200 PASS (2 passed) | +| VAE ([`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py)) | `both`; direct raw decode plus production `VAELoader`, config values, and raw decode | PINNED GB200 PASS (2 passed) | +| Production text encoder ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `production_loader`; `TextEncoderLoader` materializes a body-only Transformers `AutoModel` matching official `src/utils/loader.py`; required by the pipeline | PINNED GB200 PASS (reference/production outputs exact) | +| Native Qwen quality ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `implementation_subcomponent`; direct FastVideo-native construction, strict loading, and fp32/bf16 diagnostics; not loaded by the pipeline | PINNED GB200 PASS; fp32 `atol=2e-4`, `rtol=1e-4`; bf16 thresholds unchanged | +| 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 parity | PINNED GB200 PASS (5 passed; real-weight full forward exact) | + +## Pinned official pipeline context + +| Field | Contract | +|---|---| +| Oracle | `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`, `src/zimage/pipeline.py::generate` and `src/config/inference.py`; never Diffusers | +| Pipeline/modules | `_class_name=ZImagePipeline`; `transformer`, `vae`, `text_encoder`, `tokenizer`, `scheduler` | +| Inputs/outputs | text-to-image prompt with optional negative prompt; PIL images by default or latents when requested | +| Official defaults | `1024x1024`, `num_inference_steps=8`, `guidance_scale=0.0`, `cfg_truncation=1.0`, `max_sequence_length=512` | +| Reproducible RNG | `torch.Generator("cuda").manual_seed(42)`; `generate()` accepts the generator rather than fixing a seed internally | +| Scheduler | `use_reference_discrete_timesteps=True`, runtime `sigma_min=0.0`, and the endpoint-preserving `num_steps + 1` schedule with the final zero-timestep DiT call skipped | +| Decode | apply `(latents / 0.3611) + 0.1159` before official VAE decode | +| FastVideo status | Native pipeline/config/preset/registry/example and PNG SSIM test are implemented. One-step RTX 5090 parity is bit-exact; the reviewed L40S PNG generated at `9a6dbcadc` is published at dataset revision `7b27d000...f01c`. Later PR changes affect only tests/docs. | + +## Pipeline and quality verification + +| Scope | Command | Status | +|---|---|---| +| Hardware-free pipeline surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS | +| Hardware-free native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS | +| Pinned native full GPU pipeline parity | `ZIMAGE_REFERENCE_REPO= ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS on RTX 5090; one-step latents bit-exact (`max=0`, `mean=0`) | +| 8-step 1024x1024 PNG SSIM | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | Reviewed L40S PNG uploaded; remote SHA-256 `6c2be648...f7a6`; end-to-end SSIM CI pending | + +## Component contracts + +### Text encoder and tokenizer + +- The production `TextEncoderLoader` path deliberately loads Transformers + `AutoModel`, even when the checkpoint advertises `Qwen3ForCausalLM`. FastVideo + consumes hidden states only, so the production path is body-only and does not + materialize or execute an LM head. +- The parity oracle is the pinned official Z-Image source and pipeline contract. + A separate Transformers `AutoModel` instance materializes the official + loader behavior independently; it is not the object returned by FastVideo's + production loader. The native `Qwen3ForCausalLM` implementation is compared + separately. +- CPU/MPS text-encoder offload remains on the requested device. The loader records + `_fastvideo_input_device`, and `TextEncodingStage` sends token tensors there. +- Native loading must account for every destination parameter and every required + fused source shard. Q/K/V and gate/up shards are all required unless the exact + destination is supplied already fused; quantized auxiliary parameters such as + `scale_weight` are not silently ignored. +- The pinned official prompt path enables thinking, tokenizes to 512, and consumes + `hidden_states[-2]` at valid-token positions. A 36-block Qwen model exposes 37 + hidden-state entries in this contract (embedding/intermediate entries plus the + final normalized state). +- `TextEncoderConfig.chat_template_enable_thinking` is keyword-only and defaults + to `False` for existing model families. The Z-Image pipeline config opts in + with `True`. + +### VAE + +The test covers both direct implementation parity and production `VAELoader` +resolution/strict loading, confirms the official config's scaling/shift values, +and compares raw decode outputs. The official pipeline's +`(latents / 0.3611) + 0.1159` transformation is pipeline behavior, so it remains +a Phase 7 Z-Image pipeline-stage and pipeline-parity requirement. This port +reuses the +existing shared `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` wrapper, +which subclasses Diffusers `AutoencoderKL`, as a narrowly scoped exception to +the native-component boundary. That runtime inheritance does not make Diffusers +the reference implementation: parity is against pinned +`Tongyi-MAI/Z-Image/src/zimage/autoencoder.py`. This is not precedent for the +transformer or pipeline implementation. + +### Scheduler + +The pinned official pipeline mutates `scheduler.sigma_min = 0.0` before building +its `num_steps + 1` schedule. The Z-Image pipeline applies 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. The pinned GB200 run strict-loaded the + 6.15B model and matched the official full-forward output exactly. + +## Remaining work + +- Run the end-to-end SSIM/full suite after the external-fork documentation + approval gate permits the repo-wide launcher. +- Keep PR approval and merge separate from the completed reference upload. diff --git a/tests/local_tests/zimage/__init__.py b/tests/local_tests/zimage/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/tests/local_tests/zimage/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/zimage/test_zimage_encoder_parity.py b/tests/local_tests/zimage/test_zimage_encoder_parity.py new file mode 100644 index 000000000..284fe5b72 --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_encoder_parity.py @@ -0,0 +1,769 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Parity test for Z-Image Qwen3 text encoder support in FastVideo. + +This compares: +1) direct transformers AutoModel output, and +2) FastVideo's production ``TextEncoderLoader`` passthrough, and +3) FastVideo's shared native Qwen3 encoder (``Qwen3ForCausalLM``, added for Flux2 + Klein and reused here — Z-Image-Turbo's ``Qwen3Model`` checkpoint routes + to it via the model registry), + +using identical local checkpoint, tokenization, and inputs. + +Usage: + pytest tests/local_tests/zimage/test_zimage_encoder_parity.py -v -s +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +import gc +import json +import os + +import pytest +import torch + + +def _reclaim_vram() -> None: + """Actually reclaim VRAM after the caller has ``del``'d its model refs. + + HF transformer models hold reference cycles, so ``del`` alone does not free + them — a ``gc.collect()`` is required before the CUDA caching allocator will + release the blocks. Without this the parity test keeps two full encoder + copies resident and OOMs on a 44 GB GPU. + """ + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _gpu_mem_mb() -> float: + if torch.cuda.is_available(): + return torch.cuda.memory_allocated() / (1024 ** 2) + return 0.0 +from torch.testing import assert_close +from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer +from safetensors.torch import safe_open + +from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig +from fastvideo.distributed.parallel_state import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel +from fastvideo.layers.quantization.absmax_fp8 import AbsMaxFP8MergedParameter +from fastvideo.models.encoders.qwen3 import Qwen3ForCausalLM +from fastvideo.models.loader.component_loader import TextEncoderLoader +from fastvideo.models.loader.utils import set_default_torch_dtype +from fastvideo.models.registry import ModelRegistry + +PARITY_SCOPE = "both" + +# Strict-load contract: the shared Qwen3 encoder is body-only (embed_tokens + +# layers + norm, no lm_head). Z-Image-Turbo ships a full Qwen3 checkpoint, so +# `lm_head.weight` is the only key that may go unmatched; anything else means a +# real silent drop. Enforced here in the test since the shared encoder's loader +# is intentionally lenient (it's used by multiple models). +_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"} + + +REPO_ROOT = Path(__file__).resolve().parents[3] +ZIMAGE_TEXT_ENCODER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "text_encoder" +ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer" + + +@pytest.fixture(scope="module", autouse=True) +def _init_dist_and_tp_groups(): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29531") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + cleanup_dist_env_and_memory() + + +def _load_json(path: Path) -> dict: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def _iter_safetensors(path: Path): + with safe_open(str(path), framework="pt", device="cpu") as f: + for k in f.keys(): + yield k, f.get_tensor(k) + + +def _iter_pretrained_safetensors(model_dir: Path): + single = model_dir / "model.safetensors" + if single.exists(): + yield from _iter_safetensors(single) + return + + index = model_dir / "model.safetensors.index.json" + if index.exists(): + idx = _load_json(index) + shard_names = sorted(set(idx["weight_map"].values())) + for shard in shard_names: + yield from _iter_safetensors(model_dir / shard) + return + + raise FileNotFoundError( + f"Missing safetensors checkpoint in {model_dir} (expected model.safetensors or model.safetensors.index.json)" + ) + + +def _pretrained_safetensor_keys(model_dir: Path) -> set[str]: + return {name for name, _ in _iter_pretrained_safetensors(model_dir)} + + +def _load_qwen3_config() -> Qwen3TextConfig: + cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json") + for key in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"): + cfg_raw.pop(key, None) + config = Qwen3TextConfig() + config.update_model_arch(cfg_raw) + assert config.num_hidden_layers == 36 + assert config.hidden_size == 2560 + assert config.intermediate_size == 9728 + assert config.num_attention_heads == 32 + assert config.num_key_value_heads == 8 + return config + + +def _loader_args(cpu_offload: bool) -> SimpleNamespace: + return SimpleNamespace( + text_encoder_cpu_offload=cpu_offload, + override_text_encoder_quant=None, + override_text_encoder_safetensors=None, + pin_cpu_memory=False, + ) + + +def _precision_name(dtype: torch.dtype) -> str: + if dtype == torch.float32: + return "fp32" + if dtype == torch.bfloat16: + return "bf16" + raise ValueError(f"Unsupported parity dtype: {dtype}") + + +def _tokenize_official_prompts(tokenizer, prompts: list[str]): + formatted = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + for prompt in prompts + ] + return tokenizer( + formatted, + padding="max_length", + truncation=True, + max_length=512, + return_tensors="pt", + ) + + +def _assert_native_load_surface( + model, + checkpoint_keys: set[str], + loaded_params: set[str], +) -> None: + param_names = {name for name, _ in model.named_parameters()} + missing = param_names - loaded_params + assert not missing, f"FastVideo parameters not loaded: {sorted(missing)}" + + normalized_keys = { + name[len("model."):] if name.startswith("model.") else name + for name in checkpoint_keys + } + mapped_keys: set[str] = set() + for checkpoint_name in normalized_keys: + for param_name, weight_name, _ in model.config.arch_config.stacked_params_mapping: + if weight_name in checkpoint_name: + mapped_keys.add(checkpoint_name.replace(weight_name, param_name)) + break + else: + mapped_keys.add(checkpoint_name) + + unexpected = mapped_keys - param_names - { + "rotary_emb.inv_freq", + "rotary_emb.cos_cached", + "rotary_emb.sin_cached", + } + assert unexpected <= _ALLOWED_UNEXPECTED_KEYS, ( + "Unexpected checkpoint keys not in allowlist: " + f"{sorted(unexpected - _ALLOWED_UNEXPECTED_KEYS)}" + ) + + +def _fake_qwen3_weight_target(include_quant_scales: bool = False): + stacked = Qwen3TextConfig().arch_config.stacked_params_mapping + params = { + "layers.0.self_attn.qkv_proj.weight": torch.nn.Parameter(torch.empty(1)), + "layers.0.mlp.gate_up_proj.weight": torch.nn.Parameter(torch.empty(1)), + } + if include_quant_scales: + qkv_scale = AbsMaxFP8MergedParameter(torch.zeros(3), requires_grad=False) + qkv_scale.output_partition_sizes = [1, 1, 1] + gate_up_scale = AbsMaxFP8MergedParameter(torch.zeros(2), requires_grad=False) + gate_up_scale.output_partition_sizes = [1, 1] + params.update({ + "layers.0.self_attn.qkv_proj.scale_weight": qkv_scale, + "layers.0.mlp.gate_up_proj.scale_weight": gate_up_scale, + }) + for param in params.values(): + if not hasattr(param, "weight_loader"): + param.weight_loader = ( + lambda target, loaded, *args: target.data.copy_(loaded) + ) + + model = SimpleNamespace( + config=SimpleNamespace( + arch_config=SimpleNamespace(stacked_params_mapping=stacked) + ), + named_parameters=lambda: params.items(), + ) + return model, params, stacked + + +# The provisional bf16 thresholds below come from the historical unpinned +# 35-layer run. Per add-model-02-parity's calibration block, bf16 max-based +# asserts are meaningless for deep encoders — fused-vs-unfused linear order +# alone produces a long tail, so the test uses mean + median instead. The +# pinned snapshot has 36 blocks and must be rerun before these thresholds are +# treated as current evidence. +# +# Historical empirical numbers (unrecorded snapshot, 35 layers, CUDA): +# last_hidden_state (post-norm): mean=0.0152, median=0.0117 +# hidden_states[-2] (pre-norm): mean=0.0754, median=0.0625 +# Thresholds below add ~1.6x headroom over observed. +_BF16_MEAN_DRIFT_POST_NORM = 0.025 +_BF16_MEAN_DRIFT_PRE_NORM = 0.120 +_BF16_MEDIAN_DRIFT_POST_NORM = 0.020 +_BF16_MEDIAN_DRIFT_PRE_NORM = 0.100 +_FP32_ATOL = 2e-4 +_FP32_RTOL = 1e-4 + + +def _print_diag(label: str, ref: torch.Tensor, fv: torch.Tensor) -> torch.Tensor: + diff = (ref.float() - fv.float()).abs() + flat = diff.flatten() + p99 = flat.kthvalue(max(1, int(0.99 * flat.numel()))).values + print( + f"[{label}] max_diff={diff.max():.4f} mean_diff={diff.mean():.4f} " + f"median_diff={diff.median():.4f} p99_diff={p99:.4f}", + flush=True, + ) + return diff + + +@pytest.mark.parametrize( + ("cpu_offload", "target_device", "expected_device"), + [ + (True, torch.device("cuda"), torch.device("cpu")), + (False, torch.device("meta"), torch.device("meta")), + ], +) +def test_qwen3_production_loader_avoids_device_context_and_honors_placement( + monkeypatch: pytest.MonkeyPatch, + cpu_offload: bool, + target_device: torch.device, + expected_device: torch.device, +): + """HF loading stays ambient while final placement honors target/offload.""" + import fastvideo.platforms as platforms + + placements: list[torch.device] = [] + auto_model_calls: list[dict] = [] + + class FakeHFModel: + def eval(self): + return self + + def to(self, device: torch.device): + placements.append(torch.device(device)) + return self + + fake_model = FakeHFModel() + + ambient_device = torch.empty(0).device + default_devices: list[torch.device] = [] + + def fake_auto_model_from_pretrained(*args, **kwargs): + auto_model_calls.append(kwargs) + default_devices.append(torch.empty(0).device) + return fake_model + + def fail_causal_lm_from_pretrained(*args, **kwargs): + raise AssertionError("Qwen text encoding must not load AutoModelForCausalLM") + + monkeypatch.setattr( + AutoModel, + "from_pretrained", + staticmethod(fake_auto_model_from_pretrained), + ) + monkeypatch.setattr( + AutoModelForCausalLM, + "from_pretrained", + staticmethod(fail_causal_lm_from_pretrained), + ) + monkeypatch.setattr( + platforms, + "_current_platform", + SimpleNamespace( + is_mps=lambda: False, + verify_model_arch=lambda _arch: None, + ), + ) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + "fastvideo.distributed.get_local_torch_device", + lambda: torch.device("cuda"), + ) + + config = Qwen3TextConfig() + # The pinned checkpoint names the causal wrapper even though the official + # Z-Image loader intentionally asks Transformers for the body-only model. + config.arch_config.architectures = ["Qwen3ForCausalLM"] + loaded = TextEncoderLoader().load_model( + "unused-by-mock", + config, + target_device, + _loader_args(cpu_offload=cpu_offload), + dtype="fp32", + use_text_encoder_override=True, + ) + + assert loaded is fake_model + assert auto_model_calls == [{ + "local_files_only": True, + "dtype": torch.float32, + "low_cpu_mem_usage": True, + }] + assert default_devices == [ambient_device] + assert placements == [expected_device] + assert loaded._fastvideo_input_device == expected_device + + +def test_set_default_torch_dtype_restores_after_exception(): + original = torch.get_default_dtype() + + with pytest.raises(RuntimeError, match="expected failure"): + with set_default_torch_dtype(torch.bfloat16): + raise RuntimeError("expected failure") + + assert torch.get_default_dtype() == original + + +@pytest.mark.parametrize("checkpoint_prefix", ["", "model."]) +def test_qwen3_native_load_accepts_already_fused_weights_and_scales( + checkpoint_prefix: str, +): + """Direct FastVideo/state-dict keys satisfy each fused target atomically.""" + fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True) + expected = { + name: torch.arange(1, param.numel() + 1, dtype=param.dtype) + for name, param in params.items() + } + weights = [ + (f"{checkpoint_prefix}{name}", expected[name]) + for name in params + ] + + loaded = Qwen3ForCausalLM.load_weights(fake_model, weights) + + assert loaded == set(params) + for name, param in params.items(): + assert torch.equal(param.data, expected[name]) + _assert_native_load_surface( + fake_model, + {name for name, _ in weights}, + loaded, + ) + + +def test_qwen3_native_load_rejects_incompatible_fused_quant_scale_shape(): + fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True) + weights = [ + (name, torch.ones_like(param)) + for name, param in params.items() + if name != "layers.0.self_attn.qkv_proj.scale_weight" + ] + weights.append(("layers.0.self_attn.qkv_proj.scale_weight", torch.ones(1))) + + with pytest.raises(AssertionError, match="Attempted to load weight"): + Qwen3ForCausalLM.load_weights(fake_model, weights) + + +def test_qwen3_native_load_uses_real_shard_loader_for_split_quant_scales(): + fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True) + weights = [ + (name, torch.ones_like(param)) + for name, param in params.items() + if name.endswith(".weight") + ] + weights.extend( + (f"model.layers.0.self_attn{weight_name}.scale_weight", torch.tensor(float(index + 1))) + for index, (_, weight_name, _) in enumerate(stacked[:3]) + ) + weights.extend( + (f"model.layers.0.mlp{weight_name}.scale_weight", torch.tensor(float(index + 4))) + for index, (_, weight_name, _) in enumerate(stacked[3:]) + ) + + loaded = Qwen3ForCausalLM.load_weights(fake_model, weights) + + assert loaded == set(params) + assert torch.equal( + params["layers.0.self_attn.qkv_proj.scale_weight"].data, + torch.tensor([1.0, 2.0, 3.0]), + ) + assert torch.equal( + params["layers.0.mlp.gate_up_proj.scale_weight"].data, + torch.tensor([4.0, 5.0]), + ) + + +@pytest.mark.parametrize( + ("missing_weight_name", "missing_shard_id"), + [ + (".q_proj", "q"), + (".k_proj", "k"), + (".v_proj", "v"), + (".gate_proj", 0), + (".up_proj", 1), + ], +) +def test_qwen3_native_load_rejects_missing_fused_weight_shard( + missing_weight_name: str, + missing_shard_id: str | int, +): + fake_model, _, stacked = _fake_qwen3_weight_target() + weights = [ + (f"model.layers.0.self_attn{weight_name}.weight", torch.empty(1)) + for _, weight_name, _ in stacked[:3] + if weight_name != missing_weight_name + ] + weights.extend( + (f"model.layers.0.mlp{weight_name}.weight", torch.empty(1)) + for _, weight_name, _ in stacked[3:] + if weight_name != missing_weight_name + ) + + with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info: + Qwen3ForCausalLM.load_weights(fake_model, weights) + + destination = ( + "layers.0.self_attn.qkv_proj.weight" + if isinstance(missing_shard_id, str) + else "layers.0.mlp.gate_up_proj.weight" + ) + assert f"{destination}[{missing_shard_id}]" in str(exc_info.value) + + +@pytest.mark.parametrize( + ("missing_weight_name", "missing_shard_id"), + [ + (".q_proj", "q"), + (".k_proj", "k"), + (".v_proj", "v"), + (".gate_proj", 0), + (".up_proj", 1), + ], +) +def test_qwen3_native_load_rejects_missing_fused_quant_scale_shard( + missing_weight_name: str, + missing_shard_id: str | int, +): + fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True) + weights = [ + (name, torch.empty(1)) + for name in params + if name.endswith(".weight") + ] + weights.extend( + (f"model.layers.0.self_attn{weight_name}.scale_weight", torch.empty(1)) + for _, weight_name, _ in stacked[:3] + if weight_name != missing_weight_name + ) + weights.extend( + (f"model.layers.0.mlp{weight_name}.scale_weight", torch.empty(1)) + for _, weight_name, _ in stacked[3:] + if weight_name != missing_weight_name + ) + + with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info: + Qwen3ForCausalLM.load_weights(fake_model, weights) + + destination = ( + "layers.0.self_attn.qkv_proj.scale_weight" + if isinstance(missing_shard_id, str) + else "layers.0.mlp.gate_up_proj.scale_weight" + ) + assert f"{destination}[{missing_shard_id}]" in str(exc_info.value) + + +@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required") +@pytest.mark.parametrize( + "dtype", + [ + pytest.param(torch.float32, id="fp32"), + pytest.param( + torch.bfloat16, + id="bf16", + marks=pytest.mark.skipif( + not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()), + reason="bf16 parity requires a bf16-capable CUDA device", + ), + ), + ], +) +def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype): + if not ZIMAGE_TOKENIZER_DIR.exists(): + pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + torch.manual_seed(11) + + tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True) + + prompts = [ + "A cinematic shot of a rainy neon street at night.", + "A watercolor illustration of a fox in autumn leaves.", + ] + toks = _tokenize_official_prompts(tokenizer, prompts) + input_ids = toks["input_ids"].to(device) + attention_mask = toks["attention_mask"].to(device) + config = _load_qwen3_config() + + ref = AutoModel.from_pretrained( + str(ZIMAGE_TEXT_ENCODER_DIR), + local_files_only=True, + trust_remote_code=True, + dtype=dtype, + low_cpu_mem_usage=True, + ).eval().to(device=device) + + with torch.no_grad(): + ref_out = ref( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + ref_last = ref_out.last_hidden_state.detach().float().cpu() + assert ref_out.hidden_states is not None + ref_hs_m2 = ref_out.hidden_states[-2].detach().float().cpu() + + print(f"[mem] after HF ref forward: {_gpu_mem_mb():.0f} MiB", flush=True) + # Keep only CPU outputs so each full encoder is resident one at a time. + del ref, ref_out + _reclaim_vram() + print(f"[mem] after freeing HF ref: {_gpu_mem_mb():.0f} MiB", flush=True) + + production = TextEncoderLoader().load_model( + str(ZIMAGE_TEXT_ENCODER_DIR), + config, + device, + _loader_args(cpu_offload=False), + dtype=_precision_name(dtype), + use_text_encoder_override=True, + ) + assert production.__class__.__name__ == "Qwen3Model" + assert not hasattr(production, "lm_head") + assert production._fastvideo_input_device == device + + with torch.no_grad(): + production_out = production( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + production_last = production_out.last_hidden_state.detach().float().cpu() + assert production_out.hidden_states is not None + production_hs_m2 = production_out.hidden_states[-2].detach().float().cpu() + + # Production loading must preserve the independent Transformers oracle + # exactly; native-kernel drift is assessed separately below. + assert_close(ref_last, production_last, atol=0.0, rtol=0.0) + assert_close(ref_hs_m2, production_hs_m2, atol=0.0, rtol=0.0) + del production, production_out, production_last, production_hs_m2 + _reclaim_vram() + + fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model") + fv = fv_cls(config).eval() + loaded = fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR)) + assert loaded, "No Qwen3 weights were loaded into FastVideo model" + fv = fv.to(device=device, dtype=dtype) + print(f"[mem] after FastVideo load: {_gpu_mem_mb():.0f} MiB", flush=True) + + _assert_native_load_surface( + fv, + _pretrained_safetensor_keys(ZIMAGE_TEXT_ENCODER_DIR), + loaded, + ) + + with torch.no_grad(): + fv_out = fv( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + fv_last = fv_out.last_hidden_state.detach().float().cpu() + assert fv_out.hidden_states is not None + fv_hs_m2 = fv_out.hidden_states[-2].detach().float().cpu() + + # Z-Image uses only valid token embeddings selected by attention mask. + # Compare parity on that exact path to avoid padded-token artifacts. + mask_cpu = attention_mask.detach().bool().cpu() + is_bf16 = dtype == torch.bfloat16 + + for i in range(mask_cpu.shape[0]): + valid = mask_cpu[i] + ref_last_v = ref_last[i][valid] + fv_last_v = fv_last[i][valid] + ref_hs_m2_v = ref_hs_m2[i][valid] + fv_hs_m2_v = fv_hs_m2[i][valid] + + last_diff = _print_diag(f"batch{i} last_hidden_state {dtype}", ref_last_v, fv_last_v) + hs_m2_diff = _print_diag(f"batch{i} hidden_states[-2] {dtype}", ref_hs_m2_v, fv_hs_m2_v) + + if is_bf16: + # Distribution checks instead of element-wise assert_close: cross- + # kernel bf16 produces a long max tail (fused vs unfused linears), + # but a real op bug pushes mean *and* median high. Mean catches + # systematic drift; median catches "wrong weights / swapped + # layers" where most elements diverge. Pre-norm tensors get a + # looser bound because the final RMSNorm compresses drift ~5x. + assert last_diff.mean().item() < _BF16_MEAN_DRIFT_POST_NORM, ( + f"batch{i} last_hidden_state bf16 mean drift " + f"{last_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_POST_NORM}" + ) + assert last_diff.median().item() < _BF16_MEDIAN_DRIFT_POST_NORM, ( + f"batch{i} last_hidden_state bf16 median drift " + f"{last_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_POST_NORM}" + ) + assert hs_m2_diff.mean().item() < _BF16_MEAN_DRIFT_PRE_NORM, ( + f"batch{i} hidden_states[-2] bf16 mean drift " + f"{hs_m2_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_PRE_NORM}" + ) + assert hs_m2_diff.median().item() < _BF16_MEDIAN_DRIFT_PRE_NORM, ( + f"batch{i} hidden_states[-2] bf16 median drift " + f"{hs_m2_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_PRE_NORM}" + ) + else: + assert_close(ref_last_v, fv_last_v, atol=_FP32_ATOL, rtol=_FP32_RTOL) + assert_close(ref_hs_m2_v, fv_hs_m2_v, atol=_FP32_ATOL, rtol=_FP32_RTOL) + + +@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required") +@pytest.mark.skipif( + not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()), + reason="bf16 per-layer diagnostic requires a bf16-capable CUDA device", +) +def test_zimage_qwen3_encoder_per_layer_bf16_diagnostic(): + """Diagnostic-only: prints per-layer FastVideo-vs-HF drift for the bf16 + forward to distinguish monotonic accumulation (bf16 tail) from a + single-layer spike (real op mismatch). + + This test always passes; review the captured stdout to interpret. + Expected pattern for healthy bf16 accumulation: max/mean/median grow + smoothly with layer depth. Suspect pattern: one layer's diff is + >>2x the previous, suggesting a real bug at that block. + """ + if not ZIMAGE_TOKENIZER_DIR.exists(): + pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(11) + + tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True) + toks = _tokenize_official_prompts( + tokenizer, + ["A cinematic shot of a rainy neon street at night."], + ) + input_ids = toks["input_ids"].to(device) + attention_mask = toks["attention_mask"].to(device) + + ref = AutoModel.from_pretrained( + str(ZIMAGE_TEXT_ENCODER_DIR), + local_files_only=True, + trust_remote_code=True, + dtype=dtype, + low_cpu_mem_usage=True, + ).eval().to(device=device) + + with torch.no_grad(): + ref_out = ref( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + # Cache the reference hidden states on CPU, then free the HF model so + # only one encoder is resident when FastVideo runs (avoids OOM). + ref_hidden_cpu = [h.detach().float().cpu() for h in ref_out.hidden_states] + del ref, ref_out + _reclaim_vram() + + fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model") + cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json") + for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"): + cfg_raw.pop(k, None) + cfg = Qwen3TextConfig() + cfg.update_model_arch(cfg_raw) + fv = fv_cls(cfg).eval() + fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR)) + fv = fv.to(device=device, dtype=dtype) + + with torch.no_grad(): + fv_out = fv( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + fv_hidden_cpu = [h.detach().float().cpu() for h in fv_out.hidden_states] + + assert fv_out.hidden_states is not None + assert len(ref_hidden_cpu) == 37, ( + f"expected embeddings plus 36 block outputs, got {len(ref_hidden_cpu)} entries" + ) + assert len(ref_hidden_cpu) == len(fv_hidden_cpu), ( + f"len mismatch: HF={len(ref_hidden_cpu)} FV={len(fv_hidden_cpu)}" + ) + + mask_cpu = attention_mask.detach().bool().cpu()[0] + print( + f"\n[per-layer bf16 diagnostic] num_hidden_states={len(ref_hidden_cpu)} " + f"valid_tokens={mask_cpu.sum().item()}/{mask_cpu.numel()}", + flush=True, + ) + print( + f"{'idx':>3} {'kind':<10} {'max':>8} {'mean':>8} {'median':>8} {'p99':>8}", + flush=True, + ) + for idx, (ref_h, fv_h) in enumerate(zip(ref_hidden_cpu, fv_hidden_cpu)): + ref_v = ref_h[0][mask_cpu] + fv_v = fv_h[0][mask_cpu] + diff = (ref_v - fv_v).abs() + flat = diff.flatten() + p99 = torch.quantile(flat, 0.99).item() + if idx == 0: + kind = "embedding" + elif idx == len(ref_hidden_cpu) - 1: + kind = "post-norm" + else: + kind = f"layer{idx - 1}-out" + print( + f"{idx:>3} {kind:<10} " + f"{diff.max().item():>8.4f} {diff.mean().item():>8.4f} " + f"{diff.median().item():>8.4f} {p99:>8.4f}", + flush=True, + ) diff --git a/tests/local_tests/zimage/test_zimage_scheduler_parity.py b/tests/local_tests/zimage/test_zimage_scheduler_parity.py new file mode 100644 index 000000000..5a867f94f --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_scheduler_parity.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Parity tests for Z-Image FlowMatchEulerDiscreteScheduler. + +These tests compare FastVideo's scheduler implementation against the +reference scheduler shipped in the local Z-Image repository. + +Usage: + pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v +""" + +from __future__ import annotations + +import importlib +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler as FastVideoScheduler, +) + + +REPO_ROOT = Path(__file__).resolve().parents[3] +ZIMAGE_REPO = REPO_ROOT / "Z-Image" +ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src" +ZIMAGE_SCHEDULER_CFG = ( + REPO_ROOT / "official_weights" / "Z-Image" / "scheduler" / "scheduler_config.json" +) +ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" +PARITY_SCOPE = "implementation_subcomponent" + + +def _require_pinned_reference_module(module_name: str, source_file: Path): + if not ZIMAGE_REPO.exists(): + pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}") + if not source_file.is_file(): + pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}") + + try: + result = subprocess.run( + ["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + pytest.fail(f"Cannot verify Z-Image reference revision: {exc}") + + actual_revision = result.stdout.strip() + assert actual_revision == ZIMAGE_REFERENCE_REVISION, ( + "Z-Image reference clone is not at the pinned revision: " + f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}" + ) + + if str(ZIMAGE_SRC) not in sys.path: + sys.path.insert(0, str(ZIMAGE_SRC)) + try: + module = importlib.import_module(module_name) + except Exception as exc: + pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}") + + module_file = Path(module.__file__ or "").resolve() + assert module_file.is_relative_to(ZIMAGE_SRC.resolve()), ( + f"{module_name} resolved outside the pinned clone: {module_file}" + ) + return module + + +@pytest.fixture(scope="module") +def reference_scheduler_cls(): + module = _require_pinned_reference_module( + "zimage.scheduler", + ZIMAGE_SRC / "zimage" / "scheduler.py", + ) + return module.FlowMatchEulerDiscreteScheduler + + +@pytest.fixture(scope="module") +def scheduler_config() -> dict: + if not ZIMAGE_SCHEDULER_CFG.exists(): + pytest.skip(f"Z-Image scheduler config not found: {ZIMAGE_SCHEDULER_CFG}") + with ZIMAGE_SCHEDULER_CFG.open("r", encoding="utf-8") as f: + return json.load(f) + + +# Diffusers-style scheduler-config keys we do NOT want to forward — they are +# loader/state metadata, not constructor kwargs. Anything else in +# scheduler_config.json (including future keys like `time_shift_type`, +# `invert_sigmas`, etc.) must be forwarded so parity actually exercises the +# config-on-disk and not a hand-curated subset. +_SCHEDULER_CONFIG_LOADER_KEYS = frozenset({"_class_name", "_diffusers_version", "_name_or_path"}) + + +def _scheduler_kwargs_from_config(scheduler_config: dict) -> dict: + return {k: v for k, v in scheduler_config.items() if k not in _SCHEDULER_CONFIG_LOADER_KEYS} + + +def _make_scheduler_pair(reference_scheduler_cls, scheduler_config: dict, **overrides): + kwargs = _scheduler_kwargs_from_config(scheduler_config) + kwargs.update(overrides) + + ref = reference_scheduler_cls(**kwargs) + # The pinned Z-Image pipeline mutates this immediately before scheduling. + ref.sigma_min = 0.0 + + # These values must be serialized by the future production pipeline config. + # setdefault also allows this test to consume that config once it lands. + fv_kwargs = dict(kwargs) + fv_kwargs.setdefault("use_reference_discrete_timesteps", True) + fv_kwargs.setdefault("sigma_min", 0.0) + fv = FastVideoScheduler(**fv_kwargs) + return ref, fv + + +@pytest.fixture(scope="module") +def scheduler_pair(reference_scheduler_cls, scheduler_config: dict): + return _make_scheduler_pair(reference_scheduler_cls, scheduler_config) + + +def _run_step_loop(ref_scheduler, fv_scheduler, sample_shape=(2, 4, 32, 32)): + torch.manual_seed(123) + + ref_sample = torch.randn(sample_shape, dtype=torch.float32) + fv_sample = ref_sample.clone() + + for t in ref_scheduler.timesteps: + model_output = torch.randn_like(ref_sample) + + ref_next = ref_scheduler.step( + model_output, + t, + ref_sample, + return_dict=False, + )[0] + fv_next = fv_scheduler.step( + model_output, + t, + fv_sample, + return_dict=False, + )[0] + + assert_close(ref_next, fv_next, atol=1e-6, rtol=1e-6) + + ref_sample = ref_next + fv_sample = fv_next + + +def test_zimage_scheduler_parity_default_schedule_and_step(scheduler_pair): + ref, fv = scheduler_pair + + num_inference_steps = 8 + ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu") + fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu") + + assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6) + assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6) + + _run_step_loop(ref, fv) + + +def test_scheduler_positional_args_keep_existing_bindings_and_default_schedule(): + scheduler = FastVideoScheduler(1000, 1.0, False, 0.7) + + assert scheduler.config.base_shift == 0.7 + assert scheduler.config.use_reference_discrete_timesteps is False + + scheduler.set_timesteps(num_inference_steps=8, device="cpu") + assert_close( + scheduler.sigmas[-2], + torch.tensor(scheduler.sigma_min, dtype=torch.float32), + ) + + +def test_scheduler_honors_explicit_zero_sigma_min(): + scheduler = FastVideoScheduler( + sigma_min=0.0, + use_reference_discrete_timesteps=True, + ) + + assert scheduler.sigma_min == 0.0 + scheduler.set_timesteps(num_inference_steps=8, device="cpu") + assert_close( + scheduler.timesteps[-1], + torch.tensor(125.0, dtype=torch.float32), + ) + + +def test_reference_schedule_preserves_float64_linspace_rounding(): + scheduler = FastVideoScheduler( + sigma_min=0.0, + use_reference_discrete_timesteps=True, + ) + scheduler.set_timesteps(num_inference_steps=9, device="cpu") + + expected = torch.from_numpy( + np.linspace(1000.0, 0.0, 10)[:-1] / 1000.0, + ).to(torch.float32) * 1000.0 + float32_regression = torch.from_numpy( + np.linspace(1000.0, 0.0, 10, dtype=np.float32)[:-1] / np.float32(1000.0), + ) * 1000.0 + + # The fifth timestep differs by one float32 ULP depending on where the + # linspace is rounded, so this assertion fails if dtype=np.float32 returns. + assert expected[4].item() != float32_regression[4].item() + assert scheduler.timesteps[4].item() == expected[4].item() + + +def test_zimage_scheduler_parity_dynamic_shifting_with_mu( + reference_scheduler_cls, + scheduler_config: dict, +): + ref, fv = _make_scheduler_pair( + reference_scheduler_cls, + scheduler_config, + use_dynamic_shifting=True, + ) + + mu = 0.75 + num_inference_steps = 8 + ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu) + fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu) + + assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6) + assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6) + + _run_step_loop(ref, fv) diff --git a/tests/local_tests/zimage/test_zimage_tokenizer_parity.py b/tests/local_tests/zimage/test_zimage_tokenizer_parity.py new file mode 100644 index 000000000..978ee1306 --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_tokenizer_parity.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Parity test for Z-Image tokenizer loading and tokenization behavior. + +This compares: +1) direct transformers AutoTokenizer loading from local tokenizer dir, and +2) FastVideo TokenizerLoader loading path, + +using the exact chat-template and tokenization settings from the pinned +Z-Image pipeline. + +Usage: + pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close +from transformers import AutoTokenizer + +from fastvideo.models.loader.component_loader import TokenizerLoader + + +REPO_ROOT = Path(__file__).resolve().parents[3] +ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer" +PARITY_SCOPE = "production_loader" +OFFICIAL_MAX_SEQUENCE_LENGTH = 512 +OFFICIAL_CHAT_TEMPLATE_KWARGS = { + "tokenize": False, + "add_generation_prompt": True, + "enable_thinking": True, +} + + +@dataclass +class _DummyPipelineConfig: + text_encoder_configs: tuple = field(default_factory=tuple) + + +@dataclass +class _DummyFastVideoArgs: + pipeline_config: _DummyPipelineConfig = field(default_factory=_DummyPipelineConfig) + + +def _load_reference_tokenizer(): + if not ZIMAGE_TOKENIZER_DIR.exists(): + pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") + return AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True) + + +def _load_fastvideo_tokenizer(): + if not ZIMAGE_TOKENIZER_DIR.exists(): + pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}") + loader = TokenizerLoader() + return loader.load(str(ZIMAGE_TOKENIZER_DIR), _DummyFastVideoArgs()) + + +def test_zimage_tokenizer_loader_and_tokenization_parity(): + ref_tok = _load_reference_tokenizer() + fv_tok = _load_fastvideo_tokenizer() + + prompts = [ + "A close-up portrait of a cat in warm studio lighting.", + "An oil painting of a lighthouse at sunset.", + ] + + tok_kwargs = { + "padding": "max_length", + "max_length": OFFICIAL_MAX_SEQUENCE_LENGTH, + "truncation": True, + "return_tensors": "pt", + } + + ref_out = ref_tok(prompts, **tok_kwargs) + fv_out = fv_tok(prompts, **tok_kwargs) + + assert_close(ref_out["input_ids"], fv_out["input_ids"], atol=0, rtol=0) + assert_close(ref_out["attention_mask"], fv_out["attention_mask"], atol=0, rtol=0) + + # Compare key tokenizer attributes that affect generation-time behavior. + assert ref_tok.padding_side == fv_tok.padding_side + assert ref_tok.pad_token_id == fv_tok.pad_token_id + assert ref_tok.eos_token_id == fv_tok.eos_token_id + + +@pytest.mark.skipif(not ZIMAGE_TOKENIZER_DIR.exists(), reason="Z-Image tokenizer assets are required") +def test_zimage_tokenizer_chat_template_parity(): + ref_tok = _load_reference_tokenizer() + fv_tok = _load_fastvideo_tokenizer() + + assert callable(getattr(ref_tok, "apply_chat_template", None)), ( + "Pinned Z-Image tokenizer must expose apply_chat_template" + ) + assert callable(getattr(fv_tok, "apply_chat_template", None)), ( + "FastVideo TokenizerLoader dropped the required apply_chat_template API" + ) + + messages = [{"role": "user", "content": "Describe a futuristic city skyline."}] + + ref_text = ref_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS) + fv_text = fv_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS) + + assert ref_text == fv_text + + # Z-Image deliberately enables Qwen's thinking template. Guard against a + # pipeline silently using the generic non-thinking prompt path. + non_thinking_text = fv_tok.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + assert fv_text != non_thinking_text + + tokenization_kwargs = { + "padding": "max_length", + "max_length": OFFICIAL_MAX_SEQUENCE_LENGTH, + "truncation": True, + "return_tensors": "pt", + } + ref_tokens = ref_tok(ref_text, **tokenization_kwargs) + fv_tokens = fv_tok(fv_text, **tokenization_kwargs) + + assert_close(ref_tokens["input_ids"], fv_tokens["input_ids"], atol=0, rtol=0) + assert_close(ref_tokens["attention_mask"], fv_tokens["attention_mask"], atol=0, rtol=0) 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..7a0971302 --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_transformer_parity.py @@ -0,0 +1,344 @@ +# 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): + from fastvideo.models.registry import ModelRegistry + + model_cls, _ = ModelRegistry.resolve_model_cls("ZImageTransformer2DModel") + assert model_cls is ZImageTransformer2DModel + + 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) diff --git a/tests/local_tests/zimage/test_zimage_vae_parity.py b/tests/local_tests/zimage/test_zimage_vae_parity.py new file mode 100644 index 000000000..aa6ce42e6 --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_vae_parity.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Parity test for Z-Image VAE (AutoencoderKL decode path). + +This compares FastVideo's AutoencoderKL wrapper against the Z-Image +reference AutoencoderKL implementation using the same local weights. + +Usage: + pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v +""" + +from __future__ import annotations + +import importlib +import json +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import load_file as safetensors_load_file +from torch.testing import assert_close + +from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig +from fastvideo.models.loader.component_loader import VAELoader +from fastvideo.models.vaes.autoencoder_kl import AutoencoderKL as FastVideoAutoencoderKL + + +REPO_ROOT = Path(__file__).resolve().parents[3] +ZIMAGE_REPO = REPO_ROOT / "Z-Image" +ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src" +ZIMAGE_VAE_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "vae" +ZIMAGE_VAE_CFG = ZIMAGE_VAE_DIR / "config.json" +ZIMAGE_VAE_WEIGHTS = ZIMAGE_VAE_DIR / "diffusion_pytorch_model.safetensors" +ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" +ZIMAGE_VAE_SCALING_FACTOR = 0.3611 +ZIMAGE_VAE_SHIFT_FACTOR = 0.1159 +PARITY_SCOPE = "both" + + +def _require_pinned_reference_module(module_name: str, source_file: Path): + if not ZIMAGE_REPO.exists(): + pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}") + if not source_file.is_file(): + pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}") + + try: + result = subprocess.run( + ["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + pytest.fail(f"Cannot verify Z-Image reference revision: {exc}") + + actual_revision = result.stdout.strip() + assert actual_revision == ZIMAGE_REFERENCE_REVISION, ( + "Z-Image reference clone is not at the pinned revision: " + f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}" + ) + + if str(ZIMAGE_SRC) not in sys.path: + sys.path.insert(0, str(ZIMAGE_SRC)) + try: + module = importlib.import_module(module_name) + except Exception as exc: + pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}") + + module_file = Path(module.__file__ or "").resolve() + assert module_file.is_relative_to(ZIMAGE_SRC.resolve()), ( + f"{module_name} resolved outside the pinned clone: {module_file}" + ) + return module + + +@pytest.fixture(scope="module") +def reference_autoencoder_cls(): + module = _require_pinned_reference_module( + "zimage.autoencoder", + ZIMAGE_SRC / "zimage" / "autoencoder.py", + ) + return module.AutoencoderKL + + +def _load_cfg() -> dict: + if not ZIMAGE_VAE_CFG.exists(): + pytest.skip(f"Z-Image VAE config not found: {ZIMAGE_VAE_CFG}") + with ZIMAGE_VAE_CFG.open("r", encoding="utf-8") as f: + cfg = json.load(f) + cfg.pop("_class_name", None) + cfg.pop("_diffusers_version", None) + cfg.pop("_name_or_path", None) + return cfg + + +def _load_weights() -> dict[str, torch.Tensor]: + if not ZIMAGE_VAE_WEIGHTS.exists(): + pytest.skip(f"Z-Image VAE weights not found: {ZIMAGE_VAE_WEIGHTS}") + return safetensors_load_file(str(ZIMAGE_VAE_WEIGHTS), device="cpu") + + +def _build_reference(reference_autoencoder_cls, cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module: + ref = reference_autoencoder_cls( + in_channels=cfg["in_channels"], + out_channels=cfg["out_channels"], + down_block_types=tuple(cfg["down_block_types"]), + up_block_types=tuple(cfg["up_block_types"]), + block_out_channels=tuple(cfg["block_out_channels"]), + layers_per_block=cfg["layers_per_block"], + latent_channels=cfg["latent_channels"], + norm_num_groups=cfg["norm_num_groups"], + scaling_factor=cfg["scaling_factor"], + shift_factor=cfg["shift_factor"], + use_quant_conv=cfg["use_quant_conv"], + use_post_quant_conv=cfg["use_post_quant_conv"], + mid_block_add_attention=cfg["mid_block_add_attention"], + ).eval() + ref.load_state_dict(weights, strict=True) + return ref + + +def _build_fastvideo(cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module: + fv_cfg = AutoencoderKLVAEConfig() + fv_cfg.update_model_arch(cfg) + fv = FastVideoAutoencoderKL(fv_cfg).eval() + fv.load_state_dict(weights, strict=True) + return fv + + +def _load_fastvideo_production_vae(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "fastvideo.models.loader.component_loader.get_local_torch_device", + lambda: torch.device("cpu"), + ) + pipeline_config = SimpleNamespace( + vae_config=AutoencoderKLVAEConfig(), + vae_precision="fp32", + ) + fastvideo_args = SimpleNamespace( + model_paths={}, + pipeline_config=pipeline_config, + vae_cpu_offload=False, + ) + vae = VAELoader().load(str(ZIMAGE_VAE_DIR), fastvideo_args) + assert isinstance(vae, FastVideoAutoencoderKL) + assert fastvideo_args.model_paths["vae"] == str(ZIMAGE_VAE_DIR) + return vae + + +def test_zimage_vae_decode_parity(reference_autoencoder_cls): + torch.manual_seed(7) + + cfg = _load_cfg() + weights = _load_weights() + + ref = _build_reference(reference_autoencoder_cls, cfg, weights) + fv = _build_fastvideo(cfg, weights) + + # Decode-only parity is the critical path for Z-Image pipeline integration. + latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32) + + with torch.no_grad(): + ref_out = ref.decode(latents, return_dict=False)[0].detach().float() + fv_out = fv.decode(latents, return_dict=False)[0].detach().float() + + assert ref_out.shape == fv_out.shape + assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4) + + +def test_zimage_vae_production_loader_and_raw_decode_parity( + reference_autoencoder_cls, + monkeypatch: pytest.MonkeyPatch, +): + torch.manual_seed(17) + + cfg = _load_cfg() + weights = _load_weights() + assert cfg["scaling_factor"] == ZIMAGE_VAE_SCALING_FACTOR + assert cfg["shift_factor"] == ZIMAGE_VAE_SHIFT_FACTOR + ref = _build_reference(reference_autoencoder_cls, cfg, weights) + fv = _load_fastvideo_production_vae(monkeypatch) + + assert ref.config.scaling_factor == cfg["scaling_factor"] + assert fv.config.scaling_factor == cfg["scaling_factor"] + assert ref.config.shift_factor == cfg["shift_factor"] + assert fv.config.shift_factor == cfg["shift_factor"] + + decode_latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32) + + with torch.no_grad(): + ref_out = ref.decode(decode_latents, return_dict=False)[0].detach().float() + fv_out = fv.decode(decode_latents, return_dict=False)[0].detach().float() + + assert ref_out.shape == fv_out.shape + assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4)