From 42658fad8e817b97368924cff99b4b28a5536d0f Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Sun, 23 Nov 2025 19:54:24 +0000 Subject: [PATCH 01/31] Add Support for Z-Image. --- src/diffusers/__init__.py | 3 + src/diffusers/hooks/_helpers.py | 20 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_z_image.py | 756 ++++++++++++++++++ src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/z_image/__init__.py | 51 ++ .../pipelines/z_image/pipeline_output.py | 35 + .../pipelines/z_image/pipeline_z_image.py | 622 ++++++++++++++ tests/pipelines/z_image/__init__.py | 0 10 files changed, 1492 insertions(+) create mode 100644 src/diffusers/models/transformers/transformer_z_image.py create mode 100644 src/diffusers/pipelines/z_image/__init__.py create mode 100644 src/diffusers/pipelines/z_image/pipeline_output.py create mode 100644 src/diffusers/pipelines/z_image/pipeline_z_image.py create mode 100644 tests/pipelines/z_image/__init__.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index cd7a2cb581b7..a883a361e6db 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -271,6 +271,7 @@ "WanAnimateTransformer3DModel", "WanTransformer3DModel", "WanVACETransformer3DModel", + "ZImageTransformer2DModel", "attention_backend", ] ) @@ -647,6 +648,7 @@ "WuerstchenCombinedPipeline", "WuerstchenDecoderPipeline", "WuerstchenPriorPipeline", + "ZImagePipeline", ] ) @@ -1329,6 +1331,7 @@ WuerstchenCombinedPipeline, WuerstchenDecoderPipeline, WuerstchenPriorPipeline, + ZImagePipeline, ) try: diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 790199f3c978..da7313cb4737 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -111,6 +111,7 @@ def _register_attention_processors_metadata(): from ..models.transformers.transformer_hunyuanimage import HunyuanImageAttnProcessor from ..models.transformers.transformer_qwenimage import QwenDoubleStreamAttnProcessor2_0 from ..models.transformers.transformer_wan import WanAttnProcessor2_0 + from ..models.transformers.transformer_z_image import ZSingleStreamAttnProcessor # AttnProcessor2_0 AttentionProcessorRegistry.register( @@ -158,6 +159,14 @@ def _register_attention_processors_metadata(): ), ) + # ZSingleStreamAttnProcessor + AttentionProcessorRegistry.register( + model_class=ZSingleStreamAttnProcessor, + metadata=AttentionProcessorMetadata( + skip_processor_output_fn=_skip_proc_output_fn_Attention_ZSingleStreamAttnProcessor, + ), + ) + def _register_transformer_blocks_metadata(): from ..models.attention import BasicTransformerBlock @@ -179,6 +188,7 @@ def _register_transformer_blocks_metadata(): from ..models.transformers.transformer_mochi import MochiTransformerBlock from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock from ..models.transformers.transformer_wan import WanTransformerBlock + from ..models.transformers.transformer_z_image import ZImageTransformerBlock # BasicTransformerBlock TransformerBlockRegistry.register( @@ -312,6 +322,15 @@ def _register_transformer_blocks_metadata(): ), ) + # ZImage + TransformerBlockRegistry.register( + model_class=ZImageTransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=0, + return_encoder_hidden_states_index=None, + ), + ) + # fmt: off def _skip_attention___ret___hidden_states(self, *args, **kwargs): @@ -338,4 +357,5 @@ def _skip_attention___ret___hidden_states___encoder_hidden_states(self, *args, * _skip_proc_output_fn_Attention_FluxAttnProcessor = _skip_attention___ret___hidden_states _skip_proc_output_fn_Attention_QwenDoubleStreamAttnProcessor2_0 = _skip_attention___ret___hidden_states _skip_proc_output_fn_Attention_HunyuanImageAttnProcessor = _skip_attention___ret___hidden_states +_skip_proc_output_fn_Attention_ZSingleStreamAttnProcessor = _skip_attention___ret___hidden_states # fmt: on diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index b42e981f71a9..af44feb00128 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -110,6 +110,7 @@ _import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"] _import_structure["transformers.transformer_wan_animate"] = ["WanAnimateTransformer3DModel"] _import_structure["transformers.transformer_wan_vace"] = ["WanVACETransformer3DModel"] + _import_structure["transformers.transformer_z_image"] = ["ZImageTransformer2DModel"] _import_structure["unets.unet_1d"] = ["UNet1DModel"] _import_structure["unets.unet_2d"] = ["UNet2DModel"] _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"] @@ -218,6 +219,7 @@ WanAnimateTransformer3DModel, WanTransformer3DModel, WanVACETransformer3DModel, + ZImageTransformer2DModel, ) from .unets import ( I2VGenXLUNet, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 826469237fb1..84a41c2d470d 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -44,3 +44,4 @@ from .transformer_wan import WanTransformer3DModel from .transformer_wan_animate import WanAnimateTransformer3DModel from .transformer_wan_vace import WanVACETransformer3DModel + from .transformer_z_image import ZImageTransformer2DModel diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py new file mode 100644 index 000000000000..45a609f3eeb7 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -0,0 +1,756 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +import math +from typing import List, Optional, Tuple + +from einops import rearrange +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from flash_attn import flash_attn_varlen_func +except ImportError: + flash_attn_varlen_func = None + +try: + from apex.normalization import FusedRMSNorm as RMSNorm +except ImportError: + from torch.nn import RMSNorm + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...models.modeling_utils import ModelMixin +from ...utils.torch_utils import maybe_allow_in_graph +from ...models.attention_processor import Attention +from ...models.attention_dispatch import dispatch_attention_fn + +ADALN_EMBED_DIM = 256 +SEQ_MULTI_OF = 32 + + +class TimestepEmbedder(nn.Module): + def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): + super().__init__() + if mid_size is None: + mid_size = out_size + self.mlp = nn.Sequential( + nn.Linear( + frequency_embedding_size, + mid_size, + bias=True, + ), + nn.SiLU(), + nn.Linear( + mid_size, + out_size, + bias=True, + ), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.zeros_(self.mlp[0].bias) + nn.init.normal_(self.mlp[2].weight, std=0.02) + nn.init.zeros_(self.mlp[2].bias) + + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + with torch.amp.autocast("cuda", enabled=False): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=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): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + return t_emb + + +class ZSingleStreamAttnProcessor: + """ + Processor for Z-Image single stream attention that adapts the existing Attention class + to match the behavior of the original Z-ImageAttention module. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + x_cu_seqlens: Optional[torch.Tensor] = None, + x_max_item_seqlen: Optional[int] = None, + ) -> torch.Tensor: + x_shard = hidden_states + x_freqs_cis_shard = image_rotary_emb + + query = attn.to_q(x_shard) + key = attn.to_k(x_shard) + value = attn.to_v(x_shard) + + seqlen_shard = x_shard.shape[0] + + # Reshape to [seq_len, heads, head_dim] + head_dim = query.shape[-1] // attn.heads + query = query.view(seqlen_shard, attn.heads, head_dim) + key = key.view(seqlen_shard, attn.heads, head_dim) + value = value.view(seqlen_shard, attn.heads, head_dim) + # Apply Norms + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + # Apply RoPE + 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)) + freqs_cis = freqs_cis.unsqueeze(1) + x_out = torch.view_as_real(x * freqs_cis).flatten(2) + return x_out.type_as(x_in) + + if x_freqs_cis_shard is not None: + query = apply_rotary_emb(query, x_freqs_cis_shard) + key = apply_rotary_emb(key, x_freqs_cis_shard) + + # Cast to correct dtype + dtype = query.dtype + query, key = query.to(dtype), key.to(dtype) + + # Flash Attention + softmax_scale = math.sqrt(1 / head_dim) + assert dtype in [torch.float16, torch.bfloat16] + + if x_cu_seqlens is None or x_max_item_seqlen is None: + raise ValueError("x_cu_seqlens and x_max_item_seqlen are required for ZSingleStreamAttnProcessor") + + if flash_attn_varlen_func is not None: + output = flash_attn_varlen_func( + query, + key, + value, + cu_seqlens_q=x_cu_seqlens, + cu_seqlens_k=x_cu_seqlens, + max_seqlen_q=x_max_item_seqlen, + max_seqlen_k=x_max_item_seqlen, + dropout_p=0.0, + causal=False, + softmax_scale=softmax_scale, + ) + output = output.flatten(-2) + else: + seqlens = (x_cu_seqlens[1:] - x_cu_seqlens[:-1]).cpu().tolist() + + q_split = torch.split(query, seqlens, dim=0) + k_split = torch.split(key, seqlens, dim=0) + v_split = torch.split(value, seqlens, dim=0) + + q_padded = torch.nn.utils.rnn.pad_sequence(q_split, batch_first=True) + k_padded = torch.nn.utils.rnn.pad_sequence(k_split, batch_first=True) + v_padded = torch.nn.utils.rnn.pad_sequence(v_split, batch_first=True) + + batch_size, max_seqlen, _, _ = q_padded.shape + + mask = torch.zeros((batch_size, max_seqlen), dtype=torch.bool, device=query.device) + for i, l in enumerate(seqlens): + mask[i, :l] = True + + attn_mask = torch.zeros((batch_size, 1, 1, max_seqlen), dtype=query.dtype, device=query.device) + attn_mask.masked_fill_(~mask[:, None, None, :], torch.finfo(query.dtype).min) + + q_padded = q_padded.transpose(1, 2) + k_padded = k_padded.transpose(1, 2) + v_padded = v_padded.transpose(1, 2) + + output = F.scaled_dot_product_attention( + q_padded, k_padded, v_padded, attn_mask=attn_mask, dropout_p=0.0, scale=softmax_scale + ) + + output = output.transpose(1, 2) + + out_list = [] + for i, l in enumerate(seqlens): + out_list.append(output[i, :l]) + + output = torch.cat(out_list, dim=0) + output = output.flatten(-2) + + output = attn.to_out[0](output) + if len(attn.to_out) > 1: # dropout + output = attn.to_out[1](output) + + return output + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + nn.init.xavier_uniform_(self.w1.weight) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + nn.init.xavier_uniform_(self.w2.weight) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + nn.init.xavier_uniform_(self.w3.weight) + + def _forward_silu_gating(self, x1, x3): + return F.silu(x1) * x3 + + def forward(self, x): + return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x))) + + +@maybe_allow_in_graph +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, modulation=True + ): + super().__init__() + self.dim = dim + self.head_dim = dim // n_heads + + # Refactored to use diffusers Attention with custom processor + # Original Z-Image params: dim, n_heads, n_kv_heads, qk_norm + self.attention = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // n_heads, + heads=n_heads, + qk_norm="rms_norm" if qk_norm else None, + eps=1e-6, + bias=False, + processor=ZSingleStreamAttnProcessor(), + ) + + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + self.layer_id = layer_id + + 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) + + self.modulation = modulation + if modulation: + self.adaLN_modulation = nn.Sequential( + nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True), + ) + nn.init.zeros_(self.adaLN_modulation[0].weight) + nn.init.zeros_(self.adaLN_modulation[0].bias) + + def forward( + self, + x_shard: torch.Tensor, + x_src_ids_shard: torch.Tensor, + x_freqs_cis_shard: torch.Tensor, + x_cu_seqlens: torch.Tensor, + x_max_item_seqlen: int, + adaln_input: Optional[torch.Tensor] = None, + ): + if self.modulation: + assert adaln_input is not None + scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).chunk(4, dim=1) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + scale_gate_msa = (scale_msa, gate_msa) + scale_gate_mlp = (scale_mlp, gate_mlp) + else: + scale_gate_msa = None + scale_gate_mlp = None + x_src_ids_shard = None + + x_shard = self.attn_forward( + x_shard, x_freqs_cis_shard, x_cu_seqlens, x_max_item_seqlen, scale_gate_msa, x_src_ids_shard + ) + + x_shard = self.ffn_forward(x_shard, scale_gate_mlp, x_src_ids_shard) + + return x_shard + + def attn_forward( + self, + x_shard, + x_freqs_cis_shard, + x_cu_seqlens, + x_max_item_seqlen, + scale_gate: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + x_src_ids_shard: Optional[torch.Tensor] = None, + ): + if self.modulation: + assert scale_gate is not None and x_src_ids_shard is not None + scale_msa, gate_msa = scale_gate + + # Pass extra args needed for ZSingleStreamAttnProcessor + attn_out = self.attention( + self.attention_norm1(x_shard) * scale_msa[x_src_ids_shard], + image_rotary_emb=x_freqs_cis_shard, + x_cu_seqlens=x_cu_seqlens, + x_max_item_seqlen=x_max_item_seqlen + ) + + x_shard = x_shard + gate_msa[x_src_ids_shard] * self.attention_norm2(attn_out) + else: + attn_out = self.attention( + self.attention_norm1(x_shard), + image_rotary_emb=x_freqs_cis_shard, + x_cu_seqlens=x_cu_seqlens, + x_max_item_seqlen=x_max_item_seqlen + ) + x_shard = x_shard + self.attention_norm2(attn_out) + return x_shard + + def ffn_forward( + self, + x_shard, + scale_gate: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + x_src_ids_shard: Optional[torch.Tensor] = None, + ): + if self.modulation: + assert scale_gate is not None and x_src_ids_shard is not None + scale_mlp, gate_mlp = scale_gate + x_shard = x_shard + gate_mlp[x_src_ids_shard] * self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x_shard) * scale_mlp[x_src_ids_shard], + ) + ) + + else: + x_shard = x_shard + self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x_shard), + ) + ) + return x_shard + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), + ) + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x_shard, x_src_ids_shard, c): + scale = 1.0 + self.adaLN_modulation(c) + x_shard = self.norm_final(x_shard) * scale[x_src_ids_shard] + x_shard = self.linear(x_shard) + return x_shard + + +class RopeEmbedder: + def __init__( + self, theta: float = 256.0, axes_dims: List[int] = (16, 56, 56), axes_lens: List[int] = (64, 128, 128) + ): + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length" + self.freqs_cis = None + + @staticmethod + def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): + with torch.device("cpu"): + freqs_cis = [] + for i, (d, e) in enumerate(zip(dim, end)): + freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)) + timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) + freqs = torch.outer(timestep, freqs).float() + freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64) # complex64 + freqs_cis.append(freqs_cis_i) + + return freqs_cis + + def __call__(self, ids: torch.Tensor): + assert ids.ndim == 2 + assert ids.shape[-1] == len(self.axes_dims) + + if self.freqs_cis is None: + self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) + self.freqs_cis = [freqs_cis.cuda() for freqs_cis in self.freqs_cis] + + result = [] + for i in range(len(self.axes_dims)): + index = ids[:, i] + result.append(self.freqs_cis[i][index]) + return torch.cat(result, dim=-1) + + +class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): + _supports_gradient_checkpointing = True + _no_split_modules = ["ZImageTransformerBlock"] + + @register_to_config + def __init__( + self, + 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=[1024, 512, 512], + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.all_patch_size = all_patch_size + self.all_f_patch_size = all_f_patch_size + self.dim = dim + self.n_heads = n_heads + + self.rope_theta = rope_theta + self.t_scale = t_scale + + assert len(all_patch_size) == len(all_f_patch_size) + + all_x_embedder = {} + all_final_layer = {} + for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)): + x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True) + nn.init.xavier_uniform_(x_embedder.weight) + nn.init.constant_(x_embedder.bias, 0.0) + all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder + + final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels) + all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer + + self.all_x_embedder = nn.ModuleDict(all_x_embedder) + self.all_final_layer = nn.ModuleDict(all_final_layer) + self.noise_refiner = nn.ModuleList( + [ + ZImageTransformerBlock(1000 + layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True) + for layer_id in range(n_refiner_layers) + ] + ) + self.context_refiner = nn.ModuleList( + [ + ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=False) + for layer_id in range(n_refiner_layers) + ] + ) + self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024) + self.cap_embedder = nn.Sequential( + RMSNorm(cap_feat_dim, eps=norm_eps), + nn.Linear(cap_feat_dim, dim, bias=True), + ) + nn.init.trunc_normal_(self.cap_embedder[1].weight, std=0.02) + nn.init.zeros_(self.cap_embedder[1].bias) + + self.x_pad_token = nn.Parameter(torch.empty((1, dim))) + nn.init.normal_(self.x_pad_token, std=0.02) + self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) + nn.init.normal_(self.cap_pad_token, std=0.02) + + self.layers = nn.ModuleList( + [ + ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm) + for layer_id in range(n_layers) + ] + ) + head_dim = dim // n_heads + assert head_dim == sum(axes_dims) + self.axes_dims = axes_dims + self.axes_lens = axes_lens + + self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) + + def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size) -> List[torch.Tensor]: + pH = pW = patch_size + pF = f_patch_size + bsz = len(x) + assert len(size) == bsz + for i in range(bsz): + F, H, W = size[i] + ori_len = (F // pF) * (H // pH) * (W // pW) + x[i] = rearrange( + x[i][:ori_len].view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels), + "f h w pf ph pw c -> c (f pf) (h ph) (w pw)", + ) + return x + + @staticmethod + def create_coordinate_grid(size, start=None, device=None): + if start is None: + start = (0 for _ in size) + + axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)] + grids = torch.meshgrid(axes, indexing="ij") + return torch.stack(grids, 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, + ): + + bsz = len(all_image) + pH = pW = patch_size + pF = f_patch_size + device = all_image[0].device + + all_image_out = [] + all_image_size = [] + all_image_pos_ids = [] + all_image_pad_mask = [] + all_cap_pos_ids = [] + all_cap_pad_mask = [] + all_cap_feats_out = [] + + for i, image in enumerate(all_image): + ### LLM Text Encoder + cap_ori_len = len(all_cap_feats[i]) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF + # padded position ids + cap_padded_pos_ids = self.create_coordinate_grid( + size=(cap_ori_len + cap_padding_len, 1, 1), + start=(1, 0, 0), + device=device, + ).flatten(0, 2) + all_cap_pos_ids.append(cap_padded_pos_ids) + # pad mask + all_cap_pad_mask.append( + torch.cat( + [ + torch.zeros((cap_ori_len,), dtype=torch.bool, device=device), + torch.ones((cap_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + ) + # padded feature + cap_padded_feat = torch.cat([all_cap_feats[i], all_cap_feats[i][-1:].repeat(cap_padding_len, 1)], dim=0) + all_cap_feats_out.append(cap_padded_feat) + + ### Process Image + C, F, H, W = image.size() + all_image_size.append((F, H, W)) + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + + image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) + image = rearrange(image, "c f pf h ph w pw -> (f h w) (pf ph pw c)") + + image_ori_len = len(image) + image_padding_len = (-image_ori_len) % SEQ_MULTI_OF + # padded_pos_ids + + image_ori_pos_ids = self.create_coordinate_grid( + size=(F_tokens, H_tokens, W_tokens), + start=(cap_ori_len + cap_padding_len + 1, 0, 0), + device=device, + ).flatten(0, 2) + image_padding_pos_ids = ( + self.create_coordinate_grid( + size=(1, 1, 1), + start=(0, 0, 0), + device=device, + ) + .flatten(0, 2) + .repeat(image_padding_len, 1) + ) + image_padded_pos_ids = torch.cat([image_ori_pos_ids, image_padding_pos_ids], dim=0) + all_image_pos_ids.append(image_padded_pos_ids) + # pad mask + all_image_pad_mask.append( + torch.cat( + [ + torch.zeros((image_ori_len,), dtype=torch.bool, device=device), + torch.ones((image_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + ) + # padded feature + image_padded_feat = torch.cat([image, image[-1:].repeat(image_padding_len, 1)], dim=0) + all_image_out.append(image_padded_feat) + + return ( + all_image_out, + all_cap_feats_out, + all_image_size, + all_image_pos_ids, + all_cap_pos_ids, + all_image_pad_mask, + all_cap_pad_mask, + ) + + def forward( + self, + x: List[torch.Tensor], + t, + cap_feats: List[torch.Tensor], + patch_size=2, + f_patch_size=1, + ): + + assert patch_size in self.all_patch_size + assert f_patch_size in self.all_f_patch_size + + bsz = len(x) + device = x[0].device + t = t * self.t_scale + t = self.t_embedder(t) + + adaln_input = t + + ( + x, + cap_feats, + x_size, + x_pos_ids, + cap_pos_ids, + x_pad_mask, + cap_pad_mask, + ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) + + # x embed & refine + x_item_seqlens = [len(_) for _ in x] + assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) + x_max_item_seqlen = max(x_item_seqlens) + x_cu_seqlens = F.pad( + torch.cumsum(torch.tensor(x_item_seqlens, dtype=torch.int32, device=device), dim=0, dtype=torch.int32), + (1, 0), + ) + x_src_ids = [ + torch.full((count,), i, dtype=torch.int32, device=device) for i, count in enumerate(x_item_seqlens) + ] + x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) + + x_shard = torch.cat(x, dim=0) + x_src_ids_shard = torch.cat(x_src_ids, dim=0) + x_freqs_cis_shard = torch.cat(x_freqs_cis, dim=0) + x_pad_mask_shard = torch.cat(x_pad_mask, dim=0) + del x + + x_shard = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_shard) + x_shard[x_pad_mask_shard] = self.x_pad_token + for layer in self.noise_refiner: + x_shard = layer(x_shard, x_src_ids_shard, x_freqs_cis_shard, x_cu_seqlens, x_max_item_seqlen, adaln_input) + x_flatten = x_shard + + # cap embed & refine + cap_item_seqlens = [len(_) for _ in cap_feats] + assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) + cap_max_item_seqlen = max(cap_item_seqlens) + cap_cu_seqlens = F.pad( + torch.cumsum(torch.tensor(cap_item_seqlens, dtype=torch.int32, device=device), dim=0, dtype=torch.int32), + (1, 0), + ) + cap_src_ids = [ + torch.full((count,), i, dtype=torch.int32, device=device) for i, count in enumerate(cap_item_seqlens) + ] + cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) + + cap_shard = torch.cat(cap_feats, dim=0) + cap_src_ids_shard = torch.cat(cap_src_ids, dim=0) + cap_freqs_cis_shard = torch.cat(cap_freqs_cis, dim=0) + cap_pad_mask_shard = torch.cat(cap_pad_mask, dim=0) + del cap_feats + + cap_shard = self.cap_embedder(cap_shard) + cap_shard[cap_pad_mask_shard] = self.cap_pad_token + for layer in self.context_refiner: + cap_shard = layer( + cap_shard, + cap_src_ids_shard, + cap_freqs_cis_shard, + cap_cu_seqlens, + cap_max_item_seqlen, + ) + cap_flatten = cap_shard + + # unified + def merge_interleave(l1, l2): + return list(itertools.chain(*zip(l1, l2))) + + unified = torch.cat( + merge_interleave(cap_flatten.split(cap_item_seqlens, dim=0), x_flatten.split(x_item_seqlens, dim=0)), dim=0 + ) + unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + assert len(unified) == sum(unified_item_seqlens) + unified_max_item_seqlen = max(unified_item_seqlens) + unified_cu_seqlens = F.pad( + torch.cumsum( + torch.tensor(unified_item_seqlens, dtype=torch.int32, device=device), dim=0, dtype=torch.int32 + ), + (1, 0), + ) + unified_src_ids = torch.cat(merge_interleave(cap_src_ids, x_src_ids)) + unified_freqs_cis = torch.cat(merge_interleave(cap_freqs_cis, x_freqs_cis)) + + unified_shard = unified + unified_src_ids_shard = unified_src_ids + unified_freqs_cis_shard = unified_freqs_cis + for layer in self.layers: + unified_shard = layer( + unified_shard, + unified_src_ids_shard, + unified_freqs_cis_shard, + unified_cu_seqlens, + unified_max_item_seqlen, + adaln_input, + ) + unified_shard = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( + unified_shard, unified_src_ids_shard, adaln_input + ) + unified = unified_shard.split(unified_item_seqlens, dim=0) + x = [unified[i][cap_item_seqlens[i] :] for i in range(bsz)] + assert all(len(x[i]) == x_item_seqlens[i] for i in range(bsz)) + + x = self.unpatchify(x, x_size, patch_size, f_patch_size) + + return x, {} + + def parameter_count(self) -> int: + total_params = 0 + + def _recursive_count_params(module): + nonlocal total_params + for param in module.parameters(recurse=False): + total_params += param.numel() + for submodule in module.children(): + _recursive_count_params(submodule) + + _recursive_count_params(self) + return total_params diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 69bb14b98edc..2754ffdc96d6 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -395,6 +395,7 @@ "WanVACEPipeline", "WanAnimatePipeline", ] + _import_structure["z_image"] = ["ZImagePipeline"] _import_structure["kandinsky5"] = ["Kandinsky5T2VPipeline"] _import_structure["skyreels_v2"] = [ "SkyReelsV2DiffusionForcingPipeline", @@ -819,6 +820,7 @@ WanVACEPipeline, WanVideoToVideoPipeline, ) + from .z_image import ZImagePipeline from .wuerstchen import ( WuerstchenCombinedPipeline, WuerstchenDecoderPipeline, diff --git a/src/diffusers/pipelines/z_image/__init__.py b/src/diffusers/pipelines/z_image/__init__.py new file mode 100644 index 000000000000..1f301648efd4 --- /dev/null +++ b/src/diffusers/pipelines/z_image/__init__.py @@ -0,0 +1,51 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa: F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_output"] = ["ZImagePipelineOutput"] + _import_structure["pipeline_z_image"] = ["ZImagePipeline"] + + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_output import ZImagePipelineOutput + from .pipeline_z_image import ZImagePipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) + diff --git a/src/diffusers/pipelines/z_image/pipeline_output.py b/src/diffusers/pipelines/z_image/pipeline_output.py new file mode 100644 index 000000000000..f2c3961088eb --- /dev/null +++ b/src/diffusers/pipelines/z_image/pipeline_output.py @@ -0,0 +1,35 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import List, Union + +import PIL.Image +from diffusers.utils import BaseOutput +import numpy as np + + +@dataclass +class ZImagePipelineOutput(BaseOutput): + """ + Output class for Z-Image pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] + diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py new file mode 100644 index 000000000000..9f1dcf9ac2ce --- /dev/null +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -0,0 +1,622 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from transformers import AutoTokenizer, PreTrainedModel + +from ...image_processor import VaeImageProcessor +from ...loaders import FromSingleFileMixin +from ...models.autoencoders import AutoencoderKL +from ...models.transformers import ZImageTransformer2DModel +from ...pipelines.pipeline_utils import DiffusionPipeline +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import USE_PEFT_BACKEND, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from .pipeline_output import ZImagePipelineOutput + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import ZImagePipeline + + >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. + >>> image = pipe(prompt, height=1024, width=1024, num_inference_steps=9, guidance_scale=0.0, generator=torch.Generator("cuda").manual_seed(42)).images[0] + >>> image.save("zimage.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class ZImagePipeline(DiffusionPipeline, FromSingleFileMixin): + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: PreTrainedModel, + tokenizer: AutoTokenizer, + transformer: ZImageTransformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer, + ) + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 + ) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[Union[str, List[str]]] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + max_sequence_length: int = 512, + lora_scale: Optional[float] = None, + ): + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt_embeds = self._encode_prompt( + prompt=prompt, + device=device, + dtype=dtype, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=prompt_embeds, + max_sequence_length=max_sequence_length, + ) + + if do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = ["" for _ in prompt] + else: + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + assert len(prompt) == len(negative_prompt) + negative_prompt_embeds = self._encode_prompt( + prompt=negative_prompt, + device=device, + dtype=dtype, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + ) + return prompt_embeds, negative_prompt_embeds + + def _encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + max_sequence_length: int = 512, + ) -> List[torch.FloatTensor]: + + assert num_images_per_prompt == 1 + device = device or self._execution_device + + if prompt_embeds is not None: + return prompt_embeds + + if isinstance(prompt, str): + prompt = [prompt] + + for i, prompt_item in enumerate(prompt): + messages = [ + {"role": "user", "content": prompt_item}, + ] + prompt_item = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + prompt[i] = prompt_item + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device) + prompt_masks = text_inputs.attention_mask.to(device).bool() + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_masks, + output_hidden_states=True, + ).hidden_states[-2] + + embeddings_list = [] + + for i in range(len(prompt_embeds)): + embeddings_list.append(prompt_embeds[i][prompt_masks[i]]) + + return embeddings_list + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @staticmethod + def value_from_time_aware_config(config, t): + if isinstance(config, (float, int, str)): + return config + elif isinstance(config, torch.Tensor): + assert config.numel() == 1 + return config.item() + elif isinstance(config, (tuple, list)): + assert isinstance(config[0], (float, int, str)) + assert all([isinstance(x, (tuple, list, str)) for x in config[1:]]) + result = config[0] + for thresh, val in config[1:]: + if t >= thresh: + result = val + else: + break + return result + else: + raise ValueError(f"invalid time-aware config {config} of type {type(config)}") + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + sigmas: Optional[List[float]] = None, + guidance_scale: float = 5.0, + cfg_normalization: bool = False, + cfg_truncation: float = 1.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + cfg_normalization (`bool`, *optional*, defaults to False): + Whether to apply configuration normalization. + cfg_truncation (`float`, *optional*, defaults to 1.0): + The truncation value for configuration. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a + plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: + [`~pipelines.z_image.ZImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is a list with the generated images. + """ + height = height or 1024 + width = width or 1024 + + assert self.dtype == torch.bfloat16 + dtype = self.dtype + device = self._execution_device + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + self._cfg_normalization = cfg_normalization + self._cfg_truncation = cfg_truncation + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + lora_scale = self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + dtype=dtype, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.in_channels + + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + torch.float32, + device, + generator, + latents, + ) + image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] / 2) + + # 5. Prepare timesteps + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + self.scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]) + timestep = (1000 - timestep) / 1000 + # Normalized time for time-aware config (0 at start, 1 at end) + t_norm = timestep[0].item() + + # Handle cfg truncation + current_guidance_scale = self.guidance_scale + if ( + self.do_classifier_free_guidance + and self._cfg_truncation is not None + and float(self._cfg_truncation) <= 1 + ): + current_guidance_scale = self.value_from_time_aware_config( + (self.guidance_scale, (self._cfg_truncation, 0.0)), t_norm + ) + + # Run CFG only if configured AND scale is non-zero + apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 + + if apply_cfg: + # Prepare inputs for CFG + latent_model_input = torch.cat([latents.to(dtype)] * 2) + prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds + timestep_model_input = torch.cat([timestep] * 2) + else: + latent_model_input = latents.to(dtype) + prompt_embeds_model_input = prompt_embeds + timestep_model_input = timestep + + latent_model_input = latent_model_input.unsqueeze(2) + latent_model_input_list = list(latent_model_input.unbind(dim=0)) + + model_out_list = self.transformer( + latent_model_input_list, + timestep_model_input, + prompt_embeds_model_input, + )[0] + + if apply_cfg: + # Perform CFG + pos_out = model_out_list[:batch_size] + neg_out = model_out_list[batch_size:] + + noise_pred = [] + for j in range(batch_size): + pos = pos_out[j].float() + neg = neg_out[j].float() + + pred = pos + current_guidance_scale * (pos - neg) + + # Renormalization + if self._cfg_normalization and float(self._cfg_normalization) > 0.0: + ori_pos_norm = torch.linalg.vector_norm(pos) + new_pos_norm = torch.linalg.vector_norm(pred) + max_new_norm = ori_pos_norm * float(self._cfg_normalization) + if new_pos_norm > max_new_norm: + pred = pred * (max_new_norm / new_pos_norm) + + noise_pred.append(pred) + + noise_pred = torch.stack(noise_pred, dim=0) + else: + noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) + + noise_pred = noise_pred.squeeze(2) + noise_pred = -noise_pred + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0] + assert latents.dtype == torch.float32 + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + latents = latents.to(dtype) + if output_type == "latent": + image = latents + + else: + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ZImagePipelineOutput(images=image) diff --git a/tests/pipelines/z_image/__init__.py b/tests/pipelines/z_image/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 From 3e74bb259a8122fedc525a528e7eb1a2b051e5a1 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Sun, 23 Nov 2025 20:06:32 +0000 Subject: [PATCH 02/31] Reformatting with make style, black & isort. --- .../transformers/transformer_z_image.py | 156 ++++++++++++------ src/diffusers/pipelines/__init__.py | 2 +- src/diffusers/pipelines/z_image/__init__.py | 1 - .../pipelines/z_image/pipeline_output.py | 4 +- .../pipelines/z_image/pipeline_z_image.py | 30 ++-- 5 files changed, 128 insertions(+), 65 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 45a609f3eeb7..3333e026eac3 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -16,10 +16,11 @@ import math from typing import List, Optional, Tuple -from einops import rearrange import torch import torch.nn as nn import torch.nn.functional as F +from einops import rearrange + try: from flash_attn import flash_attn_varlen_func @@ -33,10 +34,10 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin from ...utils.torch_utils import maybe_allow_in_graph -from ...models.attention_processor import Attention -from ...models.attention_dispatch import dispatch_attention_fn + ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 @@ -88,10 +89,10 @@ def forward(self, t): class ZSingleStreamAttnProcessor: """ - Processor for Z-Image single stream attention that adapts the existing Attention class - to match the behavior of the original Z-ImageAttention module. + Processor for Z-Image single stream attention that adapts the existing Attention class to match the behavior of the + original Z-ImageAttention module. """ - + _attention_backend = None _parallel_config = None @@ -107,24 +108,24 @@ def __call__( ) -> torch.Tensor: x_shard = hidden_states x_freqs_cis_shard = image_rotary_emb - + query = attn.to_q(x_shard) key = attn.to_k(x_shard) value = attn.to_v(x_shard) - + seqlen_shard = x_shard.shape[0] - + # Reshape to [seq_len, heads, head_dim] head_dim = query.shape[-1] // attn.heads query = query.view(seqlen_shard, attn.heads, head_dim) key = key.view(seqlen_shard, attn.heads, head_dim) - value = value.view(seqlen_shard, attn.heads, head_dim) + value = value.view(seqlen_shard, attn.heads, head_dim) # Apply Norms if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) - + # Apply RoPE def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: with torch.amp.autocast("cuda", enabled=False): @@ -136,17 +137,17 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso if x_freqs_cis_shard is not None: query = apply_rotary_emb(query, x_freqs_cis_shard) key = apply_rotary_emb(key, x_freqs_cis_shard) - + # Cast to correct dtype dtype = query.dtype query, key = query.to(dtype), key.to(dtype) - + # Flash Attention softmax_scale = math.sqrt(1 / head_dim) assert dtype in [torch.float16, torch.bfloat16] - + if x_cu_seqlens is None or x_max_item_seqlen is None: - raise ValueError("x_cu_seqlens and x_max_item_seqlen are required for ZSingleStreamAttnProcessor") + raise ValueError("x_cu_seqlens and x_max_item_seqlen are required for ZSingleStreamAttnProcessor") if flash_attn_varlen_func is not None: output = flash_attn_varlen_func( @@ -164,45 +165,50 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso output = output.flatten(-2) else: seqlens = (x_cu_seqlens[1:] - x_cu_seqlens[:-1]).cpu().tolist() - + q_split = torch.split(query, seqlens, dim=0) k_split = torch.split(key, seqlens, dim=0) v_split = torch.split(value, seqlens, dim=0) - + q_padded = torch.nn.utils.rnn.pad_sequence(q_split, batch_first=True) k_padded = torch.nn.utils.rnn.pad_sequence(k_split, batch_first=True) v_padded = torch.nn.utils.rnn.pad_sequence(v_split, batch_first=True) - + batch_size, max_seqlen, _, _ = q_padded.shape - + mask = torch.zeros((batch_size, max_seqlen), dtype=torch.bool, device=query.device) for i, l in enumerate(seqlens): mask[i, :l] = True - + attn_mask = torch.zeros((batch_size, 1, 1, max_seqlen), dtype=query.dtype, device=query.device) attn_mask.masked_fill_(~mask[:, None, None, :], torch.finfo(query.dtype).min) - + q_padded = q_padded.transpose(1, 2) k_padded = k_padded.transpose(1, 2) v_padded = v_padded.transpose(1, 2) - + output = F.scaled_dot_product_attention( - q_padded, k_padded, v_padded, attn_mask=attn_mask, dropout_p=0.0, scale=softmax_scale + q_padded, + k_padded, + v_padded, + attn_mask=attn_mask, + dropout_p=0.0, + scale=softmax_scale, ) - + output = output.transpose(1, 2) - + out_list = [] for i, l in enumerate(seqlens): out_list.append(output[i, :l]) - + output = torch.cat(out_list, dim=0) output = output.flatten(-2) output = attn.to_out[0](output) - if len(attn.to_out) > 1: # dropout - output = attn.to_out[1](output) - + if len(attn.to_out) > 1: # dropout + output = attn.to_out[1](output) + return output @@ -226,12 +232,19 @@ def forward(self, x): @maybe_allow_in_graph 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, modulation=True + self, + layer_id: int, + dim: int, + n_heads: int, + n_kv_heads: int, + norm_eps: float, + qk_norm: bool, + modulation=True, ): super().__init__() self.dim = dim self.head_dim = dim // n_heads - + # Refactored to use diffusers Attention with custom processor # Original Z-Image params: dim, n_heads, n_kv_heads, qk_norm self.attention = Attention( @@ -244,7 +257,7 @@ def __init__( bias=False, processor=ZSingleStreamAttnProcessor(), ) - + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) self.layer_id = layer_id @@ -284,7 +297,12 @@ def forward( x_src_ids_shard = None x_shard = self.attn_forward( - x_shard, x_freqs_cis_shard, x_cu_seqlens, x_max_item_seqlen, scale_gate_msa, x_src_ids_shard + x_shard, + x_freqs_cis_shard, + x_cu_seqlens, + x_max_item_seqlen, + scale_gate_msa, + x_src_ids_shard, ) x_shard = self.ffn_forward(x_shard, scale_gate_mlp, x_src_ids_shard) @@ -303,22 +321,22 @@ def attn_forward( if self.modulation: assert scale_gate is not None and x_src_ids_shard is not None scale_msa, gate_msa = scale_gate - + # Pass extra args needed for ZSingleStreamAttnProcessor attn_out = self.attention( self.attention_norm1(x_shard) * scale_msa[x_src_ids_shard], image_rotary_emb=x_freqs_cis_shard, x_cu_seqlens=x_cu_seqlens, - x_max_item_seqlen=x_max_item_seqlen + x_max_item_seqlen=x_max_item_seqlen, ) - + x_shard = x_shard + gate_msa[x_src_ids_shard] * self.attention_norm2(attn_out) else: attn_out = self.attention( self.attention_norm1(x_shard), image_rotary_emb=x_freqs_cis_shard, x_cu_seqlens=x_cu_seqlens, - x_max_item_seqlen=x_max_item_seqlen + x_max_item_seqlen=x_max_item_seqlen, ) x_shard = x_shard + self.attention_norm2(attn_out) return x_shard @@ -371,7 +389,10 @@ def forward(self, x_shard, x_src_ids_shard, c): class RopeEmbedder: def __init__( - self, theta: float = 256.0, axes_dims: List[int] = (16, 56, 56), axes_lens: List[int] = (64, 128, 128) + self, + theta: float = 256.0, + axes_dims: List[int] = (16, 56, 56), + axes_lens: List[int] = (64, 128, 128), ): self.theta = theta self.axes_dims = axes_dims @@ -458,13 +479,29 @@ def __init__( self.all_final_layer = nn.ModuleDict(all_final_layer) self.noise_refiner = nn.ModuleList( [ - ZImageTransformerBlock(1000 + layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True) + ZImageTransformerBlock( + 1000 + layer_id, + dim, + n_heads, + n_kv_heads, + norm_eps, + qk_norm, + modulation=True, + ) for layer_id in range(n_refiner_layers) ] ) self.context_refiner = nn.ModuleList( [ - ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=False) + ZImageTransformerBlock( + layer_id, + dim, + n_heads, + n_kv_heads, + norm_eps, + qk_norm, + modulation=False, + ) for layer_id in range(n_refiner_layers) ] ) @@ -524,8 +561,6 @@ def patchify_and_embed( patch_size: int, f_patch_size: int, ): - - bsz = len(all_image) pH = pW = patch_size pF = f_patch_size device = all_image[0].device @@ -560,7 +595,10 @@ def patchify_and_embed( ) ) # padded feature - cap_padded_feat = torch.cat([all_cap_feats[i], all_cap_feats[i][-1:].repeat(cap_padding_len, 1)], dim=0) + cap_padded_feat = torch.cat( + [all_cap_feats[i], all_cap_feats[i][-1:].repeat(cap_padding_len, 1)], + dim=0, + ) all_cap_feats_out.append(cap_padded_feat) ### Process Image @@ -623,7 +661,6 @@ def forward( patch_size=2, f_patch_size=1, ): - assert patch_size in self.all_patch_size assert f_patch_size in self.all_f_patch_size @@ -649,7 +686,11 @@ def forward( assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) x_max_item_seqlen = max(x_item_seqlens) x_cu_seqlens = F.pad( - torch.cumsum(torch.tensor(x_item_seqlens, dtype=torch.int32, device=device), dim=0, dtype=torch.int32), + torch.cumsum( + torch.tensor(x_item_seqlens, dtype=torch.int32, device=device), + dim=0, + dtype=torch.int32, + ), (1, 0), ) x_src_ids = [ @@ -666,7 +707,14 @@ def forward( x_shard = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_shard) x_shard[x_pad_mask_shard] = self.x_pad_token for layer in self.noise_refiner: - x_shard = layer(x_shard, x_src_ids_shard, x_freqs_cis_shard, x_cu_seqlens, x_max_item_seqlen, adaln_input) + x_shard = layer( + x_shard, + x_src_ids_shard, + x_freqs_cis_shard, + x_cu_seqlens, + x_max_item_seqlen, + adaln_input, + ) x_flatten = x_shard # cap embed & refine @@ -674,7 +722,11 @@ def forward( assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) cap_max_item_seqlen = max(cap_item_seqlens) cap_cu_seqlens = F.pad( - torch.cumsum(torch.tensor(cap_item_seqlens, dtype=torch.int32, device=device), dim=0, dtype=torch.int32), + torch.cumsum( + torch.tensor(cap_item_seqlens, dtype=torch.int32, device=device), + dim=0, + dtype=torch.int32, + ), (1, 0), ) cap_src_ids = [ @@ -705,14 +757,20 @@ def merge_interleave(l1, l2): return list(itertools.chain(*zip(l1, l2))) unified = torch.cat( - merge_interleave(cap_flatten.split(cap_item_seqlens, dim=0), x_flatten.split(x_item_seqlens, dim=0)), dim=0 + merge_interleave( + cap_flatten.split(cap_item_seqlens, dim=0), + x_flatten.split(x_item_seqlens, dim=0), + ), + dim=0, ) unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] assert len(unified) == sum(unified_item_seqlens) unified_max_item_seqlen = max(unified_item_seqlens) unified_cu_seqlens = F.pad( torch.cumsum( - torch.tensor(unified_item_seqlens, dtype=torch.int32, device=device), dim=0, dtype=torch.int32 + torch.tensor(unified_item_seqlens, dtype=torch.int32, device=device), + dim=0, + dtype=torch.int32, ), (1, 0), ) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 2754ffdc96d6..4ffe85da3c9e 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -820,12 +820,12 @@ WanVACEPipeline, WanVideoToVideoPipeline, ) - from .z_image import ZImagePipeline from .wuerstchen import ( WuerstchenCombinedPipeline, WuerstchenDecoderPipeline, WuerstchenPriorPipeline, ) + from .z_image import ZImagePipeline try: if not is_onnx_available(): diff --git a/src/diffusers/pipelines/z_image/__init__.py b/src/diffusers/pipelines/z_image/__init__.py index 1f301648efd4..f95b3e5a0bed 100644 --- a/src/diffusers/pipelines/z_image/__init__.py +++ b/src/diffusers/pipelines/z_image/__init__.py @@ -48,4 +48,3 @@ for name, value in _dummy_objects.items(): setattr(sys.modules[__name__], name, value) - diff --git a/src/diffusers/pipelines/z_image/pipeline_output.py b/src/diffusers/pipelines/z_image/pipeline_output.py index f2c3961088eb..69a320fc036a 100644 --- a/src/diffusers/pipelines/z_image/pipeline_output.py +++ b/src/diffusers/pipelines/z_image/pipeline_output.py @@ -15,9 +15,10 @@ from dataclasses import dataclass from typing import List, Union +import numpy as np import PIL.Image + from diffusers.utils import BaseOutput -import numpy as np @dataclass @@ -32,4 +33,3 @@ class ZImagePipelineOutput(BaseOutput): """ images: Union[List[PIL.Image.Image], np.ndarray] - diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 9f1dcf9ac2ce..dcf4b916a403 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -import math from typing import Any, Callable, Dict, List, Optional, Union import torch @@ -25,10 +24,11 @@ from ...models.transformers import ZImageTransformer2DModel from ...pipelines.pipeline_utils import DiffusionPipeline from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import USE_PEFT_BACKEND, logging, replace_example_docstring +from ...utils import logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from .pipeline_output import ZImagePipelineOutput + logger = logging.get_logger(__name__) # pylint: disable=invalid-name EXAMPLE_DOC_STRING = """ @@ -42,7 +42,14 @@ >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" >>> # Depending on the variant being used, the pipeline call will slightly vary. >>> # Refer to the pipeline documentation for more details. - >>> image = pipe(prompt, height=1024, width=1024, num_inference_steps=9, guidance_scale=0.0, generator=torch.Generator("cuda").manual_seed(42)).images[0] + >>> image = pipe( + ... prompt, + ... height=1024, + ... width=1024, + ... num_inference_steps=9, + ... guidance_scale=0.0, + ... generator=torch.Generator("cuda").manual_seed(42), + ... ).images[0] >>> image.save("zimage.png") ``` """ @@ -162,7 +169,6 @@ def encode_prompt( max_sequence_length: int = 512, lora_scale: Optional[float] = None, ): - prompt = [prompt] if isinstance(prompt, str) else prompt prompt_embeds = self._encode_prompt( prompt=prompt, @@ -198,7 +204,6 @@ def _encode_prompt( prompt_embeds: Optional[List[torch.FloatTensor]] = None, max_sequence_length: int = 512, ) -> List[torch.FloatTensor]: - assert num_images_per_prompt == 1 device = device or self._execution_device @@ -326,7 +331,6 @@ def value_from_time_aware_config(config, t): return config.item() elif isinstance(config, (tuple, list)): assert isinstance(config[0], (float, int, str)) - assert all([isinstance(x, (tuple, list, str)) for x in config[1:]]) result = config[0] for thresh, val in config[1:]: if t >= thresh: @@ -414,8 +418,8 @@ def __call__( The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a - plain tuple. + Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain + tuple. joint_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in @@ -435,9 +439,9 @@ def __call__( Examples: Returns: - [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: - [`~pipelines.z_image.ZImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is a list with the generated images. + [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. """ height = height or 1024 width = width or 1024 @@ -459,7 +463,9 @@ def __call__( else: batch_size = prompt_embeds.shape[0] - lora_scale = self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) ( prompt_embeds, negative_prompt_embeds, From a4b89a08e9d215bf40bd3491a19e7ebba54f8f93 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Mon, 24 Nov 2025 08:12:52 +0000 Subject: [PATCH 03/31] Remove init, Modify import utils, Merge forward in transformers block, Remove once func in pipeline. --- .../transformers/transformer_z_image.py | 232 +++++++----------- .../pipelines/z_image/pipeline_z_image.py | 29 --- src/diffusers/utils/import_utils.py | 5 + 3 files changed, 94 insertions(+), 172 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 3333e026eac3..1db167024bf4 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -21,23 +21,25 @@ import torch.nn.functional as F from einops import rearrange +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...models.attention_processor import Attention +from ...models.modeling_utils import ModelMixin +from ...utils.import_utils import is_apex_available, is_flash_attn_available +from ...utils.torch_utils import maybe_allow_in_graph + -try: +if is_flash_attn_available(): from flash_attn import flash_attn_varlen_func -except ImportError: +else: flash_attn_varlen_func = None -try: +if is_apex_available(): + # Here needs apex with "APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation ." from apex.normalization import FusedRMSNorm as RMSNorm -except ImportError: +else: from torch.nn import RMSNorm -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...models.attention_processor import Attention -from ...models.modeling_utils import ModelMixin -from ...utils.torch_utils import maybe_allow_in_graph - ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 @@ -61,10 +63,6 @@ def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): bias=True, ), ) - nn.init.normal_(self.mlp[0].weight, std=0.02) - nn.init.zeros_(self.mlp[0].bias) - nn.init.normal_(self.mlp[2].weight, std=0.02) - nn.init.zeros_(self.mlp[2].bias) self.frequency_embedding_size = frequency_embedding_size @@ -106,20 +104,20 @@ def __call__( x_cu_seqlens: Optional[torch.Tensor] = None, x_max_item_seqlen: Optional[int] = None, ) -> torch.Tensor: - x_shard = hidden_states - x_freqs_cis_shard = image_rotary_emb + x = hidden_states + x_freqs_cis = image_rotary_emb - query = attn.to_q(x_shard) - key = attn.to_k(x_shard) - value = attn.to_v(x_shard) + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) - seqlen_shard = x_shard.shape[0] + seqlen = x.shape[0] # Reshape to [seq_len, heads, head_dim] head_dim = query.shape[-1] // attn.heads - query = query.view(seqlen_shard, attn.heads, head_dim) - key = key.view(seqlen_shard, attn.heads, head_dim) - value = value.view(seqlen_shard, attn.heads, head_dim) + query = query.view(seqlen, attn.heads, head_dim) + key = key.view(seqlen, attn.heads, head_dim) + value = value.view(seqlen, attn.heads, head_dim) # Apply Norms if attn.norm_q is not None: query = attn.norm_q(query) @@ -134,9 +132,9 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso x_out = torch.view_as_real(x * freqs_cis).flatten(2) return x_out.type_as(x_in) - if x_freqs_cis_shard is not None: - query = apply_rotary_emb(query, x_freqs_cis_shard) - key = apply_rotary_emb(key, x_freqs_cis_shard) + if x_freqs_cis is not None: + query = apply_rotary_emb(query, x_freqs_cis) + key = apply_rotary_emb(key, x_freqs_cis) # Cast to correct dtype dtype = query.dtype @@ -277,9 +275,9 @@ def __init__( def forward( self, - x_shard: torch.Tensor, - x_src_ids_shard: torch.Tensor, - x_freqs_cis_shard: torch.Tensor, + x: torch.Tensor, + x_src_ids: torch.Tensor, + x_freqs_cis: torch.Tensor, x_cu_seqlens: torch.Tensor, x_max_item_seqlen: int, adaln_input: Optional[torch.Tensor] = None, @@ -289,80 +287,40 @@ def forward( scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).chunk(4, dim=1) gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp - scale_gate_msa = (scale_msa, gate_msa) - scale_gate_mlp = (scale_mlp, gate_mlp) - else: - scale_gate_msa = None - scale_gate_mlp = None - x_src_ids_shard = None - - x_shard = self.attn_forward( - x_shard, - x_freqs_cis_shard, - x_cu_seqlens, - x_max_item_seqlen, - scale_gate_msa, - x_src_ids_shard, - ) - x_shard = self.ffn_forward(x_shard, scale_gate_mlp, x_src_ids_shard) - - return x_shard - - def attn_forward( - self, - x_shard, - x_freqs_cis_shard, - x_cu_seqlens, - x_max_item_seqlen, - scale_gate: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - x_src_ids_shard: Optional[torch.Tensor] = None, - ): - if self.modulation: - assert scale_gate is not None and x_src_ids_shard is not None - scale_msa, gate_msa = scale_gate - - # Pass extra args needed for ZSingleStreamAttnProcessor + # Attention block attn_out = self.attention( - self.attention_norm1(x_shard) * scale_msa[x_src_ids_shard], - image_rotary_emb=x_freqs_cis_shard, + self.attention_norm1(x) * scale_msa[x_src_ids], + image_rotary_emb=x_freqs_cis, x_cu_seqlens=x_cu_seqlens, x_max_item_seqlen=x_max_item_seqlen, ) + x = x + gate_msa[x_src_ids] * self.attention_norm2(attn_out) - x_shard = x_shard + gate_msa[x_src_ids_shard] * self.attention_norm2(attn_out) + # FFN block + x = x + gate_mlp[x_src_ids] * self.ffn_norm2( + self.feed_forward( + self.ffn_norm1(x) * scale_mlp[x_src_ids], + ) + ) else: + # Attention block attn_out = self.attention( - self.attention_norm1(x_shard), - image_rotary_emb=x_freqs_cis_shard, + self.attention_norm1(x), + image_rotary_emb=x_freqs_cis, x_cu_seqlens=x_cu_seqlens, x_max_item_seqlen=x_max_item_seqlen, ) - x_shard = x_shard + self.attention_norm2(attn_out) - return x_shard + x = x + self.attention_norm2(attn_out) - def ffn_forward( - self, - x_shard, - scale_gate: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - x_src_ids_shard: Optional[torch.Tensor] = None, - ): - if self.modulation: - assert scale_gate is not None and x_src_ids_shard is not None - scale_mlp, gate_mlp = scale_gate - x_shard = x_shard + gate_mlp[x_src_ids_shard] * self.ffn_norm2( + # FFN block + x = x + self.ffn_norm2( self.feed_forward( - self.ffn_norm1(x_shard) * scale_mlp[x_src_ids_shard], + self.ffn_norm1(x), ) ) - else: - x_shard = x_shard + self.ffn_norm2( - self.feed_forward( - self.ffn_norm1(x_shard), - ) - ) - return x_shard + return x class FinalLayer(nn.Module): @@ -380,11 +338,11 @@ def __init__(self, hidden_size, out_channels): nn.init.zeros_(self.adaLN_modulation[1].weight) nn.init.zeros_(self.adaLN_modulation[1].bias) - def forward(self, x_shard, x_src_ids_shard, c): + def forward(self, x, x_src_ids, c): scale = 1.0 + self.adaLN_modulation(c) - x_shard = self.norm_final(x_shard) * scale[x_src_ids_shard] - x_shard = self.linear(x_shard) - return x_shard + x = self.norm_final(x) * scale[x_src_ids] + x = self.linear(x) + return x class RopeEmbedder: @@ -468,8 +426,6 @@ def __init__( all_final_layer = {} for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)): x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True) - nn.init.xavier_uniform_(x_embedder.weight) - nn.init.constant_(x_embedder.bias, 0.0) all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels) @@ -698,24 +654,23 @@ def forward( ] x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) - x_shard = torch.cat(x, dim=0) - x_src_ids_shard = torch.cat(x_src_ids, dim=0) - x_freqs_cis_shard = torch.cat(x_freqs_cis, dim=0) - x_pad_mask_shard = torch.cat(x_pad_mask, dim=0) - del x + x = torch.cat(x, dim=0) + x_src_ids = torch.cat(x_src_ids, dim=0) + x_freqs_cis = torch.cat(x_freqs_cis, dim=0) + x_pad_mask = torch.cat(x_pad_mask, dim=0) - x_shard = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_shard) - x_shard[x_pad_mask_shard] = self.x_pad_token + x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) + x[x_pad_mask] = self.x_pad_token for layer in self.noise_refiner: - x_shard = layer( - x_shard, - x_src_ids_shard, - x_freqs_cis_shard, + x = layer( + x, + x_src_ids, + x_freqs_cis, x_cu_seqlens, x_max_item_seqlen, adaln_input, ) - x_flatten = x_shard + x_flatten = x # cap embed & refine cap_item_seqlens = [len(_) for _ in cap_feats] @@ -734,23 +689,23 @@ def forward( ] cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) - cap_shard = torch.cat(cap_feats, dim=0) - cap_src_ids_shard = torch.cat(cap_src_ids, dim=0) - cap_freqs_cis_shard = torch.cat(cap_freqs_cis, dim=0) - cap_pad_mask_shard = torch.cat(cap_pad_mask, dim=0) + cap = torch.cat(cap_feats, dim=0) + cap_src_ids = torch.cat(cap_src_ids, dim=0) + cap_freqs_cis = torch.cat(cap_freqs_cis, dim=0) + cap_pad_mask = torch.cat(cap_pad_mask, dim=0) del cap_feats - cap_shard = self.cap_embedder(cap_shard) - cap_shard[cap_pad_mask_shard] = self.cap_pad_token + cap = self.cap_embedder(cap) + cap[cap_pad_mask] = self.cap_pad_token for layer in self.context_refiner: - cap_shard = layer( - cap_shard, - cap_src_ids_shard, - cap_freqs_cis_shard, + cap = layer( + cap, + cap_src_ids, + cap_freqs_cis, cap_cu_seqlens, cap_max_item_seqlen, ) - cap_flatten = cap_shard + cap_flatten = cap # unified def merge_interleave(l1, l2): @@ -774,41 +729,32 @@ def merge_interleave(l1, l2): ), (1, 0), ) - unified_src_ids = torch.cat(merge_interleave(cap_src_ids, x_src_ids)) - unified_freqs_cis = torch.cat(merge_interleave(cap_freqs_cis, x_freqs_cis)) - - unified_shard = unified - unified_src_ids_shard = unified_src_ids - unified_freqs_cis_shard = unified_freqs_cis + unified_src_ids = torch.cat( + merge_interleave( + cap_src_ids.split(cap_item_seqlens, dim=0), + x_src_ids.split(x_item_seqlens, dim=0), + ) + ) + unified_freqs_cis = torch.cat( + merge_interleave( + cap_freqs_cis.split(cap_item_seqlens, dim=0), + x_freqs_cis.split(x_item_seqlens, dim=0), + ) + ) for layer in self.layers: - unified_shard = layer( - unified_shard, - unified_src_ids_shard, - unified_freqs_cis_shard, + unified = layer( + unified, + unified_src_ids, + unified_freqs_cis, unified_cu_seqlens, unified_max_item_seqlen, adaln_input, ) - unified_shard = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( - unified_shard, unified_src_ids_shard, adaln_input - ) - unified = unified_shard.split(unified_item_seqlens, dim=0) + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, unified_src_ids, adaln_input) + unified = unified.split(unified_item_seqlens, dim=0) x = [unified[i][cap_item_seqlens[i] :] for i in range(bsz)] assert all(len(x[i]) == x_item_seqlens[i] for i in range(bsz)) x = self.unpatchify(x, x_size, patch_size, f_patch_size) return x, {} - - def parameter_count(self) -> int: - total_params = 0 - - def _recursive_count_params(module): - nonlocal total_params - for param in module.parameters(recurse=False): - total_params += param.numel() - for submodule in module.children(): - _recursive_count_params(submodule) - - _recursive_count_params(self) - return total_params diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index dcf4b916a403..5abb73550c25 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -249,35 +249,6 @@ def _encode_prompt( return embeddings_list - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - def prepare_latents( self, batch_size, diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 57b0a337922a..520a6c1510d9 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -230,6 +230,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b _aiter_available, _aiter_version = _is_package_available("aiter") _kornia_available, _kornia_version = _is_package_available("kornia") _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True) +_apex_available, _apex_version = _is_package_available("apex") def is_torch_available(): @@ -420,6 +421,10 @@ def is_kornia_available(): return _kornia_available +def is_apex_available(): + return _apex_available + + # docstyle-ignore FLAX_IMPORT_ERROR = """ {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the From 7df350d0f70879cce5f31d71140841116851e733 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 17:13:50 +0800 Subject: [PATCH 04/31] modified main model forward, freqs_cis left --- .../transformers/transformer_z_image.py | 196 +++++++++--------- 1 file changed, 95 insertions(+), 101 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 3333e026eac3..a7edd8b04f31 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -27,6 +27,7 @@ except ImportError: flash_attn_varlen_func = None +# todo see how other teams do this try: from apex.normalization import FusedRMSNorm as RMSNorm except ImportError: @@ -61,10 +62,6 @@ def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): bias=True, ), ) - nn.init.normal_(self.mlp[0].weight, std=0.02) - nn.init.zeros_(self.mlp[0].bias) - nn.init.normal_(self.mlp[2].weight, std=0.02) - nn.init.zeros_(self.mlp[2].bias) self.frequency_embedding_size = frequency_embedding_size @@ -573,9 +570,9 @@ def patchify_and_embed( all_cap_pad_mask = [] all_cap_feats_out = [] - for i, image in enumerate(all_image): - ### LLM Text Encoder - cap_ori_len = len(all_cap_feats[i]) + for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): + ### Process Caption + cap_ori_len = len(cap_feat) cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF # padded position ids cap_padded_pos_ids = self.create_coordinate_grid( @@ -596,7 +593,7 @@ def patchify_and_embed( ) # padded feature cap_padded_feat = torch.cat( - [all_cap_feats[i], all_cap_feats[i][-1:].repeat(cap_padding_len, 1)], + [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0, ) all_cap_feats_out.append(cap_padded_feat) @@ -677,126 +674,123 @@ def forward( x_size, x_pos_ids, cap_pos_ids, - x_pad_mask, - cap_pad_mask, + x_inner_pad_mask, + cap_inner_pad_mask, ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) # x embed & refine x_item_seqlens = [len(_) for _ in x] assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) x_max_item_seqlen = max(x_item_seqlens) - x_cu_seqlens = F.pad( - torch.cumsum( - torch.tensor(x_item_seqlens, dtype=torch.int32, device=device), - dim=0, - dtype=torch.int32, - ), - (1, 0), + + x = torch.cat(x, dim=0) + x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) + x[torch.cat(x_inner_pad_mask)] = self.x_pad_token + x = x.split(x_item_seqlens, dim=0) + x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) # todo + + pad_tensor = torch.zeros( + (1, self.dim), + dtype=x[0].dtype, + device=device, ) - x_src_ids = [ - torch.full((count,), i, dtype=torch.int32, device=device) for i, count in enumerate(x_item_seqlens) - ] - x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) - - x_shard = torch.cat(x, dim=0) - x_src_ids_shard = torch.cat(x_src_ids, dim=0) - x_freqs_cis_shard = torch.cat(x_freqs_cis, dim=0) - x_pad_mask_shard = torch.cat(x_pad_mask, dim=0) - del x - - x_shard = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_shard) - x_shard[x_pad_mask_shard] = self.x_pad_token + x_pad_mask = torch.zeros( + (bsz, x_max_item_seqlen), + dtype=torch.bool, + device=device + ) + for i, item in enumerate(x): + seq_len = x_item_seqlens[i] + x[i] = torch.cat([item, pad_tensor.repeat(x_max_item_seqlen - seq_len, 1)]) + x_pad_mask[i, seq_len:] = 1 + x = torch.stack(x) + for layer in self.noise_refiner: - x_shard = layer( - x_shard, - x_src_ids_shard, - x_freqs_cis_shard, - x_cu_seqlens, - x_max_item_seqlen, + x = layer( + x, + x_pad_mask, + x_freqs_cis, adaln_input, - ) - x_flatten = x_shard + ) # todo # cap embed & refine cap_item_seqlens = [len(_) for _ in cap_feats] assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) cap_max_item_seqlen = max(cap_item_seqlens) - cap_cu_seqlens = F.pad( - torch.cumsum( - torch.tensor(cap_item_seqlens, dtype=torch.int32, device=device), - dim=0, - dtype=torch.int32, - ), - (1, 0), + + cap_feats = torch.cat(cap_feats, dim=0) + cap_feats = self.cap_embedder(cap_feats) + cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token + cap_feats = cap_feats.split(cap_item_seqlens, dim=0) + cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) # todo + + pad_tensor = torch.zeros( + (1, self.dim), + dtype=x[0].dtype, + device=device, ) - cap_src_ids = [ - torch.full((count,), i, dtype=torch.int32, device=device) for i, count in enumerate(cap_item_seqlens) - ] - cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) - - cap_shard = torch.cat(cap_feats, dim=0) - cap_src_ids_shard = torch.cat(cap_src_ids, dim=0) - cap_freqs_cis_shard = torch.cat(cap_freqs_cis, dim=0) - cap_pad_mask_shard = torch.cat(cap_pad_mask, dim=0) - del cap_feats - - cap_shard = self.cap_embedder(cap_shard) - cap_shard[cap_pad_mask_shard] = self.cap_pad_token + cap_pad_mask = torch.zeros( + (bsz, cap_max_item_seqlen), + dtype=torch.bool, + device=device + ) + for i, item in enumerate(cap_feats): + seq_len = cap_item_seqlens[i] + cap_feats[i] = torch.cat([item, pad_tensor.repeat(cap_max_item_seqlen - seq_len, 1)]) + cap_pad_mask[i, seq_len:] = 1 + cap_feats = torch.stack(cap_feats) for layer in self.context_refiner: - cap_shard = layer( - cap_shard, - cap_src_ids_shard, - cap_freqs_cis_shard, - cap_cu_seqlens, - cap_max_item_seqlen, + cap_feats = layer( + cap_feats, + cap_pad_mask, + cap_freqs_cis, ) - cap_flatten = cap_shard - - # unified - def merge_interleave(l1, l2): - return list(itertools.chain(*zip(l1, l2))) - unified = torch.cat( - merge_interleave( - cap_flatten.split(cap_item_seqlens, dim=0), - x_flatten.split(x_item_seqlens, dim=0), - ), - dim=0, - ) + # unified todo unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] - assert len(unified) == sum(unified_item_seqlens) unified_max_item_seqlen = max(unified_item_seqlens) - unified_cu_seqlens = F.pad( - torch.cumsum( - torch.tensor(unified_item_seqlens, dtype=torch.int32, device=device), - dim=0, - dtype=torch.int32, - ), - (1, 0), + + pad_tensor = torch.zeros( + (1, self.dim), + dtype=x[0].dtype, + device=device, ) - unified_src_ids = torch.cat(merge_interleave(cap_src_ids, x_src_ids)) - unified_freqs_cis = torch.cat(merge_interleave(cap_freqs_cis, x_freqs_cis)) + unified_pad_mask = torch.zeros( + (bsz, unified_max_item_seqlen), + dtype=torch.bool, + device=device + ) + + unified = [] + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified.append( + torch.cat( + [ + x[i][:x_item_seqlens[i]], + cap_feats[i][:cap_item_seqlens[i]], + pad_tensor.repeat(unified_max_item_seqlen - x_len - cap_len, 1) + ] + ) + ) + unified_pad_mask[i, x_len + cap_len:] = 1 + + unified_freqs_cis = torch.cat(merge_interleave(cap_freqs_cis, x_freqs_cis)) # todo - unified_shard = unified - unified_src_ids_shard = unified_src_ids - unified_freqs_cis_shard = unified_freqs_cis for layer in self.layers: unified_shard = layer( - unified_shard, - unified_src_ids_shard, - unified_freqs_cis_shard, - unified_cu_seqlens, - unified_max_item_seqlen, + unified, + unified_pad_mask, + unified_freqs_cis, adaln_input, ) - unified_shard = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( - unified_shard, unified_src_ids_shard, adaln_input - ) - unified = unified_shard.split(unified_item_seqlens, dim=0) - x = [unified[i][cap_item_seqlens[i] :] for i in range(bsz)] - assert all(len(x[i]) == x_item_seqlens[i] for i in range(bsz)) - x = self.unpatchify(x, x_size, patch_size, f_patch_size) + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( + unified, adaln_input # todo + ) + unified = unified.split(unified_item_seqlens, dim=0) + x = self.unpatchify(unified, x_size, patch_size, f_patch_size) return x, {} From aae03cf6a6f6d1562815f98b50521ec2f7d7f07e Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 18:37:36 +0800 Subject: [PATCH 05/31] refactored to add B dim --- .../transformers/transformer_z_image.py | 254 +++++++----------- 1 file changed, 99 insertions(+), 155 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index a6492c81df7d..62938270bd89 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -23,6 +23,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ..attention_dispatch import dispatch_attention_fn from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin from ...utils.import_utils import is_apex_available, is_flash_attn_available @@ -34,11 +35,7 @@ else: flash_attn_varlen_func = None -if is_apex_available(): - # Here needs apex with "APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation ." - from apex.normalization import FusedRMSNorm as RMSNorm -else: - from torch.nn import RMSNorm +from diffusers.models.normalization import RMSNorm ADALN_EMBED_DIM = 256 @@ -98,26 +95,18 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - x_cu_seqlens: Optional[torch.Tensor] = None, - x_max_item_seqlen: Optional[int] = None, + freqs_cis: Optional[torch.Tensor] = None, ) -> torch.Tensor: - x = hidden_states - x_freqs_cis = image_rotary_emb - query = attn.to_q(x) - key = attn.to_k(x) - value = attn.to_v(x) + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) - seqlen = x.shape[0] + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) - # Reshape to [seq_len, heads, head_dim] - head_dim = query.shape[-1] // attn.heads - query = query.view(seqlen, attn.heads, head_dim) - key = key.view(seqlen, attn.heads, head_dim) - value = value.view(seqlen, attn.heads, head_dim) # Apply Norms if attn.norm_q is not None: query = attn.norm_q(query) @@ -128,82 +117,35 @@ def __call__( 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)) - freqs_cis = freqs_cis.unsqueeze(1) - x_out = torch.view_as_real(x * freqs_cis).flatten(2) - return x_out.type_as(x_in) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x * freqs_cis).flatten(3) + return x_out.type_as(x_in) # todo - if x_freqs_cis is not None: - query = apply_rotary_emb(query, x_freqs_cis) - key = apply_rotary_emb(key, x_freqs_cis) + if freqs_cis is not None: + query = apply_rotary_emb(query, freqs_cis) + key = apply_rotary_emb(key, freqs_cis) # Cast to correct dtype dtype = query.dtype query, key = query.to(dtype), key.to(dtype) - # Flash Attention - softmax_scale = math.sqrt(1 / head_dim) - assert dtype in [torch.float16, torch.bfloat16] - - if x_cu_seqlens is None or x_max_item_seqlen is None: - raise ValueError("x_cu_seqlens and x_max_item_seqlen are required for ZSingleStreamAttnProcessor") - - if flash_attn_varlen_func is not None: - output = flash_attn_varlen_func( - query, - key, - value, - cu_seqlens_q=x_cu_seqlens, - cu_seqlens_k=x_cu_seqlens, - max_seqlen_q=x_max_item_seqlen, - max_seqlen_k=x_max_item_seqlen, - dropout_p=0.0, - causal=False, - softmax_scale=softmax_scale, - ) - output = output.flatten(-2) - else: - seqlens = (x_cu_seqlens[1:] - x_cu_seqlens[:-1]).cpu().tolist() - - q_split = torch.split(query, seqlens, dim=0) - k_split = torch.split(key, seqlens, dim=0) - v_split = torch.split(value, seqlens, dim=0) - - q_padded = torch.nn.utils.rnn.pad_sequence(q_split, batch_first=True) - k_padded = torch.nn.utils.rnn.pad_sequence(k_split, batch_first=True) - v_padded = torch.nn.utils.rnn.pad_sequence(v_split, batch_first=True) - - batch_size, max_seqlen, _, _ = q_padded.shape - - mask = torch.zeros((batch_size, max_seqlen), dtype=torch.bool, device=query.device) - for i, l in enumerate(seqlens): - mask[i, :l] = True - - attn_mask = torch.zeros((batch_size, 1, 1, max_seqlen), dtype=query.dtype, device=query.device) - attn_mask.masked_fill_(~mask[:, None, None, :], torch.finfo(query.dtype).min) - - q_padded = q_padded.transpose(1, 2) - k_padded = k_padded.transpose(1, 2) - v_padded = v_padded.transpose(1, 2) - - output = F.scaled_dot_product_attention( - q_padded, - k_padded, - v_padded, - attn_mask=attn_mask, - dropout_p=0.0, - scale=softmax_scale, - ) - - output = output.transpose(1, 2) - - out_list = [] - for i, l in enumerate(seqlens): - out_list.append(output[i, :l]) + # Compute joint attention + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) - output = torch.cat(out_list, dim=0) - output = output.flatten(-2) + # Reshape back + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(dtype) - output = attn.to_out[0](output) + output = attn.to_out[0](hidden_states) if len(attn.to_out) > 1: # dropout output = attn.to_out[1](output) @@ -214,11 +156,8 @@ class FeedForward(nn.Module): def __init__(self, dim: int, hidden_dim: int): super().__init__() self.w1 = nn.Linear(dim, hidden_dim, bias=False) - nn.init.xavier_uniform_(self.w1.weight) self.w2 = nn.Linear(hidden_dim, dim, bias=False) - nn.init.xavier_uniform_(self.w2.weight) self.w3 = nn.Linear(dim, hidden_dim, bias=False) - nn.init.xavier_uniform_(self.w3.weight) def _forward_silu_gating(self, x1, x3): return F.silu(x1) * x3 @@ -251,8 +190,9 @@ def __init__( dim_head=dim // n_heads, heads=n_heads, qk_norm="rms_norm" if qk_norm else None, - eps=1e-6, + eps=1e-5, bias=False, + out_bias=False, processor=ZSingleStreamAttnProcessor(), ) @@ -270,16 +210,12 @@ def __init__( self.adaLN_modulation = nn.Sequential( nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True), ) - nn.init.zeros_(self.adaLN_modulation[0].weight) - nn.init.zeros_(self.adaLN_modulation[0].bias) def forward( self, x: torch.Tensor, - x_src_ids: torch.Tensor, - x_freqs_cis: torch.Tensor, - x_cu_seqlens: torch.Tensor, - x_max_item_seqlen: int, + attn_mask: torch.Tensor, + freqs_cis: torch.Tensor, adaln_input: Optional[torch.Tensor] = None, ): if self.modulation: @@ -290,26 +226,24 @@ def forward( # Attention block attn_out = self.attention( - self.attention_norm1(x) * scale_msa[x_src_ids], - image_rotary_emb=x_freqs_cis, - x_cu_seqlens=x_cu_seqlens, - x_max_item_seqlen=x_max_item_seqlen, + self.attention_norm1(x) * scale_msa, + attention_mask=attn_mask, + freqs_cis=freqs_cis, ) - x = x + gate_msa[x_src_ids] * self.attention_norm2(attn_out) + x = x + gate_msa * self.attention_norm2(attn_out) # FFN block - x = x + gate_mlp[x_src_ids] * self.ffn_norm2( + x = x + gate_mlp * self.ffn_norm2( self.feed_forward( - self.ffn_norm1(x) * scale_mlp[x_src_ids], + self.ffn_norm1(x) * scale_mlp, ) ) else: # Attention block attn_out = self.attention( self.attention_norm1(x), - image_rotary_emb=x_freqs_cis, - x_cu_seqlens=x_cu_seqlens, - x_max_item_seqlen=x_max_item_seqlen, + attention_mask=attn_mask, + freqs_cis=freqs_cis, ) x = x + self.attention_norm2(attn_out) @@ -328,19 +262,15 @@ def __init__(self, hidden_size, out_channels): super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.linear = nn.Linear(hidden_size, out_channels, bias=True) - nn.init.zeros_(self.linear.weight) - nn.init.zeros_(self.linear.bias) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), ) - nn.init.zeros_(self.adaLN_modulation[1].weight) - nn.init.zeros_(self.adaLN_modulation[1].bias) - def forward(self, x, x_src_ids, c): + def forward(self, x, c): scale = 1.0 + self.adaLN_modulation(c) - x = self.norm_final(x) * scale[x_src_ids] + x = self.norm_final(x) * scale x = self.linear(x) return x @@ -466,13 +396,9 @@ def __init__( RMSNorm(cap_feat_dim, eps=norm_eps), nn.Linear(cap_feat_dim, dim, bias=True), ) - nn.init.trunc_normal_(self.cap_embedder[1].weight, std=0.02) - nn.init.zeros_(self.cap_embedder[1].bias) self.x_pad_token = nn.Parameter(torch.empty((1, dim))) - nn.init.normal_(self.x_pad_token, std=0.02) self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) - nn.init.normal_(self.cap_pad_token, std=0.02) self.layers = nn.ModuleList( [ @@ -646,31 +572,38 @@ def forward( x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) x[torch.cat(x_inner_pad_mask)] = self.x_pad_token x = x.split(x_item_seqlens, dim=0) - x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) # todo + x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) pad_tensor = torch.zeros( (1, self.dim), dtype=x[0].dtype, device=device, ) - x_pad_mask = torch.zeros( + freqs_pad_tensor = torch.zeros( + (1, self.dim // self.n_heads), + dtype=x_freqs_cis[0].dtype, + device=device, + ) + x_attn_mask = torch.ones( (bsz, x_max_item_seqlen), dtype=torch.bool, device=device ) - for i, item in enumerate(x): + for i, (item, freqs_item) in enumerate(zip(x, x_freqs_cis)): seq_len = x_item_seqlens[i] - x[i] = torch.cat([item, pad_tensor.repeat(x_max_item_seqlen - seq_len, 1)]) - x_pad_mask[i, seq_len:] = 1 + pad_len = x_max_item_seqlen - seq_len + x[i] = torch.cat([item, pad_tensor.repeat(pad_len, 1)]) + x_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) + x_attn_mask[i, seq_len:] = 0 x = torch.stack(x) for layer in self.noise_refiner: x = layer( x, - x_pad_mask, + x_attn_mask, x_freqs_cis, adaln_input, - ) # todo + ) # cap embed & refine cap_item_seqlens = [len(_) for _ in cap_feats] @@ -681,72 +614,83 @@ def forward( cap_feats = self.cap_embedder(cap_feats) cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token cap_feats = cap_feats.split(cap_item_seqlens, dim=0) - cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) # todo + cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) pad_tensor = torch.zeros( (1, self.dim), - dtype=x[0].dtype, + dtype=cap_feats[0].dtype, + device=device, + ) + freqs_pad_tensor = torch.zeros( + (1, self.dim // self.n_heads), + dtype=cap_freqs_cis[0].dtype, device=device, ) - cap_pad_mask = torch.zeros( + cap_attn_mask = torch.ones( (bsz, cap_max_item_seqlen), dtype=torch.bool, device=device ) - for i, item in enumerate(cap_feats): + for i, (item, freqs_item) in enumerate(zip(cap_feats, cap_freqs_cis)): seq_len = cap_item_seqlens[i] - cap_feats[i] = torch.cat([item, pad_tensor.repeat(cap_max_item_seqlen - seq_len, 1)]) - cap_pad_mask[i, seq_len:] = 1 + pad_len = cap_max_item_seqlen - seq_len + cap_feats[i] = torch.cat([item, pad_tensor.repeat(pad_len, 1)]) + cap_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) + cap_attn_mask[i, seq_len:] = 0 cap_feats = torch.stack(cap_feats) + for layer in self.context_refiner: cap_feats = layer( cap_feats, - cap_pad_mask, + cap_attn_mask, cap_freqs_cis, ) - # unified todo + # unified + unified = [] + unified_freqs_cis = [] + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) + unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]])) unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + assert unified_item_seqlens == [len(_) for _ in unified] unified_max_item_seqlen = max(unified_item_seqlens) pad_tensor = torch.zeros( (1, self.dim), - dtype=x[0].dtype, + dtype=unified[0].dtype, device=device, ) - unified_pad_mask = torch.zeros( + freqs_pad_tensor = torch.zeros( + (1, self.dim // self.n_heads), + dtype=unified_freqs_cis[0].dtype, + device=device, + ) + unified_attn_mask = torch.ones( (bsz, unified_max_item_seqlen), dtype=torch.bool, device=device ) - - unified = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - unified.append( - torch.cat( - [ - x[i][:x_item_seqlens[i]], - cap_feats[i][:cap_item_seqlens[i]], - pad_tensor.repeat(unified_max_item_seqlen - x_len - cap_len, 1) - ] - ) - ) - unified_pad_mask[i, x_len + cap_len:] = 1 - - unified_freqs_cis = torch.cat(merge_interleave(cap_freqs_cis, x_freqs_cis)) # todo + for i, (item, freqs_item) in enumerate(zip(unified, unified_freqs_cis)): + seq_len = unified_item_seqlens[i] + pad_len = unified_max_item_seqlen - seq_len + unified[i] = torch.cat([item, pad_tensor.repeat(pad_len, 1)]) + unified_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) + unified_attn_mask[i, seq_len:] = 0 + unified = torch.stack(unified) for layer in self.layers: - unified_shard = layer( + unified = layer( unified, - unified_pad_mask, + unified_attn_mask, unified_freqs_cis, adaln_input, ) unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( - unified, adaln_input # todo + unified, adaln_input ) unified = unified.split(unified_item_seqlens, dim=0) x = self.unpatchify(unified, x_size, patch_size, f_patch_size) From 21d81302206f6a78c76962c04fa8f03d72746c69 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 18:40:19 +0800 Subject: [PATCH 06/31] fixed stack issue --- src/diffusers/models/transformers/transformer_z_image.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 62938270bd89..359a5c6a04e7 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -304,10 +304,11 @@ def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): def __call__(self, ids: torch.Tensor): assert ids.ndim == 2 assert ids.shape[-1] == len(self.axes_dims) + device = ids.device if self.freqs_cis is None: self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) - self.freqs_cis = [freqs_cis.cuda() for freqs_cis in self.freqs_cis] + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] result = [] for i in range(len(self.axes_dims)): @@ -596,6 +597,7 @@ def forward( x_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) x_attn_mask[i, seq_len:] = 0 x = torch.stack(x) + x_freqs_cis = torch.stack(x_freqs_cis) for layer in self.noise_refiner: x = layer( @@ -638,6 +640,7 @@ def forward( cap_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) cap_attn_mask[i, seq_len:] = 0 cap_feats = torch.stack(cap_feats) + cap_freqs_cis = torch.stack(cap_freqs_cis) for layer in self.context_refiner: cap_feats = layer( @@ -680,6 +683,7 @@ def forward( unified_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) unified_attn_mask[i, seq_len:] = 0 unified = torch.stack(unified) + unified_freqs_cis = torch.stack(unified_freqs_cis) for layer in self.layers: unified = layer( From e3dfa9e6687ddc231e2f3fe8ee442eaff4c2760d Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 19:06:32 +0800 Subject: [PATCH 07/31] fixed modulation bug --- src/diffusers/models/transformers/transformer_z_image.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 359a5c6a04e7..1e77a0180c50 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -26,7 +26,7 @@ from ..attention_dispatch import dispatch_attention_fn from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin -from ...utils.import_utils import is_apex_available, is_flash_attn_available +from ...utils.import_utils import is_flash_attn_available from ...utils.torch_utils import maybe_allow_in_graph @@ -220,7 +220,7 @@ def forward( ): if self.modulation: assert adaln_input is not None - scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).chunk(4, dim=1) + scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(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 @@ -270,7 +270,7 @@ def __init__(self, hidden_size, out_channels): def forward(self, x, c): scale = 1.0 + self.adaLN_modulation(c) - x = self.norm_final(x) * scale + x = self.norm_final(x) * scale.unsqueeze(1) x = self.linear(x) return x @@ -696,7 +696,7 @@ def forward( unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( unified, adaln_input ) - unified = unified.split(unified_item_seqlens, dim=0) + unified = unified.unbind(dim=0) x = self.unpatchify(unified, x_size, patch_size, f_patch_size) return x, {} From a7fa73140c61e8926cdb7704e5d83cfabb080330 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 19:06:41 +0800 Subject: [PATCH 08/31] fixed modulation bug --- src/diffusers/utils/import_utils.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 520a6c1510d9..57b0a337922a 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -230,7 +230,6 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b _aiter_available, _aiter_version = _is_package_available("aiter") _kornia_available, _kornia_version = _is_package_available("kornia") _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True) -_apex_available, _apex_version = _is_package_available("apex") def is_torch_available(): @@ -421,10 +420,6 @@ def is_kornia_available(): return _kornia_available -def is_apex_available(): - return _apex_available - - # docstyle-ignore FLAX_IMPORT_ERROR = """ {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the From 1e0cefe14207c30cab03d280a95eff7088262237 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 19:24:51 +0800 Subject: [PATCH 09/31] fix bug --- .../models/transformers/transformer_z_image.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 1e77a0180c50..ca5e1036de5f 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -95,6 +95,7 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[torch.Tensor] = None, ) -> torch.Tensor: @@ -572,8 +573,8 @@ def forward( x = torch.cat(x, dim=0) x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) x[torch.cat(x_inner_pad_mask)] = self.x_pad_token - x = x.split(x_item_seqlens, dim=0) - x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) + x = list(x.split(x_item_seqlens, dim=0)) + x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) pad_tensor = torch.zeros( (1, self.dim), @@ -581,7 +582,7 @@ def forward( device=device, ) freqs_pad_tensor = torch.zeros( - (1, self.dim // self.n_heads), + (1, self.dim // self.n_heads // 2), dtype=x_freqs_cis[0].dtype, device=device, ) @@ -615,8 +616,8 @@ def forward( cap_feats = torch.cat(cap_feats, dim=0) cap_feats = self.cap_embedder(cap_feats) cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token - cap_feats = cap_feats.split(cap_item_seqlens, dim=0) - cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0) + cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) + cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) pad_tensor = torch.zeros( (1, self.dim), @@ -624,7 +625,7 @@ def forward( device=device, ) freqs_pad_tensor = torch.zeros( - (1, self.dim // self.n_heads), + (1, self.dim // self.n_heads // 2), dtype=cap_freqs_cis[0].dtype, device=device, ) @@ -667,7 +668,7 @@ def forward( device=device, ) freqs_pad_tensor = torch.zeros( - (1, self.dim // self.n_heads), + (1, self.dim // self.n_heads // 2), dtype=unified_freqs_cis[0].dtype, device=device, ) @@ -696,7 +697,7 @@ def forward( unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( unified, adaln_input ) - unified = unified.unbind(dim=0) + unified = list(unified.unbind(dim=0)) x = self.unpatchify(unified, x_size, patch_size, f_patch_size) return x, {} From 7adaae888d19d87176de645e93f00730dd70e59a Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 23:04:43 +0800 Subject: [PATCH 10/31] remove value_from_time_aware_config --- .../pipelines/z_image/pipeline_z_image.py | 24 ++----------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 5abb73550c25..4e8f78228531 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -293,25 +293,6 @@ def num_timesteps(self): def interrupt(self): return self._interrupt - @staticmethod - def value_from_time_aware_config(config, t): - if isinstance(config, (float, int, str)): - return config - elif isinstance(config, torch.Tensor): - assert config.numel() == 1 - return config.item() - elif isinstance(config, (tuple, list)): - assert isinstance(config[0], (float, int, str)) - result = config[0] - for thresh, val in config[1:]: - if t >= thresh: - result = val - else: - break - return result - else: - raise ValueError(f"invalid time-aware config {config} of type {type(config)}") - @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( @@ -507,9 +488,8 @@ def __call__( and self._cfg_truncation is not None and float(self._cfg_truncation) <= 1 ): - current_guidance_scale = self.value_from_time_aware_config( - (self.guidance_scale, (self._cfg_truncation, 0.0)), t_norm - ) + if t_norm > self._cfg_truncation: + current_guidance_scale = 0.0 # Run CFG only if configured AND scale is non-zero apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 From 5b4c907407f30d16b2fcd948981a04d40bc4198e Mon Sep 17 00:00:00 2001 From: liudongyang Date: Mon, 24 Nov 2025 23:26:15 +0800 Subject: [PATCH 11/31] styling --- .../transformers/transformer_z_image.py | 26 ++++--------------- 1 file changed, 5 insertions(+), 21 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index ca5e1036de5f..22b79c211a92 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import itertools import math from typing import List, Optional, Tuple @@ -23,11 +22,11 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ..attention_dispatch import dispatch_attention_fn from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin from ...utils.import_utils import is_flash_attn_available from ...utils.torch_utils import maybe_allow_in_graph +from ..attention_dispatch import dispatch_attention_fn if is_flash_attn_available(): @@ -99,7 +98,6 @@ def __call__( attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[torch.Tensor] = None, ) -> torch.Tensor: - query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) value = attn.to_v(hidden_states) @@ -586,11 +584,7 @@ def forward( dtype=x_freqs_cis[0].dtype, device=device, ) - x_attn_mask = torch.ones( - (bsz, x_max_item_seqlen), - dtype=torch.bool, - device=device - ) + x_attn_mask = torch.ones((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) for i, (item, freqs_item) in enumerate(zip(x, x_freqs_cis)): seq_len = x_item_seqlens[i] pad_len = x_max_item_seqlen - seq_len @@ -629,11 +623,7 @@ def forward( dtype=cap_freqs_cis[0].dtype, device=device, ) - cap_attn_mask = torch.ones( - (bsz, cap_max_item_seqlen), - dtype=torch.bool, - device=device - ) + cap_attn_mask = torch.ones((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) for i, (item, freqs_item) in enumerate(zip(cap_feats, cap_freqs_cis)): seq_len = cap_item_seqlens[i] pad_len = cap_max_item_seqlen - seq_len @@ -672,11 +662,7 @@ def forward( dtype=unified_freqs_cis[0].dtype, device=device, ) - unified_attn_mask = torch.ones( - (bsz, unified_max_item_seqlen), - dtype=torch.bool, - device=device - ) + unified_attn_mask = torch.ones((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) for i, (item, freqs_item) in enumerate(zip(unified, unified_freqs_cis)): seq_len = unified_item_seqlens[i] pad_len = unified_max_item_seqlen - seq_len @@ -694,9 +680,7 @@ def forward( adaln_input, ) - unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( - unified, adaln_input - ) + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input) unified = list(unified.unbind(dim=0)) x = self.unpatchify(unified, x_size, patch_size, f_patch_size) From 2bb39f46cdfcc90258c04ea9591a8b0cf180baa1 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Mon, 24 Nov 2025 18:26:28 +0000 Subject: [PATCH 12/31] Fix neg embed and devide / bug; Reuse pad zero tensor; Turn cat -> repeat; Add hint for attn processor. --- .../transformers/transformer_z_image.py | 50 ++++++++----------- .../pipelines/z_image/pipeline_z_image.py | 24 +++++++-- 2 files changed, 41 insertions(+), 33 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 22b79c211a92..39f8a2a02f13 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -90,6 +90,12 @@ class ZSingleStreamAttnProcessor: _attention_backend = None _parallel_config = None + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "ZSingleStreamAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to version 2.0 or higher." + ) + def __call__( self, attn: Attention, @@ -493,7 +499,6 @@ def patchify_and_embed( image_ori_len = len(image) image_padding_len = (-image_ori_len) % SEQ_MULTI_OF - # padded_pos_ids image_ori_pos_ids = self.create_coordinate_grid( size=(F_tokens, H_tokens, W_tokens), @@ -574,11 +579,7 @@ def forward( x = list(x.split(x_item_seqlens, dim=0)) x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) - pad_tensor = torch.zeros( - (1, self.dim), - dtype=x[0].dtype, - device=device, - ) + pad_tensor = torch.zeros((1, self.dim), dtype=x[0].dtype, device=device) freqs_pad_tensor = torch.zeros( (1, self.dim // self.n_heads // 2), dtype=x_freqs_cis[0].dtype, @@ -613,22 +614,19 @@ def forward( cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) - pad_tensor = torch.zeros( - (1, self.dim), - dtype=cap_feats[0].dtype, - device=device, - ) - freqs_pad_tensor = torch.zeros( - (1, self.dim // self.n_heads // 2), - dtype=cap_freqs_cis[0].dtype, - device=device, + # Reuse padding tensors (convert dtype if needed) + cap_pad_tensor = pad_tensor.to(cap_feats[0].dtype) if pad_tensor.dtype != cap_feats[0].dtype else pad_tensor + cap_freqs_pad_tensor = ( + freqs_pad_tensor.to(cap_freqs_cis[0].dtype) + if freqs_pad_tensor.dtype != cap_freqs_cis[0].dtype + else freqs_pad_tensor ) cap_attn_mask = torch.ones((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) for i, (item, freqs_item) in enumerate(zip(cap_feats, cap_freqs_cis)): seq_len = cap_item_seqlens[i] pad_len = cap_max_item_seqlen - seq_len - cap_feats[i] = torch.cat([item, pad_tensor.repeat(pad_len, 1)]) - cap_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) + cap_feats[i] = torch.cat([item, cap_pad_tensor.repeat(pad_len, 1)]) + cap_freqs_cis[i] = torch.cat([freqs_item, cap_freqs_pad_tensor.repeat(pad_len, 1)]) cap_attn_mask[i, seq_len:] = 0 cap_feats = torch.stack(cap_feats) cap_freqs_cis = torch.stack(cap_freqs_cis) @@ -652,22 +650,18 @@ def forward( assert unified_item_seqlens == [len(_) for _ in unified] unified_max_item_seqlen = max(unified_item_seqlens) - pad_tensor = torch.zeros( - (1, self.dim), - dtype=unified[0].dtype, - device=device, - ) - freqs_pad_tensor = torch.zeros( - (1, self.dim // self.n_heads // 2), - dtype=unified_freqs_cis[0].dtype, - device=device, + unified_pad_tensor = pad_tensor.to(unified[0].dtype) if pad_tensor.dtype != unified[0].dtype else pad_tensor + unified_freqs_pad_tensor = ( + freqs_pad_tensor.to(unified_freqs_cis[0].dtype) + if freqs_pad_tensor.dtype != unified_freqs_cis[0].dtype + else freqs_pad_tensor ) unified_attn_mask = torch.ones((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) for i, (item, freqs_item) in enumerate(zip(unified, unified_freqs_cis)): seq_len = unified_item_seqlens[i] pad_len = unified_max_item_seqlen - seq_len - unified[i] = torch.cat([item, pad_tensor.repeat(pad_len, 1)]) - unified_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) + unified[i] = torch.cat([item, unified_pad_tensor.repeat(pad_len, 1)]) + unified_freqs_cis[i] = torch.cat([freqs_item, unified_freqs_pad_tensor.repeat(pad_len, 1)]) unified_attn_mask[i, seq_len:] = 0 unified = torch.stack(unified) unified_freqs_cis = torch.stack(unified_freqs_cis) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 4e8f78228531..d33d72d9c735 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -193,6 +193,8 @@ def encode_prompt( prompt_embeds=negative_prompt_embeds, max_sequence_length=max_sequence_length, ) + else: + negative_prompt_embeds = [] return prompt_embeds, negative_prompt_embeds def _encode_prompt( @@ -398,6 +400,18 @@ def __call__( height = height or 1024 width = width or 1024 + vae_scale = self.vae_scale_factor * 2 + if height % vae_scale != 0: + raise ValueError( + f"Height must be divisible by {vae_scale} (got {height}). " + f"Please adjust the height to a multiple of {vae_scale}." + ) + if width % vae_scale != 0: + raise ValueError( + f"Width must be divisible by {vae_scale} (got {width}). " + f"Please adjust the width to a multiple of {vae_scale}." + ) + assert self.dtype == torch.bfloat16 dtype = self.dtype device = self._execution_device @@ -447,7 +461,7 @@ def __call__( generator, latents, ) - image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] / 2) + image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) # 5. Prepare timesteps mu = calculate_shift( @@ -495,12 +509,12 @@ def __call__( apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 if apply_cfg: - # Prepare inputs for CFG - latent_model_input = torch.cat([latents.to(dtype)] * 2) + latents_typed = latents if latents.dtype == dtype else latents.to(dtype) + latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds - timestep_model_input = torch.cat([timestep] * 2) + timestep_model_input = timestep.repeat(2) else: - latent_model_input = latents.to(dtype) + latent_model_input = latents if latents.dtype == dtype else latents.to(dtype) prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep From 71e8049a846568f3f0c3003df7decd7adaf79301 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Mon, 24 Nov 2025 18:37:24 +0000 Subject: [PATCH 13/31] Replace padding with pad_sequence; Add gradient checkpointing. --- .../transformers/transformer_z_image.py | 107 +++++++----------- 1 file changed, 38 insertions(+), 69 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 39f8a2a02f13..599046c68cfc 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -19,6 +19,7 @@ import torch.nn as nn import torch.nn.functional as F from einops import rearrange +from torch.nn.utils.rnn import pad_sequence from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin @@ -355,6 +356,7 @@ def __init__( self.rope_theta = rope_theta self.t_scale = t_scale + self.gradient_checkpointing = False assert len(all_patch_size) == len(all_f_patch_size) @@ -579,29 +581,18 @@ def forward( x = list(x.split(x_item_seqlens, dim=0)) x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) - pad_tensor = torch.zeros((1, self.dim), dtype=x[0].dtype, device=device) - freqs_pad_tensor = torch.zeros( - (1, self.dim // self.n_heads // 2), - dtype=x_freqs_cis[0].dtype, - device=device, - ) - x_attn_mask = torch.ones((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) - for i, (item, freqs_item) in enumerate(zip(x, x_freqs_cis)): - seq_len = x_item_seqlens[i] - pad_len = x_max_item_seqlen - seq_len - x[i] = torch.cat([item, pad_tensor.repeat(pad_len, 1)]) - x_freqs_cis[i] = torch.cat([freqs_item, freqs_pad_tensor.repeat(pad_len, 1)]) - x_attn_mask[i, seq_len:] = 0 - x = torch.stack(x) - x_freqs_cis = torch.stack(x_freqs_cis) - - for layer in self.noise_refiner: - x = layer( - x, - x_attn_mask, - x_freqs_cis, - adaln_input, - ) + x = pad_sequence(x, batch_first=True, padding_value=0.0) + x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) + x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(x_item_seqlens): + x_attn_mask[i, :seq_len] = 1 + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for layer in self.noise_refiner: + x = self._gradient_checkpointing_func(layer, x, x_attn_mask, x_freqs_cis, adaln_input) + else: + for layer in self.noise_refiner: + x = layer(x, x_attn_mask, x_freqs_cis, adaln_input) # cap embed & refine cap_item_seqlens = [len(_) for _ in cap_feats] @@ -614,29 +605,18 @@ def forward( cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) - # Reuse padding tensors (convert dtype if needed) - cap_pad_tensor = pad_tensor.to(cap_feats[0].dtype) if pad_tensor.dtype != cap_feats[0].dtype else pad_tensor - cap_freqs_pad_tensor = ( - freqs_pad_tensor.to(cap_freqs_cis[0].dtype) - if freqs_pad_tensor.dtype != cap_freqs_cis[0].dtype - else freqs_pad_tensor - ) - cap_attn_mask = torch.ones((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) - for i, (item, freqs_item) in enumerate(zip(cap_feats, cap_freqs_cis)): - seq_len = cap_item_seqlens[i] - pad_len = cap_max_item_seqlen - seq_len - cap_feats[i] = torch.cat([item, cap_pad_tensor.repeat(pad_len, 1)]) - cap_freqs_cis[i] = torch.cat([freqs_item, cap_freqs_pad_tensor.repeat(pad_len, 1)]) - cap_attn_mask[i, seq_len:] = 0 - cap_feats = torch.stack(cap_feats) - cap_freqs_cis = torch.stack(cap_freqs_cis) - - for layer in self.context_refiner: - cap_feats = layer( - cap_feats, - cap_attn_mask, - cap_freqs_cis, - ) + cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0) + cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) + cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(cap_item_seqlens): + cap_attn_mask[i, :seq_len] = 1 + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for layer in self.context_refiner: + cap_feats = self._gradient_checkpointing_func(layer, cap_feats, cap_attn_mask, cap_freqs_cis) + else: + for layer in self.context_refiner: + cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis) # unified unified = [] @@ -650,29 +630,18 @@ def forward( assert unified_item_seqlens == [len(_) for _ in unified] unified_max_item_seqlen = max(unified_item_seqlens) - unified_pad_tensor = pad_tensor.to(unified[0].dtype) if pad_tensor.dtype != unified[0].dtype else pad_tensor - unified_freqs_pad_tensor = ( - freqs_pad_tensor.to(unified_freqs_cis[0].dtype) - if freqs_pad_tensor.dtype != unified_freqs_cis[0].dtype - else freqs_pad_tensor - ) - unified_attn_mask = torch.ones((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, (item, freqs_item) in enumerate(zip(unified, unified_freqs_cis)): - seq_len = unified_item_seqlens[i] - pad_len = unified_max_item_seqlen - seq_len - unified[i] = torch.cat([item, unified_pad_tensor.repeat(pad_len, 1)]) - unified_freqs_cis[i] = torch.cat([freqs_item, unified_freqs_pad_tensor.repeat(pad_len, 1)]) - unified_attn_mask[i, seq_len:] = 0 - unified = torch.stack(unified) - unified_freqs_cis = torch.stack(unified_freqs_cis) - - for layer in self.layers: - unified = layer( - unified, - unified_attn_mask, - unified_freqs_cis, - adaln_input, - ) + 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 = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(unified_item_seqlens): + unified_attn_mask[i, :seq_len] = 1 + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for layer in self.layers: + unified = self._gradient_checkpointing_func(layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input) + else: + 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) unified = list(unified.unbind(dim=0)) From fbf26b7ed11d55146103c97740bad4a5f91744e0 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Mon, 24 Nov 2025 18:49:45 +0000 Subject: [PATCH 14/31] Fix flash_attn3 in dispatch attn backend by _flash_attn_forward, replace its origin implement; Add DocString in pipeline for that. --- src/diffusers/models/attention_dispatch.py | 30 ++++++++++++++++--- .../pipelines/z_image/pipeline_z_image.py | 9 ++++-- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 8504504981a3..df4e0a01220e 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -79,9 +79,11 @@ if _CAN_USE_FLASH_ATTN_3: from flash_attn_interface import flash_attn_func as flash_attn_3_func from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func + from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward else: flash_attn_3_func = None flash_attn_3_varlen_func = None + flash_attn_3_forward = None if _CAN_USE_AITER_ATTN: from aiter import flash_attn_func as aiter_flash_attn_func @@ -621,22 +623,42 @@ def _wrapped_flash_attn_3( ) -> Tuple[torch.Tensor, torch.Tensor]: # Hardcoded for now because pytorch does not support tuple/int type hints window_size = (-1, -1) - out, lse, *_ = flash_attn_3_func( + max_seqlen_q = q.shape[2] + max_seqlen_k = k.shape[2] + + out, lse, *_ = flash_attn_3_forward( q=q, k=k, v=v, - softmax_scale=softmax_scale, - causal=causal, + k_new=None, + v_new=None, qv=qv, + out=None, + cu_seqlens_q=None, + cu_seqlens_k=None, + cu_seqlens_k_new=None, + seqused_q=None, + seqused_k=None, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + page_table=None, + kv_batch_idx=None, + leftpad_k=None, + rotary_cos=None, + rotary_sin=None, + seqlens_rotary=None, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, + softmax_scale=softmax_scale, + causal=causal, window_size=window_size, attention_chunk=attention_chunk, softcap=softcap, + rotary_interleaved=True, + scheduler_metadata=None, num_splits=num_splits, pack_gqa=pack_gqa, - deterministic=deterministic, sm_margin=sm_margin, ) lse = lse.permute(0, 2, 1) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index d33d72d9c735..d4cd574a04f7 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -39,9 +39,14 @@ >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) >>> pipe.to("cuda") + + >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. + >>> # (1) Use flash attention 2 + >>> # pipe.transformer.set_attention_backend("flash") + >>> # (2) Use flash attention 3 + >>> # pipe.transformer.set_attention_backend("_flash_3") + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" - >>> # Depending on the variant being used, the pipeline call will slightly vary. - >>> # Refer to the pipeline documentation for more details. >>> image = pipe( ... prompt, ... height=1024, From 6c0c059facdb709d1c346e89bc0009cafd822473 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Mon, 24 Nov 2025 18:59:29 +0000 Subject: [PATCH 15/31] Fix Docstring and Make Style. --- src/diffusers/models/attention_dispatch.py | 129 ++++++++++++++---- .../transformers/transformer_z_image.py | 4 +- .../pipelines/z_image/pipeline_z_image.py | 3 +- 3 files changed, 106 insertions(+), 30 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index df4e0a01220e..74994bd2c62c 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -18,7 +18,17 @@ import math from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Tuple, + Union, +) import torch @@ -68,7 +78,10 @@ if _CAN_USE_FLASH_ATTN: from flash_attn import flash_attn_func, flash_attn_varlen_func - from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward + from flash_attn.flash_attn_interface import ( + _wrapped_flash_attn_backward, + _wrapped_flash_attn_forward, + ) else: flash_attn_func = None flash_attn_varlen_func = None @@ -77,9 +90,9 @@ if _CAN_USE_FLASH_ATTN_3: + from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward from flash_attn_interface import flash_attn_func as flash_attn_3_func from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func - from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward else: flash_attn_3_func = None flash_attn_3_varlen_func = None @@ -122,7 +135,9 @@ if _CAN_USE_XLA_ATTN: - from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention + from torch_xla.experimental.custom_kernel import ( + flash_attention as xla_flash_attention, + ) else: xla_flash_attention = None @@ -265,13 +280,17 @@ class _HubKernelConfig: _HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = { # TODO: temporary revision for now. Remove when merged upstream into `main`. AttentionBackendName._FLASH_3_HUB: _HubKernelConfig( - repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs" + repo_id="kernels-community/flash-attn3", + function_attr="flash_attn_func", + revision="fake-ops-return-probs", ) } @contextlib.contextmanager -def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE): +def attention_backend( + backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE, +): """ Context manager to set the active attention backend. """ @@ -416,7 +435,10 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None f"Flash Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `flash-attn>={_REQUIRED_FLASH_VERSION}`." ) - elif backend in [AttentionBackendName._FLASH_3, AttentionBackendName._FLASH_VARLEN_3]: + elif backend in [ + AttentionBackendName._FLASH_3, + AttentionBackendName._FLASH_VARLEN_3, + ]: if not _CAN_USE_FLASH_ATTN_3: raise RuntimeError( f"Flash Attention 3 backend '{backend.value}' is not usable because of missing package or the version is too old. Please build FA3 beta release from source." @@ -488,7 +510,11 @@ def _prepare_for_flash_attn_or_sage_varlen_without_mask( cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) max_seqlen_q = seqlens_q.max().item() max_seqlen_k = seqlens_k.max().item() - return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) + return ( + (seqlens_q, seqlens_k), + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_q, max_seqlen_k), + ) def _prepare_for_flash_attn_or_sage_varlen_with_mask( @@ -505,7 +531,11 @@ def _prepare_for_flash_attn_or_sage_varlen_with_mask( cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) max_seqlen_q = seqlens_q.max().item() max_seqlen_k = seqlens_k.max().item() - return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) + return ( + (seqlens_q, seqlens_k), + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_q, max_seqlen_k), + ) def _prepare_for_flash_attn_or_sage_varlen( @@ -625,7 +655,7 @@ def _wrapped_flash_attn_3( window_size = (-1, -1) max_seqlen_q = q.shape[2] max_seqlen_k = k.shape[2] - + out, lse, *_ = flash_attn_3_forward( q=q, k=k, @@ -764,7 +794,10 @@ def _native_attention_backward_op( grad_out_t = grad_out.permute(0, 2, 1, 3) grad_query_t, grad_key_t, grad_value_t = torch.autograd.grad( - outputs=out, inputs=[query_t, key_t, value_t], grad_outputs=grad_out_t, retain_graph=False + outputs=out, + inputs=[query_t, key_t, value_t], + grad_outputs=grad_out_t, + retain_graph=False, ) grad_query = grad_query_t.permute(0, 2, 1, 3) @@ -803,18 +836,26 @@ def _cudnn_attention_forward_op( value = value.transpose(1, 2).contiguous() tensors_to_save += (query, key, value) - out, lse, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask = ( - torch.ops.aten._scaled_dot_product_cudnn_attention( - query=query, - key=key, - value=value, - attn_bias=attn_mask, - compute_log_sumexp=return_lse, - dropout_p=dropout_p, - is_causal=is_causal, - return_debug_mask=False, - scale=scale, - ) + ( + out, + lse, + cum_seq_q, + cum_seq_k, + max_q, + max_k, + philox_seed, + philox_offset, + debug_attn_mask, + ) = torch.ops.aten._scaled_dot_product_cudnn_attention( + query=query, + key=key, + value=value, + attn_bias=attn_mask, + compute_log_sumexp=return_lse, + dropout_p=dropout_p, + is_causal=is_causal, + return_debug_mask=False, + scale=scale, ) tensors_to_save += (out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset) @@ -941,7 +982,11 @@ def _flash_attention_backward_op( **kwargs, ): query, key, value, out, lse, rng_state = ctx.saved_tensors - grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) + grad_query, grad_key, grad_value = ( + torch.empty_like(query), + torch.empty_like(key), + torch.empty_like(value), + ) lse_d = _wrapped_flash_attn_backward( # noqa: F841 grad_out, @@ -1165,7 +1210,19 @@ def backward( grad_query, grad_key, grad_value = (x.to(grad_out.dtype) for x in (grad_query, grad_key, grad_value)) - return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None + return ( + grad_query, + grad_key, + grad_value, + None, + None, + None, + None, + None, + None, + None, + None, + ) class TemplatedUlyssesAttention(torch.autograd.Function): @@ -1260,7 +1317,19 @@ def backward( x.flatten(0, 1).permute(1, 2, 0, 3).contiguous() for x in (grad_query, grad_key, grad_value) ) - return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None + return ( + grad_query, + grad_key, + grad_value, + None, + None, + None, + None, + None, + None, + None, + None, + ) def _templated_context_parallel_attention( @@ -1608,7 +1677,12 @@ def _native_flex_attention( block_mask = attn_mask elif is_causal: block_mask = flex_attention.create_block_mask( - _flex_attention_causal_mask_mod, batch_size, num_heads, seq_len_q, seq_len_kv, query.device + _flex_attention_causal_mask_mod, + batch_size, + num_heads, + seq_len_q, + seq_len_kv, + query.device, ) elif torch.is_tensor(attn_mask): if attn_mask.ndim == 2: @@ -1628,6 +1702,7 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx] + else: raise ValueError("Attention mask must be either None, a BlockMask, or a 2D/4D tensor.") diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 599046c68cfc..24c3189a8213 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -638,7 +638,9 @@ def forward( if torch.is_grad_enabled() and self.gradient_checkpointing: for layer in self.layers: - unified = self._gradient_checkpointing_func(layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input) + unified = self._gradient_checkpointing_func( + layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input + ) else: for layer in self.layers: unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index d4cd574a04f7..3c334facc68b 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -45,8 +45,7 @@ >>> # pipe.transformer.set_attention_backend("flash") >>> # (2) Use flash attention 3 >>> # pipe.transformer.set_attention_backend("_flash_3") - - >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" + >>> prompt = '一幅为名为"造相「Z-IMAGE-TURBO」"的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。' >>> image = pipe( ... prompt, ... height=1024, From 28685dd3d4823ddb0af11b7bb5da0f90c9dd48ed Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 13:08:57 +0800 Subject: [PATCH 16/31] Revert "Fix flash_attn3 in dispatch attn backend by _flash_attn_forward, replace its origin implement; Add DocString in pipeline for that." This reverts commit fbf26b7ed11d55146103c97740bad4a5f91744e0. --- src/diffusers/models/attention_dispatch.py | 30 +++---------------- .../pipelines/z_image/pipeline_z_image.py | 9 ++---- 2 files changed, 6 insertions(+), 33 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index df4e0a01220e..8504504981a3 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -79,11 +79,9 @@ if _CAN_USE_FLASH_ATTN_3: from flash_attn_interface import flash_attn_func as flash_attn_3_func from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func - from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward else: flash_attn_3_func = None flash_attn_3_varlen_func = None - flash_attn_3_forward = None if _CAN_USE_AITER_ATTN: from aiter import flash_attn_func as aiter_flash_attn_func @@ -623,42 +621,22 @@ def _wrapped_flash_attn_3( ) -> Tuple[torch.Tensor, torch.Tensor]: # Hardcoded for now because pytorch does not support tuple/int type hints window_size = (-1, -1) - max_seqlen_q = q.shape[2] - max_seqlen_k = k.shape[2] - - out, lse, *_ = flash_attn_3_forward( + out, lse, *_ = flash_attn_3_func( q=q, k=k, v=v, - k_new=None, - v_new=None, + softmax_scale=softmax_scale, + causal=causal, qv=qv, - out=None, - cu_seqlens_q=None, - cu_seqlens_k=None, - cu_seqlens_k_new=None, - seqused_q=None, - seqused_k=None, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - page_table=None, - kv_batch_idx=None, - leftpad_k=None, - rotary_cos=None, - rotary_sin=None, - seqlens_rotary=None, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, - softmax_scale=softmax_scale, - causal=causal, window_size=window_size, attention_chunk=attention_chunk, softcap=softcap, - rotary_interleaved=True, - scheduler_metadata=None, num_splits=num_splits, pack_gqa=pack_gqa, + deterministic=deterministic, sm_margin=sm_margin, ) lse = lse.permute(0, 2, 1) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index d4cd574a04f7..d33d72d9c735 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -39,14 +39,9 @@ >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) >>> pipe.to("cuda") - - >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. - >>> # (1) Use flash attention 2 - >>> # pipe.transformer.set_attention_backend("flash") - >>> # (2) Use flash attention 3 - >>> # pipe.transformer.set_attention_backend("_flash_3") - >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. >>> image = pipe( ... prompt, ... height=1024, From 8e391b773543dbbbd01979035e8f1bf2465b9053 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 13:10:35 +0800 Subject: [PATCH 17/31] update z-image docstring --- src/diffusers/pipelines/z_image/pipeline_z_image.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index d33d72d9c735..d4cd574a04f7 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -39,9 +39,14 @@ >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) >>> pipe.to("cuda") + + >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. + >>> # (1) Use flash attention 2 + >>> # pipe.transformer.set_attention_backend("flash") + >>> # (2) Use flash attention 3 + >>> # pipe.transformer.set_attention_backend("_flash_3") + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" - >>> # Depending on the variant being used, the pipeline call will slightly vary. - >>> # Refer to the pipeline documentation for more details. >>> image = pipe( ... prompt, ... height=1024, From 3b22e84b628d59b0ca40d0cf17f2b348e528c59d Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 13:08:57 +0800 Subject: [PATCH 18/31] Revert attention dispatcher --- src/diffusers/models/attention_dispatch.py | 155 ++++-------------- .../pipelines/z_image/pipeline_z_image.py | 10 +- 2 files changed, 32 insertions(+), 133 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 74994bd2c62c..8504504981a3 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -18,17 +18,7 @@ import math from dataclasses import dataclass from enum import Enum -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Literal, - Optional, - Tuple, - Union, -) +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch @@ -78,10 +68,7 @@ if _CAN_USE_FLASH_ATTN: from flash_attn import flash_attn_func, flash_attn_varlen_func - from flash_attn.flash_attn_interface import ( - _wrapped_flash_attn_backward, - _wrapped_flash_attn_forward, - ) + from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward else: flash_attn_func = None flash_attn_varlen_func = None @@ -90,13 +77,11 @@ if _CAN_USE_FLASH_ATTN_3: - from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward from flash_attn_interface import flash_attn_func as flash_attn_3_func from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func else: flash_attn_3_func = None flash_attn_3_varlen_func = None - flash_attn_3_forward = None if _CAN_USE_AITER_ATTN: from aiter import flash_attn_func as aiter_flash_attn_func @@ -135,9 +120,7 @@ if _CAN_USE_XLA_ATTN: - from torch_xla.experimental.custom_kernel import ( - flash_attention as xla_flash_attention, - ) + from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention else: xla_flash_attention = None @@ -280,17 +263,13 @@ class _HubKernelConfig: _HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = { # TODO: temporary revision for now. Remove when merged upstream into `main`. AttentionBackendName._FLASH_3_HUB: _HubKernelConfig( - repo_id="kernels-community/flash-attn3", - function_attr="flash_attn_func", - revision="fake-ops-return-probs", + repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs" ) } @contextlib.contextmanager -def attention_backend( - backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE, -): +def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE): """ Context manager to set the active attention backend. """ @@ -435,10 +414,7 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None f"Flash Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `flash-attn>={_REQUIRED_FLASH_VERSION}`." ) - elif backend in [ - AttentionBackendName._FLASH_3, - AttentionBackendName._FLASH_VARLEN_3, - ]: + elif backend in [AttentionBackendName._FLASH_3, AttentionBackendName._FLASH_VARLEN_3]: if not _CAN_USE_FLASH_ATTN_3: raise RuntimeError( f"Flash Attention 3 backend '{backend.value}' is not usable because of missing package or the version is too old. Please build FA3 beta release from source." @@ -510,11 +486,7 @@ def _prepare_for_flash_attn_or_sage_varlen_without_mask( cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) max_seqlen_q = seqlens_q.max().item() max_seqlen_k = seqlens_k.max().item() - return ( - (seqlens_q, seqlens_k), - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_q, max_seqlen_k), - ) + return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) def _prepare_for_flash_attn_or_sage_varlen_with_mask( @@ -531,11 +503,7 @@ def _prepare_for_flash_attn_or_sage_varlen_with_mask( cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) max_seqlen_q = seqlens_q.max().item() max_seqlen_k = seqlens_k.max().item() - return ( - (seqlens_q, seqlens_k), - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_q, max_seqlen_k), - ) + return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) def _prepare_for_flash_attn_or_sage_varlen( @@ -653,42 +621,22 @@ def _wrapped_flash_attn_3( ) -> Tuple[torch.Tensor, torch.Tensor]: # Hardcoded for now because pytorch does not support tuple/int type hints window_size = (-1, -1) - max_seqlen_q = q.shape[2] - max_seqlen_k = k.shape[2] - - out, lse, *_ = flash_attn_3_forward( + out, lse, *_ = flash_attn_3_func( q=q, k=k, v=v, - k_new=None, - v_new=None, + softmax_scale=softmax_scale, + causal=causal, qv=qv, - out=None, - cu_seqlens_q=None, - cu_seqlens_k=None, - cu_seqlens_k_new=None, - seqused_q=None, - seqused_k=None, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - page_table=None, - kv_batch_idx=None, - leftpad_k=None, - rotary_cos=None, - rotary_sin=None, - seqlens_rotary=None, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, - softmax_scale=softmax_scale, - causal=causal, window_size=window_size, attention_chunk=attention_chunk, softcap=softcap, - rotary_interleaved=True, - scheduler_metadata=None, num_splits=num_splits, pack_gqa=pack_gqa, + deterministic=deterministic, sm_margin=sm_margin, ) lse = lse.permute(0, 2, 1) @@ -794,10 +742,7 @@ def _native_attention_backward_op( grad_out_t = grad_out.permute(0, 2, 1, 3) grad_query_t, grad_key_t, grad_value_t = torch.autograd.grad( - outputs=out, - inputs=[query_t, key_t, value_t], - grad_outputs=grad_out_t, - retain_graph=False, + outputs=out, inputs=[query_t, key_t, value_t], grad_outputs=grad_out_t, retain_graph=False ) grad_query = grad_query_t.permute(0, 2, 1, 3) @@ -836,26 +781,18 @@ def _cudnn_attention_forward_op( value = value.transpose(1, 2).contiguous() tensors_to_save += (query, key, value) - ( - out, - lse, - cum_seq_q, - cum_seq_k, - max_q, - max_k, - philox_seed, - philox_offset, - debug_attn_mask, - ) = torch.ops.aten._scaled_dot_product_cudnn_attention( - query=query, - key=key, - value=value, - attn_bias=attn_mask, - compute_log_sumexp=return_lse, - dropout_p=dropout_p, - is_causal=is_causal, - return_debug_mask=False, - scale=scale, + out, lse, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask = ( + torch.ops.aten._scaled_dot_product_cudnn_attention( + query=query, + key=key, + value=value, + attn_bias=attn_mask, + compute_log_sumexp=return_lse, + dropout_p=dropout_p, + is_causal=is_causal, + return_debug_mask=False, + scale=scale, + ) ) tensors_to_save += (out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset) @@ -982,11 +919,7 @@ def _flash_attention_backward_op( **kwargs, ): query, key, value, out, lse, rng_state = ctx.saved_tensors - grad_query, grad_key, grad_value = ( - torch.empty_like(query), - torch.empty_like(key), - torch.empty_like(value), - ) + grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) lse_d = _wrapped_flash_attn_backward( # noqa: F841 grad_out, @@ -1210,19 +1143,7 @@ def backward( grad_query, grad_key, grad_value = (x.to(grad_out.dtype) for x in (grad_query, grad_key, grad_value)) - return ( - grad_query, - grad_key, - grad_value, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None class TemplatedUlyssesAttention(torch.autograd.Function): @@ -1317,19 +1238,7 @@ def backward( x.flatten(0, 1).permute(1, 2, 0, 3).contiguous() for x in (grad_query, grad_key, grad_value) ) - return ( - grad_query, - grad_key, - grad_value, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None def _templated_context_parallel_attention( @@ -1677,12 +1586,7 @@ def _native_flex_attention( block_mask = attn_mask elif is_causal: block_mask = flex_attention.create_block_mask( - _flex_attention_causal_mask_mod, - batch_size, - num_heads, - seq_len_q, - seq_len_kv, - query.device, + _flex_attention_causal_mask_mod, batch_size, num_heads, seq_len_q, seq_len_kv, query.device ) elif torch.is_tensor(attn_mask): if attn_mask.ndim == 2: @@ -1702,7 +1606,6 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx] - else: raise ValueError("Attention mask must be either None, a BlockMask, or a 2D/4D tensor.") diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 3c334facc68b..d33d72d9c735 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -39,13 +39,9 @@ >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) >>> pipe.to("cuda") - - >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. - >>> # (1) Use flash attention 2 - >>> # pipe.transformer.set_attention_backend("flash") - >>> # (2) Use flash attention 3 - >>> # pipe.transformer.set_attention_backend("_flash_3") - >>> prompt = '一幅为名为"造相「Z-IMAGE-TURBO」"的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。' + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. >>> image = pipe( ... prompt, ... height=1024, From 3d1a7aa34ea9b8aeff73ad76c4a574502e4765f1 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 13:10:35 +0800 Subject: [PATCH 19/31] update z-image docstring --- src/diffusers/pipelines/z_image/pipeline_z_image.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index d33d72d9c735..d4cd574a04f7 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -39,9 +39,14 @@ >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) >>> pipe.to("cuda") + + >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. + >>> # (1) Use flash attention 2 + >>> # pipe.transformer.set_attention_backend("flash") + >>> # (2) Use flash attention 3 + >>> # pipe.transformer.set_attention_backend("_flash_3") + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" - >>> # Depending on the variant being used, the pipeline call will slightly vary. - >>> # Refer to the pipeline documentation for more details. >>> image = pipe( ... prompt, ... height=1024, From 336c5cee44d94762cc6935007fded0ea853fae4b Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 13:43:52 +0800 Subject: [PATCH 20/31] styling --- src/diffusers/pipelines/z_image/pipeline_z_image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index d4cd574a04f7..e5aa49c28e55 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -45,7 +45,7 @@ >>> # pipe.transformer.set_attention_backend("flash") >>> # (2) Use flash attention 3 >>> # pipe.transformer.set_attention_backend("_flash_3") - + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" >>> image = pipe( ... prompt, From 38a89edee67b232be9242da9d109ebf50650e8ea Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 05:58:43 +0000 Subject: [PATCH 21/31] Recover attention_dispatch.py with its origin impl, later would special commit for fa3 compatibility. --- src/diffusers/models/attention_dispatch.py | 155 ++++----------------- 1 file changed, 29 insertions(+), 126 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 74994bd2c62c..8504504981a3 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -18,17 +18,7 @@ import math from dataclasses import dataclass from enum import Enum -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Literal, - Optional, - Tuple, - Union, -) +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch @@ -78,10 +68,7 @@ if _CAN_USE_FLASH_ATTN: from flash_attn import flash_attn_func, flash_attn_varlen_func - from flash_attn.flash_attn_interface import ( - _wrapped_flash_attn_backward, - _wrapped_flash_attn_forward, - ) + from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward else: flash_attn_func = None flash_attn_varlen_func = None @@ -90,13 +77,11 @@ if _CAN_USE_FLASH_ATTN_3: - from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward from flash_attn_interface import flash_attn_func as flash_attn_3_func from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func else: flash_attn_3_func = None flash_attn_3_varlen_func = None - flash_attn_3_forward = None if _CAN_USE_AITER_ATTN: from aiter import flash_attn_func as aiter_flash_attn_func @@ -135,9 +120,7 @@ if _CAN_USE_XLA_ATTN: - from torch_xla.experimental.custom_kernel import ( - flash_attention as xla_flash_attention, - ) + from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention else: xla_flash_attention = None @@ -280,17 +263,13 @@ class _HubKernelConfig: _HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = { # TODO: temporary revision for now. Remove when merged upstream into `main`. AttentionBackendName._FLASH_3_HUB: _HubKernelConfig( - repo_id="kernels-community/flash-attn3", - function_attr="flash_attn_func", - revision="fake-ops-return-probs", + repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs" ) } @contextlib.contextmanager -def attention_backend( - backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE, -): +def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE): """ Context manager to set the active attention backend. """ @@ -435,10 +414,7 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None f"Flash Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `flash-attn>={_REQUIRED_FLASH_VERSION}`." ) - elif backend in [ - AttentionBackendName._FLASH_3, - AttentionBackendName._FLASH_VARLEN_3, - ]: + elif backend in [AttentionBackendName._FLASH_3, AttentionBackendName._FLASH_VARLEN_3]: if not _CAN_USE_FLASH_ATTN_3: raise RuntimeError( f"Flash Attention 3 backend '{backend.value}' is not usable because of missing package or the version is too old. Please build FA3 beta release from source." @@ -510,11 +486,7 @@ def _prepare_for_flash_attn_or_sage_varlen_without_mask( cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) max_seqlen_q = seqlens_q.max().item() max_seqlen_k = seqlens_k.max().item() - return ( - (seqlens_q, seqlens_k), - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_q, max_seqlen_k), - ) + return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) def _prepare_for_flash_attn_or_sage_varlen_with_mask( @@ -531,11 +503,7 @@ def _prepare_for_flash_attn_or_sage_varlen_with_mask( cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) max_seqlen_q = seqlens_q.max().item() max_seqlen_k = seqlens_k.max().item() - return ( - (seqlens_q, seqlens_k), - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_q, max_seqlen_k), - ) + return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) def _prepare_for_flash_attn_or_sage_varlen( @@ -653,42 +621,22 @@ def _wrapped_flash_attn_3( ) -> Tuple[torch.Tensor, torch.Tensor]: # Hardcoded for now because pytorch does not support tuple/int type hints window_size = (-1, -1) - max_seqlen_q = q.shape[2] - max_seqlen_k = k.shape[2] - - out, lse, *_ = flash_attn_3_forward( + out, lse, *_ = flash_attn_3_func( q=q, k=k, v=v, - k_new=None, - v_new=None, + softmax_scale=softmax_scale, + causal=causal, qv=qv, - out=None, - cu_seqlens_q=None, - cu_seqlens_k=None, - cu_seqlens_k_new=None, - seqused_q=None, - seqused_k=None, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - page_table=None, - kv_batch_idx=None, - leftpad_k=None, - rotary_cos=None, - rotary_sin=None, - seqlens_rotary=None, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, - softmax_scale=softmax_scale, - causal=causal, window_size=window_size, attention_chunk=attention_chunk, softcap=softcap, - rotary_interleaved=True, - scheduler_metadata=None, num_splits=num_splits, pack_gqa=pack_gqa, + deterministic=deterministic, sm_margin=sm_margin, ) lse = lse.permute(0, 2, 1) @@ -794,10 +742,7 @@ def _native_attention_backward_op( grad_out_t = grad_out.permute(0, 2, 1, 3) grad_query_t, grad_key_t, grad_value_t = torch.autograd.grad( - outputs=out, - inputs=[query_t, key_t, value_t], - grad_outputs=grad_out_t, - retain_graph=False, + outputs=out, inputs=[query_t, key_t, value_t], grad_outputs=grad_out_t, retain_graph=False ) grad_query = grad_query_t.permute(0, 2, 1, 3) @@ -836,26 +781,18 @@ def _cudnn_attention_forward_op( value = value.transpose(1, 2).contiguous() tensors_to_save += (query, key, value) - ( - out, - lse, - cum_seq_q, - cum_seq_k, - max_q, - max_k, - philox_seed, - philox_offset, - debug_attn_mask, - ) = torch.ops.aten._scaled_dot_product_cudnn_attention( - query=query, - key=key, - value=value, - attn_bias=attn_mask, - compute_log_sumexp=return_lse, - dropout_p=dropout_p, - is_causal=is_causal, - return_debug_mask=False, - scale=scale, + out, lse, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask = ( + torch.ops.aten._scaled_dot_product_cudnn_attention( + query=query, + key=key, + value=value, + attn_bias=attn_mask, + compute_log_sumexp=return_lse, + dropout_p=dropout_p, + is_causal=is_causal, + return_debug_mask=False, + scale=scale, + ) ) tensors_to_save += (out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset) @@ -982,11 +919,7 @@ def _flash_attention_backward_op( **kwargs, ): query, key, value, out, lse, rng_state = ctx.saved_tensors - grad_query, grad_key, grad_value = ( - torch.empty_like(query), - torch.empty_like(key), - torch.empty_like(value), - ) + grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) lse_d = _wrapped_flash_attn_backward( # noqa: F841 grad_out, @@ -1210,19 +1143,7 @@ def backward( grad_query, grad_key, grad_value = (x.to(grad_out.dtype) for x in (grad_query, grad_key, grad_value)) - return ( - grad_query, - grad_key, - grad_value, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None class TemplatedUlyssesAttention(torch.autograd.Function): @@ -1317,19 +1238,7 @@ def backward( x.flatten(0, 1).permute(1, 2, 0, 3).contiguous() for x in (grad_query, grad_key, grad_value) ) - return ( - grad_query, - grad_key, - grad_value, - None, - None, - None, - None, - None, - None, - None, - None, - ) + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None def _templated_context_parallel_attention( @@ -1677,12 +1586,7 @@ def _native_flex_attention( block_mask = attn_mask elif is_causal: block_mask = flex_attention.create_block_mask( - _flex_attention_causal_mask_mod, - batch_size, - num_heads, - seq_len_q, - seq_len_kv, - query.device, + _flex_attention_causal_mask_mod, batch_size, num_heads, seq_len_q, seq_len_kv, query.device ) elif torch.is_tensor(attn_mask): if attn_mask.ndim == 2: @@ -1702,7 +1606,6 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx] - else: raise ValueError("Attention mask must be either None, a BlockMask, or a 2D/4D tensor.") From 69d61e5074f78c48ff33b10786be2c9ee29a1f31 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 06:05:35 +0000 Subject: [PATCH 22/31] Fix prev bug, and support for prompt_embeds pass in args after prompt pre-encode as List of torch Tensor. --- .../pipelines/z_image/pipeline_z_image.py | 41 +++++++++++-------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 3c334facc68b..90a558581cf0 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -431,26 +431,35 @@ def __call__( elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: - batch_size = prompt_embeds.shape[0] + batch_size = len(prompt_embeds) lora_scale = ( self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None ) - ( - prompt_embeds, - negative_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - dtype=dtype, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) + + # If prompt_embeds is provided and prompt is None, skip encoding + if prompt_embeds is not None and prompt is None: + if self.do_classifier_free_guidance and negative_prompt_embeds is None: + raise ValueError( + "When `prompt_embeds` is provided without `prompt`, " + "`negative_prompt_embeds` must also be provided for classifier-free guidance." + ) + else: + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + dtype=dtype, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) # 4. Prepare latent variables num_channels_latents = self.transformer.in_channels From 1dd8f3cfe3e261343f42c386115d2184c489a6fd Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 06:43:11 +0000 Subject: [PATCH 23/31] Remove einop dependency. --- .../models/transformers/transformer_z_image.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 24c3189a8213..0b20c30fa373 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -18,7 +18,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from einops import rearrange from torch.nn.utils.rnn import pad_sequence from ...configuration_utils import ConfigMixin, register_to_config @@ -429,9 +428,12 @@ def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_pat for i in range(bsz): F, H, W = size[i] ori_len = (F // pF) * (H // pH) * (W // pW) - x[i] = rearrange( - x[i][:ori_len].view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels), - "f h w pf ph pw c -> c (f pf) (h ph) (w pw)", + # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)" + x[i] = ( + x[i][:ori_len] + .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels) + .permute(6, 0, 3, 1, 4, 2, 5) + .reshape(self.out_channels, F, H, W) ) return x @@ -497,7 +499,8 @@ def patchify_and_embed( F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) - image = rearrange(image, "c f pf h ph w pw -> (f h w) (pf ph pw c)") + # "c f pf h ph w pw -> (f h w) (pf ph pw c)" + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C) image_ori_len = len(image) image_padding_len = (-image_ori_len) % SEQ_MULTI_OF From e49a1f901a036737dbd7b5b930fe202cfd3b0fe7 Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 21:15:59 +0800 Subject: [PATCH 24/31] remove redundant imports & make fix-copies --- .../models/transformers/transformer_z_image.py | 11 ++--------- .../utils/dummy_torch_and_transformers_objects.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 0b20c30fa373..b10644166120 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -20,23 +20,16 @@ import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence +from diffusers.models.normalization import RMSNorm + from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin -from ...utils.import_utils import is_flash_attn_available from ...utils.torch_utils import maybe_allow_in_graph from ..attention_dispatch import dispatch_attention_fn -if is_flash_attn_available(): - from flash_attn import flash_attn_varlen_func -else: - flash_attn_varlen_func = None - -from diffusers.models.normalization import RMSNorm - - ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 9eb123b94e9d..ec68c52c4d9b 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3645,3 +3645,18 @@ def from_config(cls, *args, **kwargs): @classmethod def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) + + +class ZImagePipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) From 1048d0a94d553978b35bec0942c408c57bb12c9e Mon Sep 17 00:00:00 2001 From: liudongyang Date: Tue, 25 Nov 2025 21:18:41 +0800 Subject: [PATCH 25/31] fix import --- src/diffusers/models/transformers/transformer_z_image.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index b10644166120..a5c1de682a74 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -20,12 +20,11 @@ import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence -from diffusers.models.normalization import RMSNorm - from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin +from ...models.normalization import RMSNorm from ...utils.torch_utils import maybe_allow_in_graph from ..attention_dispatch import dispatch_attention_fn From 266e169a9cb0746f68c59394974ffb7cac7904f4 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 15:38:37 +0000 Subject: [PATCH 26/31] Support for num_images_per_prompt>1; Remove redundant unquote variables. --- .../pipelines/z_image/pipeline_z_image.py | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index cc4e9d52019b..6cfe8b3f32fe 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -165,21 +165,16 @@ def encode_prompt( self, prompt: Union[str, List[str]], device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, negative_prompt: Optional[Union[str, List[str]]] = None, prompt_embeds: Optional[List[torch.FloatTensor]] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt prompt_embeds = self._encode_prompt( prompt=prompt, device=device, - dtype=dtype, - num_images_per_prompt=num_images_per_prompt, prompt_embeds=prompt_embeds, max_sequence_length=max_sequence_length, ) @@ -193,8 +188,6 @@ def encode_prompt( negative_prompt_embeds = self._encode_prompt( prompt=negative_prompt, device=device, - dtype=dtype, - num_images_per_prompt=num_images_per_prompt, prompt_embeds=negative_prompt_embeds, max_sequence_length=max_sequence_length, ) @@ -206,12 +199,9 @@ def _encode_prompt( self, prompt: Union[str, List[str]], device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - num_images_per_prompt: int = 1, prompt_embeds: Optional[List[torch.FloatTensor]] = None, max_sequence_length: int = 512, ) -> List[torch.FloatTensor]: - assert num_images_per_prompt == 1 device = device or self._execution_device if prompt_embeds is not None: @@ -417,8 +407,8 @@ def __call__( f"Please adjust the width to a multiple of {vae_scale}." ) - assert self.dtype == torch.bfloat16 - dtype = self.dtype + # assert self.dtype == torch.bfloat16 + dtype = self.dtype if hasattr(self, "dtype") and self.dtype is not None else torch.float32 device = self._execution_device self._guidance_scale = guidance_scale @@ -434,10 +424,6 @@ def __call__( else: batch_size = len(prompt_embeds) - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - # If prompt_embeds is provided and prompt is None, skip encoding if prompt_embeds is not None and prompt is None: if self.do_classifier_free_guidance and negative_prompt_embeds is None: @@ -455,11 +441,8 @@ def __call__( do_classifier_free_guidance=self.do_classifier_free_guidance, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, - dtype=dtype, device=device, - num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, - lora_scale=lora_scale, ) # 4. Prepare latent variables @@ -475,6 +458,12 @@ def __call__( generator, latents, ) + + # Repeat prompt_embeds for num_images_per_prompt + if num_images_per_prompt > 1: + prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] + if self.do_classifier_free_guidance and negative_prompt_embeds: + negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) # 5. Prepare timesteps From 12d2fb283bae1b7cfd5e0b26e0d834e0b3d0f8c7 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 21:09:55 +0000 Subject: [PATCH 27/31] Fix bugs for num_images_per_prompt with actual batch. --- src/diffusers/pipelines/z_image/pipeline_z_image.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 6cfe8b3f32fe..e64f10f40cd9 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -464,6 +464,8 @@ def __call__( prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] if self.do_classifier_free_guidance and negative_prompt_embeds: negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] + + actual_batch_size = batch_size * num_images_per_prompt image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) # 5. Prepare timesteps @@ -532,11 +534,11 @@ def __call__( if apply_cfg: # Perform CFG - pos_out = model_out_list[:batch_size] - neg_out = model_out_list[batch_size:] + pos_out = model_out_list[:actual_batch_size] + neg_out = model_out_list[actual_batch_size:] noise_pred = [] - for j in range(batch_size): + for j in range(actual_batch_size): pos = pos_out[j].float() neg = neg_out[j].float() From 9a049f0cce918680257625a9f79b4840431b57b9 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 21:13:19 +0000 Subject: [PATCH 28/31] Add unit tests for Z-Image. --- .../pipelines/z_image/pipeline_z_image.py | 2 +- tests/pipelines/z_image/test_z_image.py | 302 ++++++++++++++++++ 2 files changed, 303 insertions(+), 1 deletion(-) create mode 100644 tests/pipelines/z_image/test_z_image.py diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index e64f10f40cd9..e415582d5f83 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -464,7 +464,7 @@ def __call__( prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] if self.do_classifier_free_guidance and negative_prompt_embeds: negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] - + actual_batch_size = batch_size * num_images_per_prompt image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) diff --git a/tests/pipelines/z_image/test_z_image.py b/tests/pipelines/z_image/test_z_image.py new file mode 100644 index 000000000000..25098408bc21 --- /dev/null +++ b/tests/pipelines/z_image/test_z_image.py @@ -0,0 +1,302 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import os +import unittest + +import numpy as np +import torch +from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model + +from diffusers import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + ZImagePipeline, + ZImageTransformer2DModel, +) + +from ...testing_utils import torch_device +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin, to_np + + +# Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations +# Cannot use enable_full_determinism() which sets it to True +os.environ["CUDA_LAUNCH_BLOCKING"] = "1" +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" +torch.use_deterministic_algorithms(False) +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False +if hasattr(torch.backends, "cuda"): + torch.backends.cuda.matmul.allow_tf32 = False + + +class ZImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = ZImagePipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + def setUp(self): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) + + def tearDown(self): + super().tearDown() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = ZImageTransformer2DModel( + all_patch_size=(2,), + all_f_patch_size=(1,), + in_channels=16, + dim=32, + n_layers=2, + n_refiner_layers=1, + n_heads=2, + n_kv_heads=2, + norm_eps=1e-5, + qk_norm=True, + cap_feat_dim=16, + rope_theta=256.0, + t_scale=1000.0, + axes_dims=[8, 4, 4], + axes_lens=[256, 32, 32], + ) + + torch.manual_seed(0) + vae = AutoencoderKL( + in_channels=3, + out_channels=3, + down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], + up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], + block_out_channels=[32, 64], + layers_per_block=1, + latent_channels=16, + norm_num_groups=32, + sample_size=32, + scaling_factor=0.3611, + shift_factor=0.1159, + ) + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler() + + torch.manual_seed(0) + config = Qwen3Config( + hidden_size=16, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + vocab_size=151936, + max_position_embeddings=512, + ) + text_encoder = Qwen3Model(config) + tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + inputs = { + "prompt": "dance monkey", + "negative_prompt": "bad quality", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "cfg_normalization": False, + "cfg_truncation": 1.0, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "pt", + } + + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (3, 32, 32)) + + # fmt: off + expected_slice = torch.tensor([0.4521, 0.4512, 0.4693, 0.5115, 0.5250, 0.5271, 0.4776, 0.4688, 0.2765, 0.2164, 0.5656, 0.6909, 0.3831, 0.5431, 0.5493, 0.4732]) + # fmt: on + + generated_slice = generated_image.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-2)) + + def test_inference_batch_single_identical(self): + self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) + + def test_num_images_per_prompt(self): + import inspect + + sig = inspect.signature(self.pipeline_class.__call__) + + if "num_images_per_prompt" not in sig.parameters: + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + batch_sizes = [1, 2] + num_images_per_prompts = [1, 2] + + for batch_size in batch_sizes: + for num_images_per_prompt in num_images_per_prompts: + inputs = self.get_dummy_inputs(torch_device) + + for key in inputs.keys(): + if key in self.batch_params: + inputs[key] = batch_size * [inputs[key]] + + images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0] + + assert images.shape[0] == batch_size * num_images_per_prompt + + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def test_attention_slicing_forward_pass( + self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 + ): + if not self.test_attention_slicing: + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + generator_device = "cpu" + inputs = self.get_dummy_inputs(generator_device) + output_without_slicing = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=1) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing1 = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=2) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing2 = pipe(**inputs)[0] + + if test_max_difference: + max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() + max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() + self.assertLess( + max(max_diff1, max_diff2), + expected_max_diff, + "Attention slicing should not affect the inference results", + ) + + def test_vae_tiling(self, expected_diff_max: float = 0.2): + generator_device = "cpu" + components = self.get_dummy_components() + + pipe = self.pipeline_class(**components) + pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + + # Without tiling + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + output_without_tiling = pipe(**inputs)[0] + + # With tiling (standard AutoencoderKL doesn't accept parameters) + pipe.vae.enable_tiling() + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + output_with_tiling = pipe(**inputs)[0] + + self.assertLess( + (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), + expected_diff_max, + "VAE tiling should not affect the inference results", + ) + + def test_pipeline_with_accelerator_device_map(self, expected_max_difference=5e-4): + # Z-Image RoPE embeddings (complex64) have slightly higher numerical tolerance + super().test_pipeline_with_accelerator_device_map(expected_max_difference=expected_max_difference) + + def test_group_offloading_inference(self): + # Block-level offloading conflicts with RoPE cache. Pipeline-level offloading (tested separately) works fine. + self.skipTest("Using test_pipeline_level_group_offloading_inference instead") + + def test_save_load_float16(self, expected_max_diff=1e-2): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) + super().test_save_load_float16(expected_max_diff=expected_max_diff) From c4e4a57ef825862bf0fc611bdd25e41e1f1237c0 Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 21:22:59 +0000 Subject: [PATCH 29/31] Refine unitest and skip for cases needed separate test env; Fix compatibility with unitest in model, mostly precision formating. --- .../transformers/transformer_z_image.py | 19 ++++++++++++++++--- tests/pipelines/z_image/test_z_image.py | 8 +------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index a5c1de682a74..d2dec0b2c5cd 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -69,7 +69,10 @@ def timestep_embedding(t, dim, max_period=10000): def forward(self, t): t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + weight_dtype = self.mlp[0].weight.dtype + if weight_dtype in [torch.float32, torch.float16, torch.bfloat16]: + t_freq = t_freq.to(weight_dtype) + t_emb = self.mlp(t_freq) return t_emb @@ -126,6 +129,10 @@ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tenso dtype = query.dtype query, key = query.to(dtype), key.to(dtype) + # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len] + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask[:, None, None, :] + # Compute joint attention hidden_states = dispatch_attention_fn( query, @@ -306,6 +313,10 @@ def __call__(self, ids: torch.Tensor): if self.freqs_cis is None: self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + else: + # Ensure freqs_cis are on the same device as ids + if self.freqs_cis[0].device != device: + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] result = [] for i in range(len(self.axes_dims)): @@ -317,6 +328,7 @@ def __call__(self, ids: torch.Tensor): class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): _supports_gradient_checkpointing = True _no_split_modules = ["ZImageTransformerBlock"] + _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"] # precision sensitive layers @register_to_config def __init__( @@ -553,8 +565,6 @@ def forward( t = t * self.t_scale t = self.t_embedder(t) - adaln_input = t - ( x, cap_feats, @@ -572,6 +582,9 @@ def forward( x = torch.cat(x, dim=0) x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) + + # Match t_embedder output dtype to x for layerwise casting compatibility + adaln_input = t.type_as(x) x[torch.cat(x_inner_pad_mask)] = self.x_pad_token x = list(x.split(x_item_seqlens, dim=0)) x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) diff --git a/tests/pipelines/z_image/test_z_image.py b/tests/pipelines/z_image/test_z_image.py index 25098408bc21..dae7c931f870 100644 --- a/tests/pipelines/z_image/test_z_image.py +++ b/tests/pipelines/z_image/test_z_image.py @@ -291,12 +291,6 @@ def test_group_offloading_inference(self): # Block-level offloading conflicts with RoPE cache. Pipeline-level offloading (tested separately) works fine. self.skipTest("Using test_pipeline_level_group_offloading_inference instead") + @unittest.skip("Known issue: fails in full suite due to test isolation (passes when run individually)") def test_save_load_float16(self, expected_max_diff=1e-2): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - torch.manual_seed(0) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) super().test_save_load_float16(expected_max_diff=expected_max_diff) From 6f2808b6f6a76229dfc9abb86344eede04b5574b Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Tue, 25 Nov 2025 21:31:22 +0000 Subject: [PATCH 30/31] Add clean env for test_save_load_float16 separ test; Add Note; Styling. --- tests/pipelines/z_image/test_z_image.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/pipelines/z_image/test_z_image.py b/tests/pipelines/z_image/test_z_image.py index dae7c931f870..709473b0dbb8 100644 --- a/tests/pipelines/z_image/test_z_image.py +++ b/tests/pipelines/z_image/test_z_image.py @@ -42,6 +42,10 @@ if hasattr(torch.backends, "cuda"): torch.backends.cuda.matmul.allow_tf32 = False +# Note: Some tests (test_float16_inference, test_save_load_float16) may fail in full suite +# due to RopeEmbedder cache state pollution between tests. They pass when run individually. +# This is a known test isolation issue, not a functional bug. + class ZImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = ZImagePipeline @@ -291,6 +295,12 @@ def test_group_offloading_inference(self): # Block-level offloading conflicts with RoPE cache. Pipeline-level offloading (tested separately) works fine. self.skipTest("Using test_pipeline_level_group_offloading_inference instead") - @unittest.skip("Known issue: fails in full suite due to test isolation (passes when run individually)") def test_save_load_float16(self, expected_max_diff=1e-2): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) super().test_save_load_float16(expected_max_diff=expected_max_diff) From aeed8904c1981b16e599f5cb0809fd43c5abd35a Mon Sep 17 00:00:00 2001 From: Jerry Qilong Wu Date: Wed, 26 Nov 2025 09:15:55 +0000 Subject: [PATCH 31/31] Update dtype mentioned by yiyi. --- src/diffusers/models/transformers/transformer_z_image.py | 2 +- src/diffusers/pipelines/z_image/pipeline_z_image.py | 8 +++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index d2dec0b2c5cd..3ad835ceeeb0 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -70,7 +70,7 @@ def timestep_embedding(t, dim, max_period=10000): def forward(self, t): t_freq = self.timestep_embedding(t, self.frequency_embedding_size) weight_dtype = self.mlp[0].weight.dtype - if weight_dtype in [torch.float32, torch.float16, torch.bfloat16]: + if weight_dtype.is_floating_point: t_freq = t_freq.to(weight_dtype) t_emb = self.mlp(t_freq) return t_emb diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index e415582d5f83..a4fcacb6eb9b 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -407,8 +407,6 @@ def __call__( f"Please adjust the width to a multiple of {vae_scale}." ) - # assert self.dtype == torch.bfloat16 - dtype = self.dtype if hasattr(self, "dtype") and self.dtype is not None else torch.float32 device = self._execution_device self._guidance_scale = guidance_scale @@ -514,12 +512,12 @@ def __call__( apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 if apply_cfg: - latents_typed = latents if latents.dtype == dtype else latents.to(dtype) + latents_typed = latents.to(self.transformer.dtype) latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds timestep_model_input = timestep.repeat(2) else: - latent_model_input = latents if latents.dtype == dtype else latents.to(dtype) + latent_model_input = latents.to(self.transformer.dtype) prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep @@ -579,11 +577,11 @@ def __call__( if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() - latents = latents.to(dtype) if output_type == "latent": image = latents else: + latents = latents.to(self.vae.dtype) latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor image = self.vae.decode(latents, return_dict=False)[0]