From 7f3fa8a99547dd2f076eda02f01d8c780a0286f7 Mon Sep 17 00:00:00 2001 From: naykun Date: Mon, 14 Sep 2026 12:55:34 +0800 Subject: [PATCH 01/20] feat: add Qwen-Image 2.1 pipeline with block-causal attention and KV cache New classes: - QwenImage21Transformer2DModel: single-stream transformer with block-causal attention and t=0 modulation for the text and condition-image prefix - AutoencoderKLQwenImage21: 64-channel VAE (z_dim=64, decoder_base_dim=144) - QwenImage21Pipeline: text-to-image and image-conditioned generation Attention: - Block-causal: the joint text/image sequence is causal while each image block (condition and target) stays internally bidirectional. Built as a compiled flex_attention BlockMask, which keeps the score matrix block-sparse and makes 2048x2048 feasible. - flex_attention is optional. Without it the mask is approximated by a two-pass prefill (prefix causally, then the target image over the cached prefix). Exact for text and for the target image, approximate for condition images. KV cache: - The text and condition-image prefix is modulated from t=0, so its activations do not change across denoising steps and its keys and values are cached after the first step. Later steps only recompute the target image's tokens. Also: separate text-to-image and image-conditioned prompt templates with image-pad token downsampling, and plain classifier-free guidance. Includes model tests, docs, and full registration. --- docs/source/en/_toctree.yml | 6 + .../api/models/autoencoderkl_qwenimage21.md | 37 + .../api/models/qwenimage21_transformer2d.md | 43 + docs/source/en/api/pipelines/qwenimage21.md | 42 + src/diffusers/__init__.py | 6 + src/diffusers/models/__init__.py | 4 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoder_kl_qwenimage21.py | 1518 +++++++++++++++++ src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_qwenimage21.py | 869 ++++++++++ src/diffusers/pipelines/__init__.py | 2 + .../pipelines/qwenimage21/__init__.py | 48 + .../qwenimage21/pipeline_qwenimage21.py | 804 +++++++++ src/diffusers/utils/dummy_pt_objects.py | 30 + .../dummy_torch_and_transformers_objects.py | 15 + .../test_models_transformer_qwenimage21.py | 235 +++ 16 files changed, 3661 insertions(+) create mode 100644 docs/source/en/api/models/autoencoderkl_qwenimage21.md create mode 100644 docs/source/en/api/models/qwenimage21_transformer2d.md create mode 100644 docs/source/en/api/pipelines/qwenimage21.md create mode 100644 src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py create mode 100644 src/diffusers/models/transformers/transformer_qwenimage21.py create mode 100644 src/diffusers/pipelines/qwenimage21/__init__.py create mode 100644 src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py create mode 100644 tests/models/transformers/test_models_transformer_qwenimage21.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f2801a0cd6a6..3bc97543ab9b 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -405,6 +405,8 @@ title: PixArtTransformer2DModel - local: api/models/prior_transformer title: PriorTransformer + - local: api/models/qwenimage21_transformer2d + title: QwenImage21Transformer2DModel - local: api/models/qwenimage_transformer2d title: QwenImageTransformer2DModel - local: api/models/sana_transformer2d @@ -489,6 +491,8 @@ title: AutoencoderKLMochi - local: api/models/autoencoderkl_qwenimage title: AutoencoderKLQwenImage + - local: api/models/autoencoderkl_qwenimage21 + title: AutoencoderKLQwenImage21 - local: api/models/autoencoder_kl_wan title: AutoencoderKLWan - local: api/models/autoencoder_rae @@ -639,6 +643,8 @@ title: PRX Pixel - local: api/pipelines/qwenimage title: QwenImage + - local: api/pipelines/qwenimage21 + title: Qwen-Image 2.1 - local: api/pipelines/sana title: Sana - local: api/pipelines/sana_sprint diff --git a/docs/source/en/api/models/autoencoderkl_qwenimage21.md b/docs/source/en/api/models/autoencoderkl_qwenimage21.md new file mode 100644 index 000000000000..eb81c7674470 --- /dev/null +++ b/docs/source/en/api/models/autoencoderkl_qwenimage21.md @@ -0,0 +1,37 @@ + + +# AutoencoderKLQwenImage21 + +The 64-channel variational auto-encoder used by Qwen-Image 2.1. It compresses 16x spatially, and its per-channel +`latents_mean` / `latents_std` are part of the config rather than a single scaling factor. + +```python +import torch +from diffusers import AutoencoderKLQwenImage21 + +vae = AutoencoderKLQwenImage21.from_pretrained("Qwen/Qwen-Image-2.1", subfolder="vae", dtype=torch.bfloat16) +``` + +## AutoencoderKLQwenImage21 + +[[autodoc]] AutoencoderKLQwenImage21 + - decode + - encode + - all + +## AutoencoderKLOutput + +[[autodoc]] models.autoencoders.vae.AutoencoderKLOutput + +## DecoderOutput + +[[autodoc]] models.autoencoders.vae.DecoderOutput diff --git a/docs/source/en/api/models/qwenimage21_transformer2d.md b/docs/source/en/api/models/qwenimage21_transformer2d.md new file mode 100644 index 000000000000..41cc29cc41e8 --- /dev/null +++ b/docs/source/en/api/models/qwenimage21_transformer2d.md @@ -0,0 +1,43 @@ + + +# QwenImage21Transformer2DModel + +The single-stream transformer used by Qwen-Image 2.1. Text and image latents share one sequence, and a single shared +`modulation` projection feeds every block. + +Two config flags set 2.1 apart from earlier QwenImage transformers. Neither adds parameters: + +- `causal_block` — attention follows `(q_idx >= kv_idx) or same_image_block`, so the joint sequence is causal while + each image block (every condition image and the target image) stays internally bidirectional. This requires the + `flex` attention backend, since the mask is a `torch.nn.attention.flex_attention.BlockMask`. +- `causal_condition` — text and condition-image tokens are modulated from `t = 0` rather than the sampled timestep. + Their activations are therefore independent of the denoising step, which is what makes the keys and values of that + prefix cacheable across steps via the `kv_cache` argument. + +The model can be loaded with the following code snippet. + +```python +import torch +from diffusers import QwenImage21Transformer2DModel + +transformer = QwenImage21Transformer2DModel.from_pretrained( + "Qwen/Qwen-Image-2.1", subfolder="transformer", dtype=torch.bfloat16 +) +``` + +## QwenImage21Transformer2DModel + +[[autodoc]] QwenImage21Transformer2DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md new file mode 100644 index 000000000000..f18a0458420a --- /dev/null +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -0,0 +1,42 @@ + + +# Qwen-Image 2.1 + +Qwen-Image 2.1 encodes the prompt and any condition images together with a Qwen3-VL model, then denoises the target +image with a single-stream block-causal transformer. See +[`QwenImage21Transformer2DModel`](../models/qwenimage21_transformer2d) for what `causal_block` and `causal_condition` +change. + +Because the text and condition-image prefix is modulated from `t = 0`, its keys and values do not change between +denoising steps. The pipeline caches them after the first step by default; pass `use_kv_cache=False` to recompute the +full sequence every step. + +Toggling `use_kv_cache` does not reproduce the same image bit-for-bit in reduced precision. The cached decode step +attends with a different sequence layout than the prefill step, so the two land on different rounding — both match an +fp32 reference to the same tolerance — and a one-ULP difference at the first block is amplified by 32 blocks and every +sampler step. Keep the flag fixed when you need a reproducible sample. + + + +This pipeline requires the `flex` attention backend when `causal_block` is enabled. + + + +## QwenImage21Pipeline + +[[autodoc]] QwenImage21Pipeline + - all + - __call__ + +## QwenImagePipelineOutput + +[[autodoc]] pipelines.qwenimage.pipeline_output.QwenImagePipelineOutput diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 3ed956c75e49..9ab9010573db 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -253,6 +253,7 @@ "AutoencoderKLMiniMaxH3Audio", "AutoencoderKLMochi", "AutoencoderKLQwenImage", + "AutoencoderKLQwenImage21", "AutoencoderKLTemporalDecoder", "AutoencoderKLWan", "AutoencoderOobleck", @@ -334,6 +335,7 @@ "PRXTransformer2DModel", "QwenImageControlNetModel", "QwenImageMultiControlNetModel", + "QwenImage21Transformer2DModel", "QwenImageTransformer2DModel", "SanaControlNetModel", "SanaTransformer2DModel", @@ -772,6 +774,7 @@ "QwenImageImg2ImgPipeline", "QwenImageInpaintPipeline", "QwenImageLayeredPipeline", + "QwenImage21Pipeline", "QwenImagePipeline", "ReduxImageEncoder", "SanaControlNetPipeline", @@ -1130,6 +1133,7 @@ AutoencoderKLMiniMaxH3Audio, AutoencoderKLMochi, AutoencoderKLQwenImage, + AutoencoderKLQwenImage21, AutoencoderKLTemporalDecoder, AutoencoderKLWan, AutoencoderOobleck, @@ -1209,6 +1213,7 @@ PixArtTransformer2DModel, PriorTransformer, PRXTransformer2DModel, + QwenImage21Transformer2DModel, QwenImageControlNetModel, QwenImageMultiControlNetModel, QwenImageTransformer2DModel, @@ -1616,6 +1621,7 @@ PixArtSigmaPipeline, PRXPipeline, PRXPixelPipeline, + QwenImage21Pipeline, QwenImageControlNetInpaintPipeline, QwenImageControlNetPipeline, QwenImageEditInpaintPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 8ba17d896434..1a396b312441 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -50,6 +50,7 @@ _import_structure["autoencoders.autoencoder_kl_minimax_h3_audio"] = ["AutoencoderKLMiniMaxH3Audio"] _import_structure["autoencoders.autoencoder_kl_mochi"] = ["AutoencoderKLMochi"] _import_structure["autoencoders.autoencoder_kl_qwenimage"] = ["AutoencoderKLQwenImage"] + _import_structure["autoencoders.autoencoder_kl_qwenimage21"] = ["AutoencoderKLQwenImage21"] _import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"] _import_structure["autoencoders.autoencoder_kl_wan"] = ["AutoencoderKLWan"] _import_structure["autoencoders.autoencoder_longcat_audio_dit"] = ["LongCatAudioDiTVae"] @@ -144,6 +145,7 @@ _import_structure["transformers.transformer_ovis_image"] = ["OvisImageTransformer2DModel"] _import_structure["transformers.transformer_prx"] = ["PRXTransformer2DModel"] _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] + _import_structure["transformers.transformer_qwenimage21"] = ["QwenImage21Transformer2DModel"] _import_structure["transformers.transformer_sana_video"] = ["SanaVideoTransformer3DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] @@ -195,6 +197,7 @@ AutoencoderKLMiniMaxH3Audio, AutoencoderKLMochi, AutoencoderKLQwenImage, + AutoencoderKLQwenImage21, AutoencoderKLTemporalDecoder, AutoencoderKLWan, AutoencoderOobleck, @@ -287,6 +290,7 @@ PixArtTransformer2DModel, PriorTransformer, PRXTransformer2DModel, + QwenImage21Transformer2DModel, QwenImageTransformer2DModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 607704343743..218785bad2e2 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -20,6 +20,7 @@ from .autoencoder_kl_minimax_h3_audio import AutoencoderKLMiniMaxH3Audio from .autoencoder_kl_mochi import AutoencoderKLMochi from .autoencoder_kl_qwenimage import AutoencoderKLQwenImage +from .autoencoder_kl_qwenimage21 import AutoencoderKLQwenImage21 from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder from .autoencoder_kl_wan import AutoencoderKLWan from .autoencoder_longcat_audio_dit import LongCatAudioDiTVae diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py new file mode 100644 index 000000000000..096dcdfd280b --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -0,0 +1,1518 @@ +# Copyright 2026 The Qwen 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 torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...utils import logging +from ...utils.accelerate_utils import apply_forward_hook +from ..activations import get_activation +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +CACHE_T = 2 + + +class QwenImage21AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class QwenImage21DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class QwenImage21CausalConv3d(nn.Conv2d): + r""" + A custom 3D causal convolution layer with feature caching support. + + This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature + caching for efficient inference. + + Args: + in_channels (int): Number of channels in the input image + out_channels (int): Number of channels produced by the convolution + kernel_size (int or tuple): Size of the convolving kernel + stride (int or tuple, optional): Stride of the convolution. Default: 1 + padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int | int | int], + stride: int | tuple[int | int | int] = 1, + padding: int | tuple[int | int | int] = 0, + ) -> None: + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + + # Set up causal padding + self._padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0]) + self.padding = (0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + assert cache_x is None + x = x.squeeze(2) # Remove the temporal dimension + x = F.pad(x, padding) + x = super().forward(x) + x = x.unsqueeze(2) # Add the temporal dimension back + return x + + +class QwenImage21RMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + channel_first (bool, optional): Whether the input tensor has channels as the first dimension. + Default is True. + images (bool, optional): Whether the input represents image data. Default is True. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class QwenImage21Upsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + Returns: + torch.Tensor: Upsampled tensor with the same data type as the input. + """ + + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class QwenImage21Resample(nn.Module): + r""" + A custom resampling module for 2D and 3D data. + + Args: + dim (int): The number of input/output channels. + mode (str): The resampling mode. Must be one of: + - 'none': No resampling (identity operation). + - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. + - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. + - 'downsample2d': 2D downsampling with zero-padding and convolution. + - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. + """ + + def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: + super().__init__() + self.dim = dim + self.mode = mode + + # default to dim //2 + if upsample_out_dim is None: + upsample_out_dim = dim // 2 + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + self.time_conv = QwenImage21CausalConv3d(dim, dim * 2, (1, 1), padding=(0, 0)) + + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = QwenImage21CausalConv3d(dim, dim, (1, 1), stride=(1, 1), padding=(0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + # cache last frame of last two chunk + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +class QwenImage21ResidualBlock(nn.Module): + r""" + A custom residual block module. + + Args: + in_dim (int): Number of input channels. + out_dim (int): Number of output channels. + dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. + non_linearity (str, optional): Type of non-linearity to use. Default is "silu". + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + non_linearity: str = "silu", + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = get_activation(non_linearity) + + # layers + self.norm1 = QwenImage21RMS_norm(in_dim, images=False) + self.conv1 = QwenImage21CausalConv3d(in_dim, out_dim, 3, padding=1) + self.norm2 = QwenImage21RMS_norm(out_dim, images=False) + self.dropout = nn.Dropout(dropout) + self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1) + self.conv_shortcut = QwenImage21CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + # Apply shortcut connection + h = self.conv_shortcut(x) + + # First normalization and activation + x = self.norm1(x) + x = self.nonlinearity(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + # Second normalization and activation + x = self.norm2(x) + x = self.nonlinearity(x) + + # Dropout + x = self.dropout(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.conv2(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv2(x) + + # Add residual connection + return x + h + + +class QwenImage21AttentionBlock(nn.Module): + r""" + Causal self-attention with a single head. + + Args: + dim (int): The number of channels in the input tensor. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = QwenImage21RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + def forward(self, x): + identity = x + batch_size, channels, time, height, width = x.size() + + x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) + x = self.norm(x) + + # compute query, key, value + qkv = self.to_qkv(x) + qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) + qkv = qkv.permute(0, 1, 3, 2).contiguous() + q, k, v = qkv.chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention(q, k, v) + + x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) + + # output projection + x = self.proj(x) + + # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] + x = x.view(batch_size, time, channels, height, width) + x = x.permute(0, 2, 1, 3, 4) + + return x + identity + + +class QwenImage21MidBlock(nn.Module): + """ + Middle block for QwenVAE encoder and decoder. + + Args: + dim (int): Number of input/output channels. + dropout (float): Dropout rate. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): + super().__init__() + self.dim = dim + + # Create the components + resnets = [QwenImage21ResidualBlock(dim, dim, dropout, non_linearity)] + attentions = [] + for _ in range(num_layers): + attentions.append(QwenImage21AttentionBlock(dim)) + resnets.append(QwenImage21ResidualBlock(dim, dim, dropout, non_linearity)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + # First residual block + x = self.resnets[0](x, feat_cache=feat_cache, feat_idx=feat_idx) + + # Process through attention and residual blocks + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + x = attn(x) + + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + + return x + + +class QwenImage21ResidualDownBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = QwenImage21AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + resnets = [] + for _ in range(num_res_blocks): + resnets.append(QwenImage21ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + self.resnets = nn.ModuleList(resnets) + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + self.downsampler = QwenImage21Resample(out_dim, mode=mode) + else: + self.downsampler = None + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for resnet in self.resnets: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + if self.downsampler is not None: + x = self.downsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class QwenImage21Encoder3d(nn.Module): + r""" + A 3D encoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_downsample (list of bool): Whether to downsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__( + self, + in_channels: int = 3, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + non_linearity: str = "silu", + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.nonlinearity = get_activation(non_linearity) + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv_in = QwenImage21CausalConv3d(in_channels, dims[0], 3, padding=1) + + # downsample blocks + self.down_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if is_residual: + self.down_blocks.append( + QwenImage21ResidualDownBlock( + in_dim, + out_dim, + dropout, + num_res_blocks, + temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False, + down_flag=i != len(dim_mult) - 1, + ) + ) + else: + for _ in range(num_res_blocks): + self.down_blocks.append(QwenImage21ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + self.down_blocks.append(QwenImage21AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + self.down_blocks.append(QwenImage21Resample(out_dim, mode=mode)) + scale /= 2.0 + + # middle blocks + self.mid_block = QwenImage21MidBlock(out_dim, dropout, non_linearity, num_layers=1) + + # output blocks + self.norm_out = QwenImage21RMS_norm(out_dim, images=False) + self.conv_out = QwenImage21CausalConv3d(out_dim, z_dim, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## downsamples + for layer in self.down_blocks: + if feat_cache is not None: + x = layer(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = layer(x) + + ## middle + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + + return x + + +class QwenImage21ResidualUpBlock(nn.Module): + """ + A block that handles upsampling for the QwenVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + temperal_upsample (bool): Whether to upsample on temporal dimension + up_flag (bool): Whether to upsample or not + non_linearity (str): Type of non-linearity to use + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + temperal_upsample: bool = False, + up_flag: bool = False, + non_linearity: str = "silu", + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + if up_flag: + self.avg_shortcut = QwenImage21DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2, + ) + else: + self.avg_shortcut = None + + # create residual blocks + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout, non_linearity)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + if up_flag: + upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" + self.upsampler = QwenImage21Resample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim) + else: + self.upsampler = None + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + x_copy = x.clone() + + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = resnet(x) + + if self.upsampler is not None: + if feat_cache is not None: + x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = self.upsampler(x) + + if self.avg_shortcut is not None: + x = x + self.avg_shortcut(x_copy, first_chunk=first_chunk) + + return x + + +class QwenImage21UpBlock(nn.Module): + """ + A block that handles upsampling for the QwenVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') + non_linearity (str): Type of non-linearity to use + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: str | None = None, + non_linearity: str = "silu", + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # Create layers list + resnets = [] + # Add residual blocks and attention if needed + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout, non_linearity)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + self.upsamplers = None + if upsample_mode is not None: + self.upsamplers = nn.ModuleList([QwenImage21Resample(out_dim, mode=upsample_mode)]) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=None): + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = resnet(x) + + if self.upsamplers is not None: + if feat_cache is not None: + x = self.upsamplers[0](x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = self.upsamplers[0](x) + return x + + +class QwenImage21Decoder3d(nn.Module): + r""" + A 3D decoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_upsample (list of bool): Whether to upsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + non_linearity: str = "silu", + out_channels: int = 3, + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + self.nonlinearity = get_activation(non_linearity) + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + + # init block + self.conv_in = QwenImage21CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.mid_block = QwenImage21MidBlock(dims[0], dropout, non_linearity, num_layers=1) + + # upsample blocks + self.up_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i > 0 and not is_residual: + in_dim = in_dim // 2 + + # determine if we need upsampling + up_flag = i != len(dim_mult) - 1 + # determine upsampling mode, if not upsampling, set to None + upsample_mode = None + if up_flag and temperal_upsample[i]: + upsample_mode = "upsample3d" + elif up_flag: + upsample_mode = "upsample2d" + # Create and add the upsampling block + if is_residual: + up_block = QwenImage21ResidualUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + temperal_upsample=temperal_upsample[i] if up_flag else False, + up_flag=up_flag, + non_linearity=non_linearity, + ) + else: + up_block = QwenImage21UpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + non_linearity=non_linearity, + ) + self.up_blocks.append(up_block) + + # output blocks + self.norm_out = QwenImage21RMS_norm(out_dim, images=False) + self.conv_out = QwenImage21CausalConv3d(out_dim, out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## middle + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + ## upsamples + for up_block in self.up_blocks: + x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + return x + + +def _patchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + # x shape: [batch_size, channels, frames, height, width] + batch_size, channels, frames, height, width = x.shape + + # Ensure height and width are divisible by patch_size + if height % patch_size != 0 or width % patch_size != 0: + raise ValueError(f"Height ({height}) and width ({width}) must be divisible by patch_size ({patch_size})") + + # Reshape to [batch_size, channels, frames, height//patch_size, patch_size, width//patch_size, patch_size] + x = x.view(batch_size, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size) + + # Rearrange to [batch_size, channels * patch_size * patch_size, frames, height//patch_size, width//patch_size] + x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() + x = x.view(batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size) + + return x + + +def _unpatchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + # x shape: [batch_size, (channels * patch_size * patch_size), frame, height, width] + batch_size, c_patches, frames, height, width = x.shape + channels = c_patches // (patch_size * patch_size) + + # Reshape to [b, c, patch_size, patch_size, f, h, w] + x = x.view(batch_size, channels, patch_size, patch_size, frames, height, width) + + # Rearrange to [b, c, f, h * patch_size, w * patch_size] + x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous() + x = x.view(batch_size, channels, frames, height * patch_size, width * patch_size) + + return x + + +class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Introduced in [Qwen Image 2]. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = False + _group_offload_block_modules = ["quant_conv", "post_quant_conv", "encoder", "decoder"] + # keys toignore when AlignDeviceHook moves inputs/outputs between devices + # these are shared mutable state modified in-place + _skip_keys = ["feat_cache", "feat_idx"] + + @register_to_config + def __init__( + self, + base_dim: int = 96, + decoder_base_dim: int | None = 144, + z_dim: int = 64, + dim_mult: list[int] = [1, 2, 4, 8, 8], + num_res_blocks: int = 2, + attn_scales: list[float] = [], + temperal_downsample: list[bool] = [False, True, True, True], + dropout: float = 0.0, + latents_mean: list[float] = [ + 0.5126, + 0.7721, + -0.0631, + 1.3506, + -0.7855, + -2.1025, + -0.3458, + 1.3722, + 1.8873, + -1.7177, + -0.6510, + 0.2732, + 0.7562, + -0.6163, + -1.0277, + 3.8363, + 2.0210, + 0.0472, + 0.9320, + 2.0087, + 2.4954, + -0.1391, + -1.4249, + 1.8464, + -0.5236, + 1.2826, + 3.7046, + -1.3035, + 2.7286, + -1.4518, + -1.9036, + -1.9955, + -0.0342, + -1.0265, + -0.7636, + 3.0555, + 0.0746, + -3.0751, + -0.1076, + 1.7376, + -1.0914, + -1.9435, + -0.2784, + -1.3680, + 0.4809, + -0.4433, + 0.3764, + 0.5729, + -2.0595, + 1.0960, + -1.3260, + -2.0211, + -5.0179, + 0.5275, + 4.0162, + 1.8505, + 0.3026, + 1.9373, + 1.4937, + 0.2632, + 0.5547, + -1.7121, + -0.1562, + 0.0304, + ], + latents_std: list[float] = [ + 3.2001, + 3.2936, + 3.4321, + 3.0091, + 3.1061, + 4.0379, + 4.0705, + 3.7910, + 3.0785, + 3.6500, + 3.9308, + 3.0904, + 2.8778, + 3.7675, + 3.7320, + 5.0756, + 3.2864, + 4.0397, + 3.1317, + 4.0443, + 2.9249, + 3.9454, + 3.0988, + 4.2489, + 3.4896, + 3.8513, + 3.9323, + 3.4719, + 3.7498, + 4.2830, + 3.5694, + 4.2467, + 3.9037, + 3.2947, + 5.0770, + 3.5075, + 3.2700, + 3.4767, + 2.8063, + 5.1125, + 3.5327, + 4.7833, + 3.1286, + 4.1819, + 3.8527, + 3.8312, + 3.5605, + 4.3875, + 3.9624, + 4.0168, + 3.5643, + 4.0550, + 5.5614, + 4.2963, + 4.4080, + 3.4959, + 3.8747, + 3.7608, + 3.5735, + 3.1490, + 3.7662, + 3.6746, + 3.4563, + 3.8161, + ], + is_residual: bool = True, + in_channels: int = 3, + out_channels: int = 3, + patch_size: int | None = None, + scale_factor_temporal: int | None = 8, + scale_factor_spatial: int | None = 8, + ) -> None: + super().__init__() + + self.z_dim = z_dim + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + if decoder_base_dim is None: + decoder_base_dim = base_dim + + self.encoder = QwenImage21Encoder3d( + in_channels=in_channels, + dim=base_dim, + z_dim=z_dim * 2, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_downsample=temperal_downsample, + dropout=dropout, + is_residual=is_residual, + ) + self.quant_conv = QwenImage21CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.post_quant_conv = QwenImage21CausalConv3d(z_dim, z_dim, 1) + + self.decoder = QwenImage21Decoder3d( + dim=decoder_base_dim, + z_dim=z_dim, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_upsample=self.temperal_upsample, + dropout=dropout, + out_channels=out_channels, + is_residual=is_residual, + ) + + self.spatial_compression_ratio = scale_factor_spatial + + # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension + # to perform decoding of a single video latent at a time. + self.use_slicing = False + + # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent + # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the + # intermediate tiles together, the memory requirement can be lowered. + self.use_tiling = False + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "decoder": sum(isinstance(m, QwenImage21CausalConv3d) for m in self.decoder.modules()) + if self.decoder is not None + else 0, + "encoder": sum(isinstance(m, QwenImage21CausalConv3d) for m in self.encoder.modules()) + if self.encoder is not None + else 0, + } + + def enable_tiling( + self, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + ) -> None: + 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. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + + def clear_cache(self): + # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call + self._conv_num = self._cached_conv_counts["decoder"] + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = self._cached_conv_counts["encoder"] + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + def _encode(self, x: torch.Tensor): + _, _, num_frame, height, width = x.shape + + self.clear_cache() + if self.config.patch_size is not None: + x = _patchify(x, patch_size=self.config.patch_size) + + if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + return self.tiled_encode(x) + + iter_ = 1 + (num_frame - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + + enc = self.quant_conv(out) + self.clear_cache() + return enc + + @apply_forward_hook + def encode( + self, x: torch.Tensor, return_dict: bool = True + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + r""" + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor, return_dict: bool = True): + _, _, num_frame, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.tiled_decode(z, return_dict=return_dict) + + self.clear_cache() + x = self.post_quant_conv(z) + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True + ) + else: + out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + + if self.config.patch_size is not None: + out = _unpatchify(out, patch_size=self.config.patch_size) + + out = torch.clamp(out, min=-1.0, max=1.0) + + self.clear_cache() + if not return_dict: + return (out,) + + return DecoderOutput(sample=out) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + + _, _, num_frames, height, width = x.shape + encode_spatial_compression_ratio = self.spatial_compression_ratio + if self.config.patch_size is not None: + assert encode_spatial_compression_ratio % self.config.patch_size == 0 + encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size + + latent_height = height // encode_spatial_compression_ratio + latent_width = width // encode_spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // encode_spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // encode_spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // encode_spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // encode_spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + self.clear_cache() + time = [] + frame_range = 1 + (num_frames - 1) // 4 + for k in range(frame_range): + self._enc_conv_idx = [0] + if k == 0: + tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] + else: + tile = x[ + :, + :, + 1 + 4 * (k - 1) : 1 + 4 * k, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + tile = self.quant_conv(tile) + time.append(tile) + row.append(torch.cat(time, dim=2)) + rows.append(row) + self.clear_cache() + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + + enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] + return enc + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_sample_stride_height = self.tile_sample_stride_height + tile_sample_stride_width = self.tile_sample_stride_width + if self.config.patch_size is not None: + sample_height = sample_height // self.config.patch_size + sample_width = sample_width // self.config.patch_size + tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size + tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size + blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height + blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width + else: + blend_height = self.tile_sample_min_height - tile_sample_stride_height + blend_width = self.tile_sample_min_width - tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + self.clear_cache() + time = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + tile = self.post_quant_conv(tile) + decoded = self.decoder( + tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0) + ) + time.append(decoded) + row.append(torch.cat(time, dim=2)) + rows.append(row) + self.clear_cache() + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :tile_sample_stride_height, :tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] + + if self.config.patch_size is not None: + dec = _unpatchify(dec, patch_size=self.config.patch_size) + + dec = torch.clamp(dec, min=-1.0, max=1.0) + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: + """ + Args: + sample (`torch.Tensor`): Input sample. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, return_dict=return_dict) + return dec diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 0e167812ad88..ffb0cbc0318b 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -60,6 +60,7 @@ from .transformer_ovis_image import OvisImageTransformer2DModel from .transformer_prx import PRXTransformer2DModel from .transformer_qwenimage import QwenImageTransformer2DModel + from .transformer_qwenimage21 import QwenImage21Transformer2DModel from .transformer_sana_video import SanaVideoTransformer3DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py new file mode 100644 index 000000000000..3788f193cad7 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -0,0 +1,869 @@ +# Copyright 2026 Qwen-Image Team, 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 math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils.torch_utils import maybe_allow_in_graph +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..embeddings import TimestepEmbedding +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# `create_block_mask` quantizes the mask to 128-token blocks. +_FLEX_BLOCK_SIZE = 128 + +# flex_attention is optional. When available and `causal_block=True`, we use a compiled +# flex_attention with a BlockMask for efficient block-causal attention. When unavailable, +# we fall back to a two-pass prefill: causal attention over the prefix, then full attention +# over the target image attending to the cached prefix + itself. +_FLEX_AVAILABLE = False +_compiled_flex_attention = None +try: + from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention + + _compiled_flex_attention = torch.compile(flex_attention) + _FLEX_AVAILABLE = True +except ImportError: + BlockMask = None + + +# Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen +def apply_rotary_emb_qwen( + x: torch.Tensor, + freqs_cis: torch.Tensor | tuple[torch.Tensor], + use_real: bool = True, + use_real_unbind_dim: int = -1, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings + to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are + reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting + tensors contain rotary embeddings and are returned as real tensors. + + Args: + x (`torch.Tensor`): + Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply + freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + + Returns: + tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings. + """ + if use_real: + cos, sin = freqs_cis # [S, D] + cos = cos[None, None] + sin = sin[None, None] + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, OmniGen, CogView4 and Cosmos + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2] + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(1) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) + + +class QwenImage21TemporalTimesteps(nn.Module): + r"""Sinusoidal timestep embedding. `cos` occupies the first half of the channels and `sin` the second.""" + + def __init__(self, timestep_dim: int, max_period: int = 10000, time_factor: float = 1000.0): + super().__init__() + self.timestep_dim = timestep_dim + self.time_factor = time_factor + + half = timestep_dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) + self.register_buffer("freqs", freqs, persistent=False) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + timestep = self.time_factor * timestep.float() + args = timestep[:, None] * self.freqs[None].to(timestep.device) + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if self.timestep_dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding.to(timestep.dtype) + + +class QwenImage21TimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim: int): + super().__init__() + self.time_proj = QwenImage21TemporalTimesteps(timestep_dim=256) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim, sample_proj_bias=False + ) + + def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + return self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype)) + + +class QwenImage21ZeroCenterRMSNorm(nn.Module): + r""" + RMSNorm whose learnable weight is stored zero-centered: the effective scale is `weight + 1`, computed in fp32. + Checkpoints therefore store `scale - 1`. + """ + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + rrms = torch.rsqrt(torch.mean(hidden_states**2, dim=-1, keepdim=True) + self.eps) + return (hidden_states * rrms * (self.weight.float() + 1)).to(input_dtype) + + +class QwenImage21TextProjection(nn.Module): + def __init__(self, context_in_dim: int, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps=eps) + self.in_layer = nn.Linear(context_in_dim, hidden_size, bias=False) + self.act = nn.GELU(approximate="tanh") + self.out_layer = nn.Linear(hidden_size, hidden_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.text_norm(hidden_states) + hidden_states = self.in_layer(hidden_states) + hidden_states = self.act(hidden_states) + return self.out_layer(hidden_states) + + +class QwenImage21SwiGLUFeedForward(nn.Module): + def __init__(self, hidden_size: int, mlp_hidden_size: int): + super().__init__() + self.proj = nn.Linear(hidden_size, mlp_hidden_size, bias=False) + self.out = nn.Linear(mlp_hidden_size, hidden_size, bias=False) + self.gate_layer = nn.Linear(hidden_size, mlp_hidden_size, bias=False) + self.activation_fn = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.out(self.activation_fn(self.gate_layer(hidden_states)) * self.proj(hidden_states)) + + +class QwenImage21AdaLayerNormContinuous(nn.Module): + r""" + Final adaptive norm. Scale only — this variant emits no shift, so `linear` maps to `embedding_dim` rather than + `2 * embedding_dim`. + """ + + def __init__(self, embedding_dim: int, conditioning_embedding_dim: int, eps: float = 1e-6): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=False) + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine=False, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + conditioning_embedding: torch.Tensor, + target_token_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + scale = self.linear(self.silu(conditioning_embedding).to(hidden_states.dtype)) + scale = _select_modulation_rows(scale, target_token_mask) + return self.norm(hidden_states) * (1 + scale) + + +def _select_modulation_rows(params: torch.Tensor, target_token_mask: torch.Tensor | None) -> torch.Tensor: + r""" + Broadcast per-sample modulation `params` over the token axis. + + With `causal_condition`, `params` holds `batch_size + 1` rows: rows `[0, batch_size)` come from the real timestep + and the trailing row from `t = 0`. Text and condition-image tokens take the `t = 0` row, target-image tokens take + their own sample's row. + + Args: + params (`torch.Tensor`): `(batch_size, dim)` without `causal_condition`, else `(batch_size + 1, dim)`. + target_token_mask (`torch.Tensor`, *optional*): `(seq_len,)` bool, `True` at target-image positions. `None` + disables the split and every token uses its own sample's row. + """ + if target_token_mask is None: + return params.unsqueeze(1) + real, zero = params[:-1].unsqueeze(1), params[-1:].unsqueeze(0) + return torch.where(target_token_mask.view(1, -1, 1), real, zero) + + +def build_qwenimage21_block_causal_mask( + image_ids: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None, + batch_size: int, + device: torch.device, +): + r""" + Build the block-causal [`~torch.nn.attention.flex_attention.BlockMask`] for Qwen-Image 2.1. + + The mask is `(q_idx >= kv_idx) or same_image_block`: the joint text/image sequence is causal, while every image + block — each condition image and the target image — is internally bidirectional. Text tokens are strictly causal. + Cross-sample isolation is implicit in diffusers because samples live on the batch axis. + + Positions masked out by `encoder_hidden_states_mask` are excluded as *keys* so right-padded prompts cannot be + attended to. They are kept as queries so their rows are never fully masked. + + Provided as a module-level function so callers can build the mask outside the transformer's compiled region. + + Args: + image_ids (`torch.Tensor`): `(seq_len,)` int, `-1` at text positions and a unique non-negative id per image + block. See [`~QwenImage21Transformer2DModel.build_token_metadata`]. + encoder_hidden_states_mask (`torch.Tensor`, *optional*): `(batch_size, seq_len)` bool over the joint sequence. + batch_size (`int`): Number of samples; the mask varies across the batch only through + `encoder_hidden_states_mask`. + """ + seq_len = image_ids.shape[0] + padded_seq_len = int(math.ceil(seq_len / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) + + image_ids = F.pad(image_ids, (0, padded_seq_len - seq_len), value=-1) + if encoder_hidden_states_mask is None: + key_valid = torch.ones(batch_size, padded_seq_len, dtype=torch.bool, device=device) + else: + key_valid = F.pad(encoder_hidden_states_mask.bool(), (0, padded_seq_len - seq_len), value=False) + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + is_padding = (q_idx >= seq_len) | (kv_idx >= seq_len) + q_image_id, kv_image_id = image_ids[q_idx], image_ids[kv_idx] + same_image_block = (q_image_id == kv_image_id) & (q_image_id >= 0) + allowed = ((q_idx >= kv_idx) | same_image_block) & key_valid[batch_idx, kv_idx] + return allowed & ~is_padding + + return create_block_mask( + mask_mod, + B=batch_size, + H=None, + Q_LEN=padded_seq_len, + KV_LEN=padded_seq_len, + device=device, + _compile=False, + ) + + +class QwenImage21AttnProcessor: + r""" + Single-stream attention processor for Qwen-Image 2.1. Text and image tokens share one sequence, so there is no + separate context projection. + + Two attention paths are supported: + + - **flex** (default when available): a compiled `flex_attention` with a `BlockMask` for efficient block-sparse + block-causal attention. Required for high resolutions (2048²+) where the dense score matrix would OOM. + - **SDPA fallback** (when flex is unavailable): the block-causal mask is implemented via a two-pass prefill + orchestrated by the model's `forward` — pass 1 runs the prefix with `is_causal=True`, pass 2 runs the target + image attending fully to the cached prefix + itself. The processor receives `is_causal` and a padding mask. + """ + + _attention_backend = "flex" if _FLEX_AVAILABLE else None + _parallel_config = None + + _SUPPORTED_FLEX_BACKENDS = ("flex", "_native_flex") + + def __call__( + self, + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + attention_mask: Any | None = None, + rotary_emb: torch.Tensor | None = None, + kv_cache: dict[str, torch.Tensor] | None = None, + cache_write_slice: slice | None = None, + is_causal: bool = False, + ) -> torch.Tensor: + if ( + _FLEX_AVAILABLE + and isinstance(attention_mask, BlockMask) + and self._attention_backend not in self._SUPPORTED_FLEX_BACKENDS + ): + raise ValueError( + f"QwenImage21AttnProcessor requires the 'flex' attention backend when a BlockMask is used " + f"(got {self._attention_backend!r})." + ) + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query).to(value.dtype) + key = attn.norm_k(key).to(value.dtype) + + if rotary_emb is not None: + query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) + key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) + + if kv_cache is not None: + if cache_write_slice is not None: + kv_cache["key"] = key[:, cache_write_slice].contiguous() + kv_cache["value"] = value[:, cache_write_slice].contiguous() + else: + key = torch.cat([kv_cache["key"], key], dim=1) + value = torch.cat([kv_cache["value"], value], dim=1) + + seq_len_q, seq_len_kv = query.shape[1], key.shape[1] + if _FLEX_AVAILABLE and isinstance(attention_mask, BlockMask): + pad_q = int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_q + pad_kv = int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_kv + if pad_q: + query = F.pad(query.transpose(1, 3), (0, pad_q)).transpose(1, 3) + if pad_kv: + key = F.pad(key.transpose(1, 3), (0, pad_kv)).transpose(1, 3) + value = F.pad(value.transpose(1, 3), (0, pad_kv)).transpose(1, 3) + + hidden_states = _compiled_flex_attention( + query.transpose(1, 2).contiguous(), + key.transpose(1, 2).contiguous(), + value.transpose(1, 2).contiguous(), + block_mask=attention_mask, + ).transpose(1, 2) + else: + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask if not isinstance(attention_mask, type(None)) and not is_causal else None, + dropout_p=0.0, + is_causal=is_causal, + backend=None, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states[:, :seq_len_q] + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +class QwenImage21Attention(torch.nn.Module, AttentionModuleMixin): + r""" + Attention module for [`QwenImage21TransformerBlock`]. Projection layout matches the legacy + [`~models.attention_processor.Attention`] so Qwen-Image 2.x checkpoints load into it unchanged. + """ + + _default_processor_cls = QwenImage21AttnProcessor + _available_processors = [QwenImage21AttnProcessor] + + def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-6, processor: Any | None = None): + super().__init__() + self.heads = heads + self.inner_dim = heads * dim_head + # Read by `AttentionModuleMixin.fuse_projections`; 2.1 has no biases anywhere. + self.use_bias = False + + self.to_q = nn.Linear(dim, self.inner_dim, bias=False) + self.to_k = nn.Linear(dim, self.inner_dim, bias=False) + self.to_v = nn.Linear(dim, self.inner_dim, bias=False) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)]) + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.set_processor(processor if processor is not None else self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.processor(self, hidden_states, **kwargs) + + +@maybe_allow_in_graph +class QwenImage21TransformerBlock(nn.Module): + r""" + Single-stream block. Modulation is not learned per block — the parent model computes one shared `modulation` tensor + and every block slices its own scales and gates out of it. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: int = 3, + eps: float = 1e-6, + ): + super().__init__() + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = QwenImage21Attention(dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, eps=eps) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = QwenImage21SwiGLUFeedForward(hidden_size=dim, mlp_hidden_size=dim * mlp_ratio) + + def _modulate( + self, + hidden_states: torch.Tensor, + mod_params: torch.Tensor, + target_token_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + scale, gate = mod_params.chunk(2, dim=-1) + scale = _select_modulation_rows(scale, target_token_mask) + gate = _select_modulation_rows(gate, target_token_mask) + return hidden_states * (1 + scale), gate + + def forward( + self, + hidden_states: torch.Tensor, + modulation: torch.Tensor, + rotary_emb: torch.Tensor | None = None, + attention_mask: Any | None = None, + target_token_mask: torch.Tensor | None = None, + kv_cache: dict[str, torch.Tensor] | None = None, + cache_write_slice: slice | None = None, + is_causal: bool = False, + ) -> torch.Tensor: + mod1, mod2 = modulation.chunk(2, dim=-1) + + img_modulated, img_gate1 = self._modulate(self.img_norm1(hidden_states), mod1, target_token_mask) + attn_output = self.attn( + hidden_states=img_modulated, + attention_mask=attention_mask, + rotary_emb=rotary_emb, + kv_cache=kv_cache, + cache_write_slice=cache_write_slice, + is_causal=is_causal, + ) + hidden_states = hidden_states + img_gate1.tanh() * attn_output + + img_modulated2, img_gate2 = self._modulate(self.img_norm2(hidden_states), mod2, target_token_mask) + hidden_states = hidden_states + img_gate2.tanh() * self.img_mlp(img_modulated2) + + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + return hidden_states + + +class QwenImage21Rope(nn.Module): + r""" + 3-axis (frame, height, width) rotary embedding over the joint text/image sequence. + + Text tokens advance a shared position on all three axes. Every image block freezes the frame axis at the position + reached by the preceding text and lays its tokens out on a height/width grid centred on zero, so a block's spatial + positions do not depend on where it sits in the sequence. + """ + + def __init__(self, theta: int, axes_dim: list[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + pos_index = torch.arange(8192) + neg_index = torch.arange(1024).flip(0) * -1 - 1 + self.freqs = [ + torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0) + for dim in axes_dim + ] + + def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor: + freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) + return torch.polar(torch.ones_like(freqs), freqs) + + def forward( + self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device + ) -> torch.Tensor: + self.freqs = [freq.to(device) for freq in self.freqs] + + frame_index, height_index, width_index = [], [], [] + image_height_index, image_width_index = [], [] + cursor, position = 0, 0 + total_len = image_pad_mask.shape[-1] + is_image_token = image_pad_mask.tolist() + + for _, height, width in img_shapes: + block_start = is_image_token.index(True, cursor) + text_len = block_start - cursor + frame_index.extend(range(position, position + text_len)) + position += text_len + + cursor = block_start + height * width + frame_index.extend([position] * (height * width)) + position += max(height, width) + + image_height_index.extend([h for h in range(-(height - height // 2), height // 2) for _ in range(width)]) + image_width_index.extend([w for _ in range(height) for w in range(-(width - width // 2), width // 2)]) + + if cursor < total_len: + frame_index.extend(range(position, position + total_len - cursor)) + + frame_index = torch.tensor(frame_index, dtype=torch.long, device=device) + height_index = frame_index.clone() + width_index = frame_index.clone() + height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device=device) + width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device=device) + + return torch.cat([self.freqs[0][frame_index], self.freqs[1][height_index], self.freqs[2][width_index]], dim=-1) + + +class QwenImage21Transformer2DModel( + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin +): + r""" + The single-stream Transformer used by Qwen-Image 2.1. + + Text and image latents share one sequence: condition-image tokens are substituted into the text stream at the + positions the vision-language encoder reserved for them, and the target image's tokens are appended. A single + shared `modulation` projection feeds every block, so blocks hold no modulation parameters of their own. + + Two behaviours distinguish 2.1 from 2.0, both switched on by config and neither adding parameters: + + - `causal_block` — attention follows `(q_idx >= kv_idx) or same_image_block`, so the sequence is causal while each + image block stays internally bidirectional. This requires the `flex` attention backend. + - `causal_condition` — text and condition-image tokens are modulated from `t = 0` instead of the sampled timestep, + which also makes their activations timestep-independent and so cacheable across denoising steps. + + Args: + patch_size (`int`, defaults to `1`): + Side length of the latent patch folded into the channel dim. 2.1 consumes latents unpatched. + in_channels (`int`, defaults to `64`): + Latent channels of the input. + out_channels (`int`, *optional*, defaults to `64`): + Latent channels of the output. Falls back to `in_channels`. + num_layers (`int`, defaults to `32`): + Number of single-stream blocks. + attention_head_dim (`int`, defaults to `128`): + Channels per attention head. + num_attention_heads (`int`, defaults to `32`): + Number of attention heads. + context_in_dim (`int`, defaults to `4096`): + Channel dim of `encoder_hidden_states`. + mlp_ratio (`int`, defaults to `3`): + Feed-forward expansion factor. + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): + Rotary dims for the frame, height and width axes. + eps (`float`, defaults to `1e-6`): + Epsilon for the norm layers. + causal_condition (`bool`, defaults to `True`): + Modulate text and condition-image tokens from `t = 0`. Required for KV caching. + causal_block (`bool`, defaults to `True`): + Use block-causal attention. Requires the `flex` attention backend. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["QwenImage21TransformerBlock"] + _skip_layerwise_casting_patterns = ["pos_embed", "norm"] + _repeated_blocks = ["QwenImage21TransformerBlock"] + _skip_keys = ["kv_cache"] + + @register_to_config + def __init__( + self, + patch_size: int = 1, + in_channels: int = 64, + out_channels: int | None = 64, + num_layers: int = 32, + attention_head_dim: int = 128, + num_attention_heads: int = 32, + context_in_dim: int = 4096, + mlp_ratio: int = 3, + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), + eps: float = 1e-6, + causal_condition: bool = True, + causal_block: bool = True, + ): + super().__init__() + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope)) + self.time_text_embed = QwenImage21TimestepProjEmbeddings(embedding_dim=self.inner_dim) + self.txt_in = QwenImage21TextProjection(context_in_dim, self.inner_dim, eps=eps) + self.img_in = nn.Linear(in_channels * patch_size * patch_size, self.inner_dim, bias=False) + + # One shared modulation for every block: [mod1.scale, mod1.gate, mod2.scale, mod2.gate]. + self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 4 * self.inner_dim, bias=False)) + + self.transformer_blocks = nn.ModuleList( + [ + QwenImage21TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = QwenImage21AdaLayerNormContinuous(self.inner_dim, self.inner_dim, eps=eps) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False) + + self.gradient_checkpointing = False + + @staticmethod + def build_token_metadata( + image_pad_mask: torch.Tensor, img_shapes: list[tuple[int, int, int]] + ) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Label every token of the joint sequence with the image block it belongs to. + + Block boundaries come from the token counts in `img_shapes`, not from runs of `True` in `image_pad_mask`: two + condition images that happen to sit next to each other with no text between them form one run but must stay + separate blocks, otherwise they would attend to each other bidirectionally. + + Args: + image_pad_mask (`torch.Tensor`): `(seq_len,)` bool, `True` at image-token positions. + img_shapes (`list[tuple[int, int, int]]`): Per-image `(frame, height, width)` in latent tokens, condition + images first and the target image last. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: `image_ids` `(seq_len,)` with `-1` at text positions and a unique id + per image block, and `target_token_mask` `(seq_len,)` marking the target image's tokens. + """ + image_positions = image_pad_mask.nonzero(as_tuple=True)[0] + block_lengths = [math.prod(shape) for shape in img_shapes] + if sum(block_lengths) != image_positions.numel(): + raise ValueError( + f"img_shapes accounts for {sum(block_lengths)} image tokens but image_pad_mask marks " + f"{image_positions.numel()}." + ) + + image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) + block_ids = torch.repeat_interleave( + torch.arange(len(block_lengths), device=image_pad_mask.device), + torch.tensor(block_lengths, device=image_pad_mask.device), + ) + image_ids[image_positions] = block_ids + + target_token_mask = torch.zeros_like(image_pad_mask) + target_token_mask[image_positions[-block_lengths[-1] :]] = True + return image_ids, target_token_mask + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + img_shapes: list[list[tuple[int, int, int]]], + img_mask: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + kv_cache: list[dict[str, torch.Tensor]] | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): + Packed latents, condition images first and the target image last. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, context_in_dim)`): + Text embeddings from the vision-language encoder. + timestep (`torch.Tensor`): + Current denoising step, scaled to `[0, 1]`. + img_shapes (`list[list[tuple[int, int, int]]]`): + Per-sample list of `(frame, height, width)` in latent tokens, condition images first and the target + image last. All samples must share a layout. + img_mask (`torch.Tensor` of shape `(batch_size, vlm_sequence_length)`): + `True` at the vision-language encoder's image slots, each standing for a `2x2` group of latent tokens. + encoder_hidden_states_mask (`torch.Tensor`, *optional*): + `(batch_size, text_sequence_length)` bool marking valid text tokens. Padded positions are excluded + from attention. + kv_cache (`list[dict[str, torch.Tensor]]`, *optional*): + One dict per block. Empty dicts prefill the text and condition-image keys and values; populated dicts + switch to decode, where only the target image's tokens are recomputed. Requires `causal_condition`. + """ + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + scale_lora_layers(self, lora_scale) + elif attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: + logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") + + batch_size = hidden_states.shape[0] + hidden_states = self.img_in(hidden_states) + encoder_hidden_states = self.txt_in(encoder_hidden_states) + + # Each vision-language image slot stands for 2x2 latent tokens, so expand those positions four-fold and drop + # the actual latents into them. Samples share a layout, hence the single row. + repeats = torch.where(img_mask, 4, 1)[0] + image_pad_mask = torch.repeat_interleave(img_mask[0], repeats) + + target_tokens = math.prod(img_shapes[0][-1]) + joint_hidden_states = torch.cat( + [ + encoder_hidden_states, + encoder_hidden_states.new_zeros(batch_size, target_tokens // 4, encoder_hidden_states.shape[2]), + ], + dim=1, + ) + joint_hidden_states = joint_hidden_states.repeat_interleave(repeats, dim=1) + joint_hidden_states[:, image_pad_mask] = hidden_states + + rotary_emb = self.pos_embed(img_shapes[0], image_pad_mask, device=hidden_states.device) + image_ids, target_token_mask = self.build_token_metadata(image_pad_mask, img_shapes[0]) + + timestep = timestep.to(hidden_states.dtype) + if self.config.causal_condition: + # Extra t=0 row; text and condition-image tokens modulate from it. `modulation_mask` selects which row + # each token reads, and is `None` when every token shares the sampled timestep. + timestep = torch.cat([timestep, timestep.new_zeros(1)], dim=0) + modulation_mask = target_token_mask + else: + modulation_mask = None + temb = self.time_text_embed(timestep, hidden_states) + modulation = self.modulation(temb) + + if kv_cache is not None and not self.config.causal_condition: + raise ValueError( + "kv_cache requires `causal_condition=True`. The cache is only valid because text and condition-image " + "tokens modulate from t=0, which makes their activations independent of the denoising step." + ) + + # Right-padded prompt positions must never be attended to, on any path. Text positions of the joint sequence + # line up, in order, with the non-image positions of the vision-language sequence — the two are interleaved, + # so the mask cannot be sliced off as a prefix. + joint_key_valid = None + if encoder_hidden_states_mask is not None: + joint_key_valid = torch.ones( + batch_size, image_pad_mask.shape[0], dtype=torch.bool, device=hidden_states.device + ) + text_positions = (~image_pad_mask).nonzero(as_tuple=True)[0] + vlm_text_positions = ~img_mask[0][: encoder_hidden_states_mask.shape[1]] + joint_key_valid[:, text_positions] = encoder_hidden_states_mask.bool()[:, vlm_text_positions] + + prefix_len = int((~target_token_mask).sum()) + is_decode = kv_cache is not None and len(kv_cache[0]) > 0 + use_flex = _FLEX_AVAILABLE and self.config.causal_block + + if is_decode: + # Only the target image's queries are recomputed. The block-causal mask degenerates to full attention + # for target rows (they can see the entire prefix + their own block), so no structural mask is needed. + joint_hidden_states = joint_hidden_states[:, prefix_len:] + rotary_emb = rotary_emb[prefix_len:] + modulation_mask = modulation_mask[prefix_len:] + attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] + cache_write_slice = None + use_two_pass = False + elif use_flex: + # flex path: single-pass with a compiled BlockMask + cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None + attention_mask = build_qwenimage21_block_causal_mask( + image_ids, joint_key_valid, batch_size, hidden_states.device + ) + use_two_pass = False + elif self.config.causal_block: + # No flex_attention available: approximate the block-causal mask with a two-pass prefill. + # Pass 1 runs the prefix (text + condition images) causally and caches its keys and values; + # pass 2 runs the target image attending fully over that prefix plus itself. + # + # This is exact for the target image and for text, but not for condition images: under the real + # block-causal mask a condition image attends within its own block bidirectionally, whereas a + # single causal pass only lets each of its tokens see earlier ones. Text-only prompts are + # therefore unaffected; prompts with condition images differ slightly. Install a PyTorch build + # with flex_attention for the exact mask. + use_two_pass = True + cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None + else: + # causal_block disabled: full attention + cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None + attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] + use_two_pass = False + + if use_two_pass: + # Two-pass prefill: split into prefix and target, run prefix causally then target fully. + prefix_hs = joint_hidden_states[:, :prefix_len] + target_hs = joint_hidden_states[:, prefix_len:] + prefix_rope = rotary_emb[:prefix_len] + target_rope = rotary_emb[prefix_len:] + prefix_mod_mask = modulation_mask[:prefix_len] if modulation_mask is not None else None + target_mod_mask = modulation_mask[prefix_len:] if modulation_mask is not None else None + for index_block, block in enumerate(self.transformer_blocks): + block_kv_cache = kv_cache[index_block] if kv_cache is not None else None + # Pass 1: prefix with causal attention. No padding mask — padded text tokens have zero embeddings + # from right-padding and is_causal handles the structural mask. + prefix_cache = {} if kv_cache is not None else None + prefix_hs = block( + hidden_states=prefix_hs, + modulation=modulation, + rotary_emb=prefix_rope, + attention_mask=None, + target_token_mask=prefix_mod_mask, + kv_cache=prefix_cache, + cache_write_slice=slice(0, prefix_len), + is_causal=True, + ) + # Pass 2: target image with full attention over [cached prefix, target]. + # No attention mask needed: the prefix cache already excludes padded positions, and the target + # image should see everything (block-causal degenerates to full attention for target rows). + if prefix_cache is not None and block_kv_cache is not None: + block_kv_cache.update(prefix_cache) + target_hs = block( + hidden_states=target_hs, + modulation=modulation, + rotary_emb=target_rope, + attention_mask=None, + target_token_mask=target_mod_mask, + kv_cache=block_kv_cache, + cache_write_slice=None, + ) + + joint_hidden_states = torch.cat([prefix_hs, target_hs], dim=1) + else: + for index_block, block in enumerate(self.transformer_blocks): + block_kv_cache = kv_cache[index_block] if kv_cache is not None else None + if torch.is_grad_enabled() and self.gradient_checkpointing: + joint_hidden_states = self._gradient_checkpointing_func( + block, + joint_hidden_states, + modulation, + rotary_emb, + attention_mask, + modulation_mask, + block_kv_cache, + cache_write_slice, + ) + else: + joint_hidden_states = block( + hidden_states=joint_hidden_states, + modulation=modulation, + rotary_emb=rotary_emb, + attention_mask=attention_mask, + target_token_mask=modulation_mask, + kv_cache=block_kv_cache, + cache_write_slice=cache_write_slice, + ) + + joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) + output = self.proj_out(joint_hidden_states) + + if USE_PEFT_BACKEND: + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index fed3e449a7e1..32f193a03080 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -467,6 +467,7 @@ "SkyReelsV2Pipeline", ] _import_structure["nucleusmoe_image"] = ["NucleusMoEImagePipeline"] + _import_structure["qwenimage21"] = ["QwenImage21Pipeline"] _import_structure["qwenimage"] = [ "QwenImagePipeline", "QwenImageImg2ImgPipeline", @@ -864,6 +865,7 @@ QwenImageLayeredPipeline, QwenImagePipeline, ) + from .qwenimage21 import QwenImage21Pipeline from .sana import ( SanaControlNetPipeline, SanaPipeline, diff --git a/src/diffusers/pipelines/qwenimage21/__init__.py b/src/diffusers/pipelines/qwenimage21/__init__.py new file mode 100644 index 000000000000..9faa8bb9f0c0 --- /dev/null +++ b/src/diffusers/pipelines/qwenimage21/__init__.py @@ -0,0 +1,48 @@ +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 = {} +_additional_imports = {} +_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_qwenimage21"] = ["QwenImage21Pipeline"] + +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 * # noqa F403 + else: + from .pipeline_qwenimage21 import QwenImage21Pipeline +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) + for name, value in _additional_imports.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py new file mode 100644 index 000000000000..ad98fa84e59b --- /dev/null +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -0,0 +1,804 @@ +# Copyright 2026 Qwen-Image Team, 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 +import random +from typing import Any, Callable + +import numpy as np +import torch +from PIL import Image as PILImage +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import QwenImageLoraLoaderMixin +from ...models import AutoencoderKLQwenImage21, QwenImage21Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from ..qwenimage.pipeline_output import QwenImagePipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import QwenImage21Pipeline + + >>> pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + >>> prompt = "A capybara wearing a wizard hat, reading a book by candlelight, oil painting" + >>> image = pipe(prompt, num_inference_steps=50).images[0] + >>> image.save("qwenimage21.png") + ``` +""" + + +# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.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.qwenimage.pipeline_qwenimage.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = 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 + + +def retrieve_latents(encoder_output, generator=None, sample_mode="sample"): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +def calculate_dimensions(target_area, ratio): + width = math.sqrt(target_area * ratio) + height = width / ratio + return round(width / 32) * 32, round(height / 32) * 32 + + +class QwenImage21Pipeline(DiffusionPipeline, QwenImageLoraLoaderMixin): + r""" + Text-to-image and image-conditioned generation with Qwen-Image 2.1. + + Prompt and condition images are encoded together by a Qwen3-VL model, so a condition image occupies the vision + slots the encoder reserved for it and the transformer sees one interleaved text/image sequence. + + Args: + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Scheduler used to denoise the encoded image latents. + vae ([`AutoencoderKLQwenImage21`]): + Variational auto-encoder mapping images to and from the 64-channel latent space. + text_encoder ([`Qwen3VLForConditionalGeneration`]): + Qwen3-VL model producing the joint text/image embeddings. + processor ([`Qwen3VLProcessor`]): + Processor that builds the chat template and tokenizes prompt and condition images. + transformer ([`QwenImage21Transformer2DModel`]): + The single-stream block-causal transformer that denoises the latents. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLQwenImage21, + text_encoder: Qwen3VLForConditionalGeneration, + processor: Qwen3VLProcessor, + transformer: QwenImage21Transformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + processor=processor, + transformer=transformer, + scheduler=scheduler, + ) + # The VAE compresses 16x spatially and the transformer consumes latents unpatched, so one token covers a 16x16 + # pixel tile. + self.vae_scale_factor = 16 + self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 64 + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor, vae_latent_channels=self.latent_channels + ) + self.sys_prompt = "Comprehend and analyze the provided prompt." + # The prompt is built as a raw template string and passed straight to + # `self.processor(text=..., images=...)`, rather than going through `apply_chat_template`: + # the two tokenize differently and the checkpoint expects this one. The "Picture 1: ..." + # vision prefix only appears in the image-conditioned template. + self.prompt_template_t2i = ( + f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" + f"<|im_start|>user\n{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + self.prompt_template_ti2i = ( + f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" + f"<|im_start|>user\nPicture 1: <|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + self.ref_token_list = ["Picture ", "Image ", "图 ", "图片 "] + # Number of leading system-role tokens to drop from the hidden states. Derived from the + # tokenized system message rather than hardcoded, so it tracks the processor's template. + sys_message = [{"role": "system", "content": [{"type": "text", "text": self.sys_prompt}]}] + sys_tokens = self.processor.apply_chat_template(sys_message, tokenize=True, return_dict=False) + self._drop_idx_t2i = len(sys_tokens[0]) + self._drop_idx_ti2i = self._drop_idx_t2i + self._img_token_id = self.processor.tokenizer.encode("<|image_pad|>")[0] + self._max_length = 8192 + + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + return torch.split(selected, valid_lengths.tolist(), dim=0) + + @staticmethod + def _downsample_image_pad_tokens(hidden_states_list, image_pad_mask_list): + """Collapse consecutive `<|image_pad|>` tokens into one per contiguous region. + + The vision-language processor expands each condition image into many vision tokens, but the + transformer expects one token per image slot, which it then expands 4x. Keep the first token of + each contiguous image-pad region and drop the rest. + """ + out_hs, out_mask = [], [] + for hidden_state, pad_mask in zip(hidden_states_list, image_pad_mask_list): + non_pad = ~pad_mask + non_pad_tokens = hidden_state[non_pad] + non_pad_mask = pad_mask[non_pad] + + pad_indices = torch.where(pad_mask)[0] + insert_tokens, insert_positions = [], [] + if len(pad_indices) > 0: + region_starts = [pad_indices[0].item()] + if len(pad_indices) > 1: + diff = torch.diff(pad_indices) + for j, d in enumerate(diff): + if d > 1: + region_starts.append(pad_indices[j + 1].item()) + + for start_idx in region_starts: + insert_pos = non_pad[:start_idx].sum().item() + insert_tokens.append(hidden_state[start_idx]) + insert_positions.append(insert_pos) + + result_tokens = non_pad_tokens + result_mask = non_pad_mask + if insert_positions: + for idx in sorted(range(len(insert_positions)), key=lambda i: insert_positions[i], reverse=True): + pos = insert_positions[idx] + result_tokens = torch.cat( + [result_tokens[:pos], insert_tokens[idx].unsqueeze(0), result_tokens[pos:]] + ) + result_mask = torch.cat( + [ + result_mask[:pos], + torch.tensor([True], dtype=torch.bool, device=result_mask.device), + result_mask[pos:], + ] + ) + + out_hs.append(result_tokens) + out_mask.append(result_mask) + return out_hs, out_mask + + def _get_qwen_prompt_embeds( + self, + prompt: str | list[str] = None, + image: list | None = None, + device: torch.device | None = None, + ): + device = device or self._execution_device + prompt = [prompt] if isinstance(prompt, str) else prompt + is_t2i = image is None + + if is_t2i: + prompts = [self.prompt_template_t2i.format(t) for t in prompt] + drop_idx = self._drop_idx_t2i + else: + prompts = [] + condition_pil_list = [] + for t in prompt: + n_imgs = len(image) + replace = "Picture 1: <|vision_start|><|image_pad|><|vision_end|>" + for i in range(2, n_imgs + 1): + replace += f" Picture {i}: <|vision_start|><|image_pad|><|vision_end|>" + template = self.prompt_template_ti2i.replace( + "Picture 1: <|vision_start|><|image_pad|><|vision_end|>", + replace.replace("Picture ", random.choice(self.ref_token_list)), + ) + prompts.append(template.format(t)) + for img in image: + if not isinstance(img, PILImage.Image): + img = PILImage.fromarray(img) + condition_pil_list.append(img) + drop_idx = self._drop_idx_ti2i + + processor_kwargs = {"text": prompts, "padding": True, "return_tensors": "pt"} + if not is_t2i: + processor_kwargs["images"] = condition_pil_list + + model_inputs = self.processor(**processor_kwargs).to(device) + + forward_kwargs = { + "input_ids": model_inputs.input_ids, + "attention_mask": model_inputs.attention_mask, + "output_hidden_states": True, + } + if not is_t2i and hasattr(model_inputs, "pixel_values"): + forward_kwargs.update(pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw) + + outputs = self.text_encoder(**forward_kwargs) + hidden_states = outputs.hidden_states[-1] + + split_hidden_states = list(self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + + image_pad_mask = [ + (sample_ids[sample_mask.bool()] == self._img_token_id) + for sample_ids, sample_mask in zip(model_inputs.input_ids, model_inputs.attention_mask) + ] + image_pad_mask = [e[drop_idx:] for e in image_pad_mask] + + if not is_t2i: + split_hidden_states, image_pad_mask = self._downsample_image_pad_tokens( + split_hidden_states, image_pad_mask + ) + + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max(e.size(0) for e in split_hidden_states) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + image_pad_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in image_pad_mask]) + + return prompt_embeds, encoder_attention_mask, image_pad_mask + + def encode_prompt( + self, + prompt: str | list[str], + image: list[PipelineImageInput] | None = None, + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + image_pad_mask: torch.Tensor | None = None, + ): + r""" + Args: + prompt (`str` or `list[str]`, *optional*): + Prompt to be encoded. + image (`list[PipelineImageInput]`, *optional*): + Condition images to encode alongside the prompt. + device (`torch.device`): + Torch device. + num_images_per_prompt (`int`): + Number of images generated per prompt. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Skips encoding when provided. + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask, image_pad_mask = self._get_qwen_prompt_embeds(prompt, image, device) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + return prompt_embeds, prompt_embeds_mask, image_pad_mask + + def check_inputs(self, prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and " + f"{width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found " + f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError("Pass either `prompt` or `prompt_embeds`, not both.") + if prompt is None and prompt_embeds is None: + raise ValueError("Pass one of `prompt` or `prompt_embeds`.") + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + # 2.1 consumes latents unpatched, so packing is a plain spatial flatten. + return latents.view(batch_size, num_channels_latents, height * width).transpose(1, 2) + + @staticmethod + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, _, channels = latents.shape + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + latents = latents.transpose(1, 2).reshape(batch_size, channels, 1, height, width) + return latents + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax") + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + return (image_latents - latents_mean) / latents_std + + def prepare_latents( + self, images, 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)) + + image_latents = None + if images is not None: + all_image_latents = [] + for image in images: + image = image.to(device=device, dtype=dtype) + encoded = image if image.shape[1] == self.latent_channels else self._encode_vae_image(image, generator) + if batch_size > encoded.shape[0]: + if batch_size % encoded.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {encoded.shape[0]} to {batch_size} text prompts." + ) + encoded = torch.cat([encoded] * (batch_size // encoded.shape[0]), dim=0) + image_latent_height, image_latent_width = encoded.shape[3:] + all_image_latents.append( + self._pack_latents( + encoded, batch_size, num_channels_latents, image_latent_height, image_latent_width + ) + ) + image_latents = torch.cat(all_image_latents, dim=1) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + shape = (batch_size, 1, num_channels_latents, height, width) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + else: + latents = latents.to(device=device, dtype=dtype) + + return latents, image_latents + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + image: PipelineImageInput | None = None, + negative_prompt: str | list[str] = None, + true_cfg_scale: float = 4.0, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + output_resolution: int = 1024, + use_kv_cache: bool = True, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt to guide image generation. Pass `prompt_embeds` instead to supply embeddings directly. + image (`PipelineImageInput`, *optional*): + One or more condition images. They are encoded by the text encoder as vision context and by the VAE + into latent tokens prepended to the noise. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt not to guide image generation. Ignored when `true_cfg_scale` is not greater than 1. + true_cfg_scale (`float`, *optional*, defaults to 4.0): + Classifier-free guidance scale. Enabled by `true_cfg_scale > 1` together with a negative prompt. + height (`int`, *optional*): + Height in pixels of the generated image. Derived from the condition image's aspect ratio if omitted. + width (`int`, *optional*): + Width in pixels of the generated image. Derived from the condition image's aspect ratio if omitted. + num_inference_steps (`int`, *optional*, defaults to 50): + Number of denoising steps. + sigmas (`list[float]`, *optional*): + Custom sigmas for the denoising schedule. + num_images_per_prompt (`int`, *optional*, defaults to 1): + Number of images generated per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Generator(s) to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents. + output_type (`str`, *optional*, defaults to `"pil"`): + Output format, `"pil"`, `"np"`, `"pt"` or `"latent"`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + Passed through to the attention processor. + callback_on_step_end (`Callable`, *optional*): + Called at the end of each denoising step. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive `height`/`width` and to resize condition images. + use_kv_cache (`bool`, *optional*, defaults to `True`): + Cache the text and condition-image keys and values after the first step. Valid because + `causal_condition` modulates those tokens from `t = 0`, making their activations step-independent. + + Toggling this does not reproduce the same image bit-for-bit in reduced precision. Caching makes the + decode step attend with a different sequence layout than the prefill step, so the two tile + differently and land on different rounding; both agree with an fp32 reference to the same tolerance. + A one-ULP difference at the first block is then amplified by 32 blocks and every sampler step, so the + two settings give equally valid but visibly distinct samples. Fix a sample by fixing this flag. + + Examples: + + Returns: + [`~pipelines.qwenimage.QwenImagePipelineOutput`] or `tuple`: + [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple` whose first + element is a list with the generated images. + """ + if image is not None: + image_size = image[-1].size if isinstance(image, list) else image.size + calculated_width, calculated_height = calculate_dimensions( + output_resolution * output_resolution, image_size[0] / image_size[1] + ) + height = height or calculated_height + width = width or calculated_width + height = height or output_resolution + width = width or output_resolution + + multiple_of = self.vae_scale_factor * 2 + width = width // multiple_of * multiple_of + height = height // multiple_of * multiple_of + + self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs) + + self._attention_kwargs = attention_kwargs or {} + self._current_timestep = None + self._interrupt = False + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 1. Preprocess condition images: one resize feeds both the text encoder and the VAE. + input_image_sizes, input_images, vae_images = [], None, None + if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): + image = image if isinstance(image, list) else [image] + input_images, vae_images = [], [] + for img in image: + image_width, image_height = img.size + input_width, input_height = calculate_dimensions( + output_resolution * output_resolution, image_width / image_height + ) + input_image_sizes.append((input_width, input_height)) + input_images.append(self.image_processor.resize(img, width=input_width, height=input_height)) + vae_images.append( + self.image_processor.preprocess(img, width=input_width, height=input_height).unsqueeze(2) + ) + + # 2. Encode prompt + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None + ) + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + if true_cfg_scale > 1 and not has_neg_prompt: + logger.warning( + f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no " + f"negative_prompt is provided." + ) + elif true_cfg_scale <= 1 and has_neg_prompt: + logger.warning( + "negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1" + ) + + prompt_embeds, prompt_embeds_mask, image_pad_mask = self.encode_prompt( + image=input_images, + prompt=prompt, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + device=device, + num_images_per_prompt=num_images_per_prompt, + ) + if do_true_cfg: + negative_prompt_embeds, negative_prompt_embeds_mask, negative_image_pad_mask = self.encode_prompt( + image=input_images, + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + device=device, + num_images_per_prompt=num_images_per_prompt, + ) + + # 3. Prepare latents + num_channels_latents = self.transformer.config.in_channels + latents, input_images_latents = self.prepare_latents( + vae_images, + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + img_shapes = [ + [ + *[ + (1, vae_height // self.vae_scale_factor, vae_width // self.vae_scale_factor) + for vae_width, vae_height in input_image_sizes + ], + (1, height // self.vae_scale_factor, width // self.vae_scale_factor), + ] + ] * batch_size + + # 4. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + latents.shape[1], + 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), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # The transformer's `img_mask` spans the joint sequence, so append one slot per 2x2 group of target latents. + target_slots = torch.ones( + [latents.shape[0], latents.shape[1] // 4], dtype=image_pad_mask.dtype, device=image_pad_mask.device + ) + image_pad_mask = torch.cat([image_pad_mask, target_slots], dim=1) + if do_true_cfg: + negative_image_pad_mask = torch.cat([negative_image_pad_mask, target_slots], dim=1) + + # Text and condition-image keys and values are step-independent under `causal_condition`, so the first step + # prefills them and later steps only recompute the target image's tokens. + num_blocks = len(self.transformer.transformer_blocks) + cache_enabled = use_kv_cache and self.transformer.config.causal_condition + cond_cache = [{} for _ in range(num_blocks)] if cache_enabled else None + neg_cache = [{} for _ in range(num_blocks)] if cache_enabled and do_true_cfg else None + + # 5. Denoising loop + self.scheduler.set_begin_index(0) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + latent_model_input = latents + if input_images_latents is not None: + latent_model_input = torch.cat([input_images_latents, latents], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + img_shapes=img_shapes, + img_mask=image_pad_mask, + attention_kwargs=self.attention_kwargs, + kv_cache=cond_cache, + return_dict=False, + )[0] + noise_pred = noise_pred[:, -latents.size(1) :] + + if do_true_cfg: + with self.transformer.cache_context("uncond"): + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states_mask=negative_prompt_embeds_mask, + img_shapes=img_shapes, + img_mask=negative_image_pad_mask, + attention_kwargs=self.attention_kwargs, + kv_cache=neg_cache, + return_dict=False, + )[0] + neg_noise_pred = neg_noise_pred[:, -latents.size(1) :] + noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + if latents.dtype != latents_dtype and torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: + # https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs} + 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) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] + image = self.image_processor.postprocess(image, output_type=output_type) + + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return QwenImagePipelineOutput(images=image) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 1598814f835a..56a44f0fe360 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -840,6 +840,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderKLQwenImage21(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderKLTemporalDecoder(metaclass=DummyObject): _backends = ["torch"] @@ -2025,6 +2040,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class QwenImage21Transformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class QwenImageControlNetModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 17bca9f23414..ed724e7de751 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3707,6 +3707,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class QwenImage21Pipeline(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"]) + + class QwenImageControlNetInpaintPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py new file mode 100644 index 000000000000..ae3a43fff77a --- /dev/null +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -0,0 +1,235 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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 pytest +import torch +from torch.nn.attention.flex_attention import create_mask + +from diffusers import QwenImage21Transformer2DModel +from diffusers.models.transformers.transformer_qwenimage21 import build_qwenimage21_block_causal_mask +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class QwenImage21TransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return QwenImage21Transformer2DModel + + @property + def output_shape(self) -> tuple[int, int]: + return (8, 4) + + @property + def input_shape(self) -> tuple[int, int]: + return (4, 4) + + @property + def model_split_percents(self) -> list: + return [0.7, 0.6, 0.6] + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + # `attention_head_dim` is 16 because flex_attention requires a head dim of at least 16, and + # `axes_dims_rope` must sum to it. + return { + "patch_size": 1, + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "attention_head_dim": 16, + "num_attention_heads": 2, + "context_in_dim": 8, + "mlp_ratio": 2, + "axes_dims_rope": (4, 6, 6), + } + + def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]: + text_len, target_height, target_width = 4, 2, 2 + target_tokens = target_height * target_width + + hidden_states = randn_tensor((batch_size, target_tokens, 4), generator=self.generator, device=device) + encoder_hidden_states = randn_tensor((batch_size, text_len, 8), generator=self.generator, device=device) + encoder_hidden_states_mask = torch.ones((batch_size, text_len), device=device, dtype=torch.long) + # One vision slot per 2x2 group of target latents, appended after the text tokens. + img_mask = torch.zeros((batch_size, text_len + target_tokens // 4), device=device, dtype=torch.bool) + img_mask[:, text_len:] = True + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "encoder_hidden_states_mask": encoder_hidden_states_mask, + "timestep": torch.tensor([1.0], device=device).expand(batch_size), + "img_shapes": [[(1, target_height, target_width)]] * batch_size, + "img_mask": img_mask, + } + + +class TestQwenImage21Transformer(QwenImage21TransformerTesterConfig, ModelTesterMixin): + @pytest.mark.skip( + reason="The block-causal BlockMask's mask_mod closes over per-token id tensors bound to one device. " + "`BlockMask.to()` relocates the mask's own index tensors but not those captures, so sharding a single " + "forward across devices mixes them. Same limitation as AnyFlowFARTransformer3DModel." + ) + def test_model_parallelism(self): + pass + + def test_causal_block_changes_output(self): + """Block-causal attention must actually change the result relative to full attention.""" + inputs = self.get_dummy_inputs() + + torch.manual_seed(0) + causal = self.model_class(**self.get_init_dict(), causal_block=True).to(torch_device).eval() + torch.manual_seed(0) + full = self.model_class(**self.get_init_dict(), causal_block=False).to(torch_device).eval() + + with torch.no_grad(): + causal_out = causal(**inputs, return_dict=False)[0] + full_out = full(**inputs, return_dict=False)[0] + + assert causal_out.shape == full_out.shape + assert not torch.allclose(causal_out, full_out, atol=1e-5) + + def test_kv_cache_matches_full_forward(self): + """ + Decoding from a cache prefilled at a different timestep must match a full forward. This only holds because + `causal_condition` modulates text and condition-image tokens from t=0, making their activations independent of + the denoising step. + """ + inputs = self.get_dummy_inputs() + target_tokens = inputs["hidden_states"].shape[1] + + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + kv_cache = [{} for _ in range(self.get_init_dict()["num_layers"])] + + prefill_inputs = dict(inputs, timestep=torch.tensor([0.9], device=torch_device)) + decode_inputs = dict(inputs, timestep=torch.tensor([0.4], device=torch_device)) + with torch.no_grad(): + model(**prefill_inputs, kv_cache=kv_cache, return_dict=False) + decoded = model(**decode_inputs, kv_cache=kv_cache, return_dict=False)[0] + reference = model(**decode_inputs, return_dict=False)[0] + + assert decoded.shape[1] == target_tokens + torch.testing.assert_close(decoded, reference[:, -target_tokens:], atol=2e-5, rtol=2e-5) + + def test_kv_cache_requires_causal_condition(self): + init_dict = dict(self.get_init_dict(), causal_condition=False) + model = self.model_class(**init_dict).to(torch_device).eval() + with pytest.raises(ValueError, match="causal_condition"): + model(**self.get_dummy_inputs(), kv_cache=[{} for _ in range(init_dict["num_layers"])]) + + def test_non_flex_backend_rejected_when_causal(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model.set_attention_backend("native") + with pytest.raises(ValueError, match="flex"): + model(**self.get_dummy_inputs()) + + +class TestQwenImage21BlockCausalMask: + """ + The mask is `(q_idx >= kv_idx) or same_image_block`: causal over the joint sequence, bidirectional inside each + image block. Verified elementwise, since a wrong mask degrades quality silently rather than raising. + """ + + def _layout(self): + # text, two *adjacent* condition images, text, target image. + img_shapes = [(1, 2, 2), (1, 2, 2), (1, 4, 4)] + vlm_mask = torch.tensor([False] * 3 + [True, True] + [False, False] + [True] * 4) + image_pad_mask = torch.repeat_interleave(vlm_mask, torch.where(vlm_mask, 4, 1)) + return img_shapes, image_pad_mask + + def _dense_mask(self, image_ids, key_valid=None, batch_size=1): + block_mask = build_qwenimage21_block_causal_mask(image_ids, key_valid, batch_size, torch.device("cpu")) + padded = block_mask.shape[-1] + seq_len = image_ids.shape[0] + dense = create_mask(block_mask.mask_mod, batch_size, 1, padded, padded, device=torch.device("cpu")) + return dense[:, 0, :seq_len, :seq_len] + + def test_block_ids_keep_adjacent_condition_images_separate(self): + img_shapes, image_pad_mask = self._layout() + image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes) + + assert int(image_ids.max()) + 1 == 3, "two adjacent condition images must not merge into one block" + assert (image_ids[~image_pad_mask] == -1).all(), "text tokens must carry no block id" + target_len = img_shapes[-1][1] * img_shapes[-1][2] + assert int(target_token_mask.sum()) == target_len + assert target_token_mask[-target_len:].all() + + def test_block_ids_reject_inconsistent_shapes(self): + img_shapes, image_pad_mask = self._layout() + with pytest.raises(ValueError, match="image tokens"): + QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes[:-1]) + + def test_mask_is_causal_across_text_and_full_within_images(self): + img_shapes, image_pad_mask = self._layout() + image_ids, _ = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes) + mask = self._dense_mask(image_ids)[0] + + query = torch.arange(mask.shape[0]).view(-1, 1) + key = torch.arange(mask.shape[0]).view(1, -1) + same_block = (image_ids.view(-1, 1) == image_ids.view(1, -1)) & (image_ids.view(-1, 1) >= 0) + expected = (query >= key) | same_block + torch.testing.assert_close(mask, expected) + + text = (~image_pad_mask).nonzero(as_tuple=True)[0] + assert (mask[text][:, text] == (text.view(-1, 1) >= text.view(1, -1))).all() + + first = (image_ids == 0).nonzero(as_tuple=True)[0] + second = (image_ids == 1).nonzero(as_tuple=True)[0] + assert mask[first][:, first].all() and mask[second][:, second].all() + assert not mask[first][:, second].any(), "an image block must not see a later block" + assert mask[second][:, first].all(), "a later block must see an earlier one" + + def test_padded_text_is_never_attended(self): + img_shapes, image_pad_mask = self._layout() + image_ids, _ = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes) + + key_valid = torch.ones(1, image_ids.shape[0], dtype=torch.bool) + key_valid[0, 1] = False + mask = self._dense_mask(image_ids, key_valid)[0] + + assert not mask[:, 1].any(), "a padded position must never be attended as a key" + assert mask[1].any(), "a padded position must keep at least one key so its row is not fully masked" + + +class TestQwenImage21TransformerMemory(QwenImage21TransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestQwenImage21TransformerTraining(QwenImage21TransformerTesterConfig, TrainingTesterMixin): + pass + + +class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, AttentionTesterMixin): + pass From 35b34cdcac876e126de5b0f0b7e511fd002186e6 Mon Sep 17 00:00:00 2001 From: naykun Date: Mon, 14 Sep 2026 17:01:19 +0800 Subject: [PATCH 02/20] fix style --- docs/source/en/_toctree.yml | 4 +- examples/qwenimage21/run_qwenimage21.py | 152 ++++++++++++++++++ src/diffusers/__init__.py | 4 +- .../transformers/transformer_qwenimage21.py | 8 +- .../qwenimage21/pipeline_qwenimage21.py | 14 +- 5 files changed, 167 insertions(+), 15 deletions(-) create mode 100644 examples/qwenimage21/run_qwenimage21.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 3bc97543ab9b..525734a2e4cb 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -641,10 +641,10 @@ title: PRX - local: api/pipelines/prx_pixel title: PRX Pixel - - local: api/pipelines/qwenimage - title: QwenImage - local: api/pipelines/qwenimage21 title: Qwen-Image 2.1 + - local: api/pipelines/qwenimage + title: QwenImage - local: api/pipelines/sana title: Sana - local: api/pipelines/sana_sprint diff --git a/examples/qwenimage21/run_qwenimage21.py b/examples/qwenimage21/run_qwenimage21.py new file mode 100644 index 000000000000..68a2cbae775e --- /dev/null +++ b/examples/qwenimage21/run_qwenimage21.py @@ -0,0 +1,152 @@ +"""Sample script for Qwen-Image 2.1: text-to-image and image editing. + +Both tasks run 40 denoising steps with classifier-free guidance off, so each step is a single +forward pass through the transformer, and with the prefix KV cache on. + + # both tasks; `edit` reuses the text-to-image result as its condition image + python run_qwenimage21.py --model Qwen/Qwen-Image-2.1 + + # editing your own image + python run_qwenimage21.py --task edit --image cat.png --edit-prompt "make it snow" + +Requires `diffusers` with Qwen-Image 2.1 support, plus `transformers`, `accelerate` and `torch`. +Block-causal attention uses `torch.nn.attention.flex_attention` when the installed torch provides +it; otherwise the script still runs on the two-pass SDPA fallback, which is exact for the target +image but only approximate for condition images. +""" + +import argparse +import time +from pathlib import Path + +import torch + +from diffusers import QwenImage21Pipeline +from diffusers.utils import load_image + + +T2I_PROMPT = "A capybara wearing a wizard hat, reading a book by candlelight, oil painting" +EDIT_PROMPT = "Move the capybara to a snowy mountain top at sunrise, keep the wizard hat" + +DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + +# Classifier-free guidance is off: `true_cfg_scale <= 1` skips the negative-prompt pass entirely. +TRUE_CFG_SCALE = 1.0 + +# The text and condition-image prefix is modulated from `t = 0`, so its keys and values are step-independent and are +# cached after the first step. Required for these samples, so it is not exposed as a flag. +USE_KV_CACHE = True + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model", default="Qwen/Qwen-Image-2.1", help="Hub repo id or local checkpoint directory.") + parser.add_argument( + "--task", choices=["t2i", "edit", "both"], default="both", help="Which sample(s) to run. Default: both." + ) + parser.add_argument("--prompt", default=T2I_PROMPT, help="Text-to-image prompt.") + parser.add_argument("--edit-prompt", default=EDIT_PROMPT, help="Instruction applied to the condition image.") + parser.add_argument( + "--image", + default=None, + help="Condition image for `edit` (path or URL). Defaults to the text-to-image result, " + "which is generated first if needed.", + ) + parser.add_argument("--steps", type=int, default=40, help="Denoising steps. Default: 40.") + parser.add_argument("--resolution", type=int, default=2048, help="Target side length in pixels. Default: 1024.") + parser.add_argument("--seed", type=int, default=0, help="Seed for the latent noise.") + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--dtype", choices=list(DTYPES), default="bf16") + parser.add_argument("--output-dir", type=Path, default=Path("qwenimage21_samples")) + parser.add_argument( + "--cpu-offload", + action="store_true", + help="Keep components on CPU until needed. Cuts peak VRAM at the cost of speed.", + ) + return parser.parse_args() + + +def flex_attention_available(): + try: + from torch.nn.attention.flex_attention import flex_attention # noqa: F401 + except ImportError: + return False + return True + + +def load_pipeline(args): + print(f"Loading {args.model} ({args.dtype}) ...") + pipe = QwenImage21Pipeline.from_pretrained(args.model, torch_dtype=DTYPES[args.dtype]) + # The pipeline silently drops the cache when the checkpoint sets `causal_condition=False`, since the prefix is then + # modulated from the sampled timestep and its keys and values change every step. Fail instead of running uncached. + if not pipe.transformer.config.causal_condition: + raise ValueError( + f"{args.model} was configured with `causal_condition=False`, which makes the prefix KV cache invalid. " + "This script requires the cache." + ) + if args.cpu_offload: + pipe.enable_model_cpu_offload(device=args.device) + else: + pipe.to(args.device) + return pipe + + +def run(pipe, args, tag, prompt, image=None): + label = "edit" if image is not None else "text-to-image" + print(f"\n[{label}] {args.steps} steps, true_cfg_scale={TRUE_CFG_SCALE} (CFG off), kv cache on") + print(f"[{label}] prompt: {prompt}") + + # A CPU generator keeps the sample reproducible for a given seed on any device. + generator = torch.Generator(device="cpu").manual_seed(args.seed) + if args.device.startswith("cuda"): + torch.cuda.reset_peak_memory_stats() + + start = time.perf_counter() + result = pipe( + prompt=prompt, + image=image, + num_inference_steps=args.steps, + true_cfg_scale=TRUE_CFG_SCALE, + output_resolution=args.resolution, + use_kv_cache=USE_KV_CACHE, + generator=generator, + ) + elapsed = time.perf_counter() - start + + out_path = args.output_dir / f"{tag}.png" + result.images[0].save(out_path) + report = f"[{label}] {elapsed:.1f}s ({elapsed / args.steps:.2f}s/step) -> {out_path}" + if args.device.startswith("cuda"): + report += f", peak VRAM {torch.cuda.max_memory_allocated() / 1024**3:.1f} GiB" + print(report) + return result.images[0] + + +def main(): + args = parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + + print(f"torch {torch.__version__}, device {args.device}") + print(f"flex_attention: {'available' if flex_attention_available() else 'unavailable, using SDPA fallback'}") + + pipe = load_pipeline(args) + + t2i_image = None + if args.task in ("t2i", "both"): + t2i_image = run(pipe, args, "t2i", args.prompt) + + if args.task in ("edit", "both"): + if args.image is not None: + condition = load_image(args.image) + elif t2i_image is not None: + condition = t2i_image + else: + # `--task edit` on its own with no `--image`: produce a condition image first. + condition = run(pipe, args, "t2i", args.prompt) + run(pipe, args, "edit", args.edit_prompt, image=condition) + + print(f"\nDone. Images written to {args.output_dir.resolve()}") + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 9ab9010573db..cee61551a5c8 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -333,9 +333,9 @@ "PixArtTransformer2DModel", "PriorTransformer", "PRXTransformer2DModel", + "QwenImage21Transformer2DModel", "QwenImageControlNetModel", "QwenImageMultiControlNetModel", - "QwenImage21Transformer2DModel", "QwenImageTransformer2DModel", "SanaControlNetModel", "SanaTransformer2DModel", @@ -766,6 +766,7 @@ "PixArtSigmaPipeline", "PRXPipeline", "PRXPixelPipeline", + "QwenImage21Pipeline", "QwenImageControlNetInpaintPipeline", "QwenImageControlNetPipeline", "QwenImageEditInpaintPipeline", @@ -774,7 +775,6 @@ "QwenImageImg2ImgPipeline", "QwenImageInpaintPipeline", "QwenImageLayeredPipeline", - "QwenImage21Pipeline", "QwenImagePipeline", "ReduxImageEncoder", "SanaControlNetPipeline", diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 3788f193cad7..a5028789f1e6 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -182,8 +182,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class QwenImage21AdaLayerNormContinuous(nn.Module): r""" - Final adaptive norm. Scale only — this variant emits no shift, so `linear` maps to `embedding_dim` rather than - `2 * embedding_dim`. + Final adaptive norm. Scale only — this variant emits no shift, so `linear` maps to `embedding_dim` rather than `2 * + embedding_dim`. """ def __init__(self, embedding_dim: int, conditioning_embedding_dim: int, eps: float = 1e-6): @@ -686,8 +686,8 @@ def forward( img_mask (`torch.Tensor` of shape `(batch_size, vlm_sequence_length)`): `True` at the vision-language encoder's image slots, each standing for a `2x2` group of latent tokens. encoder_hidden_states_mask (`torch.Tensor`, *optional*): - `(batch_size, text_sequence_length)` bool marking valid text tokens. Padded positions are excluded - from attention. + `(batch_size, text_sequence_length)` bool marking valid text tokens. Padded positions are excluded from + attention. kv_cache (`list[dict[str, torch.Tensor]]`, *optional*): One dict per block. Empty dicts prefill the text and condition-image keys and values; populated dicts switch to decode, where only the target image's tokens are recomputed. Requires `causal_condition`. diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index ad98fa84e59b..985cf4cd70f1 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -230,9 +230,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor def _downsample_image_pad_tokens(hidden_states_list, image_pad_mask_list): """Collapse consecutive `<|image_pad|>` tokens into one per contiguous region. - The vision-language processor expands each condition image into many vision tokens, but the - transformer expects one token per image slot, which it then expands 4x. Keep the first token of - each contiguous image-pad region and drop the rest. + The vision-language processor expands each condition image into many vision tokens, but the transformer expects + one token per image slot, which it then expands 4x. Keep the first token of each contiguous image-pad region + and drop the rest. """ out_hs, out_mask = [], [] for hidden_state, pad_mask in zip(hidden_states_list, image_pad_mask_list): @@ -569,10 +569,10 @@ def __call__( `causal_condition` modulates those tokens from `t = 0`, making their activations step-independent. Toggling this does not reproduce the same image bit-for-bit in reduced precision. Caching makes the - decode step attend with a different sequence layout than the prefill step, so the two tile - differently and land on different rounding; both agree with an fp32 reference to the same tolerance. - A one-ULP difference at the first block is then amplified by 32 blocks and every sampler step, so the - two settings give equally valid but visibly distinct samples. Fix a sample by fixing this flag. + decode step attend with a different sequence layout than the prefill step, so the two tile differently + and land on different rounding; both agree with an fp32 reference to the same tolerance. A one-ULP + difference at the first block is then amplified by 32 blocks and every sampler step, so the two + settings give equally valid but visibly distinct samples. Fix a sample by fixing this flag. Examples: From b25d7e21814011267d31183e4187c4b5d0ae8a2e Mon Sep 17 00:00:00 2001 From: naykun Date: Tue, 15 Sep 2026 22:59:51 +0800 Subject: [PATCH 03/20] refactor: address PR review feedback for Qwen-Image 2.1 - Remove `causal_block` config flag (always on for released checkpoint) - Split attention into QwenImage21FlexAttnProcessor and QwenImage21SDPAAttnProcessor - Replace two-pass approximate SDPA prefill with exact multi-pass prefill (each image block gets bidirectional attention, text gets causal mask) - Refactor KV cache to QwenImage21KVCache/QwenImage21KVLayerCache classes with explicit kv_cache_mode="extract"/"cached"/"extend" - Lazy-compile flex_attention on first use (fixes OOM on uncompiled path) - Fix edit pipeline: remove broken _downsample_image_pad_tokens, add mm_token_type_ids for transformers 5.x, auto-convert RGB to RGBA - Use @apply_lora_scale decorator, extract _IMG_TOKENS_PER_SLOT constant - Add # Copied from markers for retrieve_latents and _encode_vae_image - Fix mutable default feat_idx=[0] in all 8 VAE forward methods - Update docs: add usage snippet, remove stale causal_block references - Delete examples/qwenimage21/ (snippet moved to docs) --- .../api/models/qwenimage21_transformer2d.md | 9 +- docs/source/en/api/pipelines/qwenimage21.md | 28 +- examples/qwenimage21/run_qwenimage21.py | 152 ------- .../autoencoder_kl_qwenimage21.py | 32 +- .../transformers/transformer_qwenimage21.py | 405 ++++++++++++------ .../qwenimage21/pipeline_qwenimage21.py | 28 +- .../test_models_transformer_qwenimage21.py | 32 +- 7 files changed, 350 insertions(+), 336 deletions(-) delete mode 100644 examples/qwenimage21/run_qwenimage21.py diff --git a/docs/source/en/api/models/qwenimage21_transformer2d.md b/docs/source/en/api/models/qwenimage21_transformer2d.md index 41cc29cc41e8..278572747937 100644 --- a/docs/source/en/api/models/qwenimage21_transformer2d.md +++ b/docs/source/en/api/models/qwenimage21_transformer2d.md @@ -14,11 +14,12 @@ specific language governing permissions and limitations under the License. --> The single-stream transformer used by Qwen-Image 2.1. Text and image latents share one sequence, and a single shared `modulation` projection feeds every block. -Two config flags set 2.1 apart from earlier QwenImage transformers. Neither adds parameters: +Two behaviours distinguish 2.1 from earlier QwenImage transformers: -- `causal_block` — attention follows `(q_idx >= kv_idx) or same_image_block`, so the joint sequence is causal while - each image block (every condition image and the target image) stays internally bidirectional. This requires the - `flex` attention backend, since the mask is a `torch.nn.attention.flex_attention.BlockMask`. +- **Block-causal attention** — attention follows `(q_idx >= kv_idx) or same_image_block`, so the joint sequence is + causal while each image block stays internally bidirectional. The `flex` attention backend gives efficient + single-pass attention; without it the model uses an exact multi-pass SDPA prefill that processes each block + separately. Both paths produce the same results. - `causal_condition` — text and condition-image tokens are modulated from `t = 0` rather than the sampled timestep. Their activations are therefore independent of the denoising step, which is what makes the keys and values of that prefix cacheable across steps via the `kv_cache` argument. diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md index f18a0458420a..a20da77ecb68 100644 --- a/docs/source/en/api/pipelines/qwenimage21.md +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -13,23 +13,27 @@ specific language governing permissions and limitations under the License. --> Qwen-Image 2.1 encodes the prompt and any condition images together with a Qwen3-VL model, then denoises the target image with a single-stream block-causal transformer. See -[`QwenImage21Transformer2DModel`](../models/qwenimage21_transformer2d) for what `causal_block` and `causal_condition` -change. +[`QwenImage21Transformer2DModel`](../models/qwenimage21_transformer2d) for details on block-causal attention and +`causal_condition`. -Because the text and condition-image prefix is modulated from `t = 0`, its keys and values do not change between -denoising steps. The pipeline caches them after the first step by default; pass `use_kv_cache=False` to recompute the -full sequence every step. +The `flex` attention backend (`torch.nn.attention.flex_attention`) gives efficient single-pass block-causal attention. +Without it, the model uses an exact multi-pass SDPA prefill that processes each image block with bidirectional +attention and text with causal attention, matching the block-causal mask exactly. Both paths produce the same results. -Toggling `use_kv_cache` does not reproduce the same image bit-for-bit in reduced precision. The cached decode step -attends with a different sequence layout than the prefill step, so the two land on different rounding — both match an -fp32 reference to the same tolerance — and a one-ULP difference at the first block is amplified by 32 blocks and every -sampler step. Keep the flag fixed when you need a reproducible sample. +```python +import torch +from diffusers import QwenImage21Pipeline - +pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16).to("cuda") -This pipeline requires the `flex` attention backend when `causal_block` is enabled. +# Text-to-image +image = pipe("A capybara wearing a wizard hat, oil painting", num_inference_steps=40).images[0] +image.save("t2i.png") - +# Image-conditioned editing +edited = pipe("Move it to a snowy mountain top", image=image, num_inference_steps=40).images[0] +edited.save("edit.png") +``` ## QwenImage21Pipeline diff --git a/examples/qwenimage21/run_qwenimage21.py b/examples/qwenimage21/run_qwenimage21.py deleted file mode 100644 index 68a2cbae775e..000000000000 --- a/examples/qwenimage21/run_qwenimage21.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Sample script for Qwen-Image 2.1: text-to-image and image editing. - -Both tasks run 40 denoising steps with classifier-free guidance off, so each step is a single -forward pass through the transformer, and with the prefix KV cache on. - - # both tasks; `edit` reuses the text-to-image result as its condition image - python run_qwenimage21.py --model Qwen/Qwen-Image-2.1 - - # editing your own image - python run_qwenimage21.py --task edit --image cat.png --edit-prompt "make it snow" - -Requires `diffusers` with Qwen-Image 2.1 support, plus `transformers`, `accelerate` and `torch`. -Block-causal attention uses `torch.nn.attention.flex_attention` when the installed torch provides -it; otherwise the script still runs on the two-pass SDPA fallback, which is exact for the target -image but only approximate for condition images. -""" - -import argparse -import time -from pathlib import Path - -import torch - -from diffusers import QwenImage21Pipeline -from diffusers.utils import load_image - - -T2I_PROMPT = "A capybara wearing a wizard hat, reading a book by candlelight, oil painting" -EDIT_PROMPT = "Move the capybara to a snowy mountain top at sunrise, keep the wizard hat" - -DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} - -# Classifier-free guidance is off: `true_cfg_scale <= 1` skips the negative-prompt pass entirely. -TRUE_CFG_SCALE = 1.0 - -# The text and condition-image prefix is modulated from `t = 0`, so its keys and values are step-independent and are -# cached after the first step. Required for these samples, so it is not exposed as a flag. -USE_KV_CACHE = True - - -def parse_args(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--model", default="Qwen/Qwen-Image-2.1", help="Hub repo id or local checkpoint directory.") - parser.add_argument( - "--task", choices=["t2i", "edit", "both"], default="both", help="Which sample(s) to run. Default: both." - ) - parser.add_argument("--prompt", default=T2I_PROMPT, help="Text-to-image prompt.") - parser.add_argument("--edit-prompt", default=EDIT_PROMPT, help="Instruction applied to the condition image.") - parser.add_argument( - "--image", - default=None, - help="Condition image for `edit` (path or URL). Defaults to the text-to-image result, " - "which is generated first if needed.", - ) - parser.add_argument("--steps", type=int, default=40, help="Denoising steps. Default: 40.") - parser.add_argument("--resolution", type=int, default=2048, help="Target side length in pixels. Default: 1024.") - parser.add_argument("--seed", type=int, default=0, help="Seed for the latent noise.") - parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") - parser.add_argument("--dtype", choices=list(DTYPES), default="bf16") - parser.add_argument("--output-dir", type=Path, default=Path("qwenimage21_samples")) - parser.add_argument( - "--cpu-offload", - action="store_true", - help="Keep components on CPU until needed. Cuts peak VRAM at the cost of speed.", - ) - return parser.parse_args() - - -def flex_attention_available(): - try: - from torch.nn.attention.flex_attention import flex_attention # noqa: F401 - except ImportError: - return False - return True - - -def load_pipeline(args): - print(f"Loading {args.model} ({args.dtype}) ...") - pipe = QwenImage21Pipeline.from_pretrained(args.model, torch_dtype=DTYPES[args.dtype]) - # The pipeline silently drops the cache when the checkpoint sets `causal_condition=False`, since the prefix is then - # modulated from the sampled timestep and its keys and values change every step. Fail instead of running uncached. - if not pipe.transformer.config.causal_condition: - raise ValueError( - f"{args.model} was configured with `causal_condition=False`, which makes the prefix KV cache invalid. " - "This script requires the cache." - ) - if args.cpu_offload: - pipe.enable_model_cpu_offload(device=args.device) - else: - pipe.to(args.device) - return pipe - - -def run(pipe, args, tag, prompt, image=None): - label = "edit" if image is not None else "text-to-image" - print(f"\n[{label}] {args.steps} steps, true_cfg_scale={TRUE_CFG_SCALE} (CFG off), kv cache on") - print(f"[{label}] prompt: {prompt}") - - # A CPU generator keeps the sample reproducible for a given seed on any device. - generator = torch.Generator(device="cpu").manual_seed(args.seed) - if args.device.startswith("cuda"): - torch.cuda.reset_peak_memory_stats() - - start = time.perf_counter() - result = pipe( - prompt=prompt, - image=image, - num_inference_steps=args.steps, - true_cfg_scale=TRUE_CFG_SCALE, - output_resolution=args.resolution, - use_kv_cache=USE_KV_CACHE, - generator=generator, - ) - elapsed = time.perf_counter() - start - - out_path = args.output_dir / f"{tag}.png" - result.images[0].save(out_path) - report = f"[{label}] {elapsed:.1f}s ({elapsed / args.steps:.2f}s/step) -> {out_path}" - if args.device.startswith("cuda"): - report += f", peak VRAM {torch.cuda.max_memory_allocated() / 1024**3:.1f} GiB" - print(report) - return result.images[0] - - -def main(): - args = parse_args() - args.output_dir.mkdir(parents=True, exist_ok=True) - - print(f"torch {torch.__version__}, device {args.device}") - print(f"flex_attention: {'available' if flex_attention_available() else 'unavailable, using SDPA fallback'}") - - pipe = load_pipeline(args) - - t2i_image = None - if args.task in ("t2i", "both"): - t2i_image = run(pipe, args, "t2i", args.prompt) - - if args.task in ("edit", "both"): - if args.image is not None: - condition = load_image(args.image) - elif t2i_image is not None: - condition = t2i_image - else: - # `--task edit` on its own with no `--image`: produce a condition image first. - condition = run(pipe, args, "t2i", args.prompt) - run(pipe, args, "edit", args.edit_prompt, image=condition) - - print(f"\nDone. Images written to {args.output_dir.resolve()}") - - -if __name__ == "__main__": - main() diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index 096dcdfd280b..34525fae329a 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -259,7 +259,9 @@ def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: else: self.resample = nn.Identity() - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] b, c, t, h, w = x.size() if self.mode == "upsample3d": if feat_cache is not None: @@ -336,7 +338,9 @@ def __init__( self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1) self.conv_shortcut = QwenImage21CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] # Apply shortcut connection h = self.conv_shortcut(x) @@ -449,7 +453,9 @@ def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", self.gradient_checkpointing = False - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] # First residual block x = self.resnets[0](x, feat_cache=feat_cache, feat_idx=feat_idx) @@ -489,7 +495,9 @@ def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample else: self.downsampler = None - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] x_copy = x.clone() for resnet in self.resnets: x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) @@ -580,7 +588,9 @@ def __init__( self.gradient_checkpointing = False - def forward(self, x, feat_cache=None, feat_idx=[0]): + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -677,7 +687,9 @@ def __init__( self.gradient_checkpointing = False - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] """ Forward pass through the upsampling block. @@ -752,7 +764,9 @@ def __init__( self.gradient_checkpointing = False - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=None): + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=None): + if feat_idx is None: + feat_idx = [0] """ Forward pass through the upsampling block. @@ -868,7 +882,9 @@ def __init__( self.gradient_checkpointing = False - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] ## conv1 if feat_cache is not None: idx = feat_idx[0] diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index a5028789f1e6..760117a0ae76 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -21,7 +21,8 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import logging +from ...utils.peft_utils import apply_lora_scale from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn @@ -34,22 +35,80 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# Each vision-language image slot represents a 2×2 group of latent tokens. +_IMG_TOKENS_PER_SLOT = 4 + # `create_block_mask` quantizes the mask to 128-token blocks. _FLEX_BLOCK_SIZE = 128 -# flex_attention is optional. When available and `causal_block=True`, we use a compiled -# flex_attention with a BlockMask for efficient block-causal attention. When unavailable, -# we fall back to a two-pass prefill: causal attention over the prefix, then full attention -# over the target image attending to the cached prefix + itself. +# flex_attention is optional. When available we use a compiled flex_attention with a BlockMask for +# efficient single-pass block-causal attention. When unavailable, we fall back to an exact multi-pass +# SDPA prefill that processes each image block with bidirectional attention and text segments with +# causal attention, matching the block-causal mask exactly. _FLEX_AVAILABLE = False _compiled_flex_attention = None try: from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention - _compiled_flex_attention = torch.compile(flex_attention) _FLEX_AVAILABLE = True except ImportError: BlockMask = None + flex_attention = None + + +def _get_compiled_flex_attention(): + """Return a compiled flex_attention, compiling on first call. + + Compiling is required for the block-sparse kernel that avoids materializing the full Q@K^T matrix. Without it + flex_attention falls back to a dense fp32 math path that OOMs on long sequences. + """ + global _compiled_flex_attention + if _compiled_flex_attention is None: + _compiled_flex_attention = torch.compile(flex_attention) + return _compiled_flex_attention + + +class QwenImage21KVLayerCache: + """Per-layer KV cache for text and condition-image prefix tokens. + + Stores K and V projections (post-RoPE) for the prefix extracted during the first denoising step. Tensor format: + ``(batch_size, num_prefix_tokens, num_heads, head_dim)``. + """ + + def __init__(self): + self.k: torch.Tensor | None = None + self.v: torch.Tensor | None = None + + def store(self, k: torch.Tensor, v: torch.Tensor): + self.k = k + self.v = v + + def get(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.k is None: + raise RuntimeError("KV cache has not been populated yet.") + return self.k, self.v + + @property + def is_populated(self) -> bool: + return self.k is not None + + def clear(self): + self.k = None + self.v = None + + +class QwenImage21KVCache: + """Container for all transformer blocks' prefix KV caches.""" + + def __init__(self, num_layers: int): + self.layer_caches = [QwenImage21KVLayerCache() for _ in range(num_layers)] + + def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache: + return self.layer_caches[layer_idx] + + def clear(self): + for cache in self.layer_caches: + cache.clear() # Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen @@ -274,70 +333,83 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): ) -class QwenImage21AttnProcessor: +def _qwenimage21_prepare_qkv( + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor | None, + layer_cache: QwenImage21KVLayerCache | None, + kv_cache_mode: str | None, + cache_write_slice: slice | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Shared QKV projection, norm, RoPE and KV-cache bookkeeping for both processors.""" + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query).to(value.dtype) + key = attn.norm_k(key).to(value.dtype) + + if rotary_emb is not None: + query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) + key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) + + if layer_cache is not None: + if kv_cache_mode == "extract" and cache_write_slice is not None: + layer_cache.store( + key[:, cache_write_slice].contiguous(), + value[:, cache_write_slice].contiguous(), + ) + elif kv_cache_mode == "cached": + cached_k, cached_v = layer_cache.get() + key = torch.cat([cached_k, key], dim=1) + value = torch.cat([cached_v, value], dim=1) + elif kv_cache_mode == "extend": + # Read existing cache (if any), prepend to current KV for attention, then store the full + # concatenated KV back. Used by the multi-pass SDPA prefill to accumulate segment-by-segment. + if layer_cache.is_populated: + cached_k, cached_v = layer_cache.get() + key = torch.cat([cached_k, key], dim=1) + value = torch.cat([cached_v, value], dim=1) + layer_cache.store(key.contiguous(), value.contiguous()) + + seq_len_q = query.shape[1] + return query, key, value, seq_len_q + + +class QwenImage21FlexAttnProcessor: r""" - Single-stream attention processor for Qwen-Image 2.1. Text and image tokens share one sequence, so there is no - separate context projection. - - Two attention paths are supported: + Attention processor for Qwen-Image 2.1 using compiled `flex_attention` with a `BlockMask` for exact block-causal + attention. This is the recommended path and is required for high resolutions (2048²+) where a dense score matrix + would OOM. - - **flex** (default when available): a compiled `flex_attention` with a `BlockMask` for efficient block-sparse - block-causal attention. Required for high resolutions (2048²+) where the dense score matrix would OOM. - - **SDPA fallback** (when flex is unavailable): the block-causal mask is implemented via a two-pass prefill - orchestrated by the model's `forward` — pass 1 runs the prefix with `is_causal=True`, pass 2 runs the target - image attending fully to the cached prefix + itself. The processor receives `is_causal` and a padding mask. + ``flex_attention`` is compiled on the first forward pass so the block-sparse kernel is used instead of the dense + fallback. The first call will be slower due to compilation. """ - _attention_backend = "flex" if _FLEX_AVAILABLE else None + _attention_backend = "flex" _parallel_config = None - _SUPPORTED_FLEX_BACKENDS = ("flex", "_native_flex") - def __call__( self, attn: "QwenImage21Attention", hidden_states: torch.Tensor, attention_mask: Any | None = None, rotary_emb: torch.Tensor | None = None, - kv_cache: dict[str, torch.Tensor] | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, is_causal: bool = False, ) -> torch.Tensor: - if ( - _FLEX_AVAILABLE - and isinstance(attention_mask, BlockMask) - and self._attention_backend not in self._SUPPORTED_FLEX_BACKENDS - ): - raise ValueError( - f"QwenImage21AttnProcessor requires the 'flex' attention backend when a BlockMask is used " - f"(got {self._attention_backend!r})." - ) - - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) - - query = attn.norm_q(query).to(value.dtype) - key = attn.norm_k(key).to(value.dtype) - - if rotary_emb is not None: - query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) - key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) - - if kv_cache is not None: - if cache_write_slice is not None: - kv_cache["key"] = key[:, cache_write_slice].contiguous() - kv_cache["value"] = value[:, cache_write_slice].contiguous() - else: - key = torch.cat([kv_cache["key"], key], dim=1) - value = torch.cat([kv_cache["value"], value], dim=1) + query, key, value, seq_len_q = _qwenimage21_prepare_qkv( + attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice + ) - seq_len_q, seq_len_kv = query.shape[1], key.shape[1] - if _FLEX_AVAILABLE and isinstance(attention_mask, BlockMask): + seq_len_kv = key.shape[1] + if isinstance(attention_mask, BlockMask): pad_q = int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_q pad_kv = int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_kv if pad_q: @@ -346,13 +418,14 @@ def __call__( key = F.pad(key.transpose(1, 3), (0, pad_kv)).transpose(1, 3) value = F.pad(value.transpose(1, 3), (0, pad_kv)).transpose(1, 3) - hidden_states = _compiled_flex_attention( + hidden_states = _get_compiled_flex_attention()( query.transpose(1, 2).contiguous(), key.transpose(1, 2).contiguous(), value.transpose(1, 2).contiguous(), block_mask=attention_mask, ).transpose(1, 2) else: + # Decode path or no BlockMask: full attention via SDPA. hidden_states = dispatch_attention_fn( query, key, @@ -370,14 +443,59 @@ def __call__( return attn.to_out[1](hidden_states) +class QwenImage21SDPAAttnProcessor: + r""" + SDPA attention processor for Qwen-Image 2.1. Use this when `flex_attention` is not available. + + The block-causal mask is implemented exactly via a multi-pass prefill orchestrated by the model's `forward`: each + image block in the prefix is processed with bidirectional attention within the block and full attention to all + preceding segments, and text segments get a causal mask. The target image attends fully to the cached prefix + + itself. This matches the block-causal mask without requiring `flex_attention`. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + attention_mask: Any | None = None, + rotary_emb: torch.Tensor | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + is_causal: bool = False, + ) -> torch.Tensor: + query, key, value, seq_len_q = _qwenimage21_prepare_qkv( + attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice + ) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask if not isinstance(attention_mask, type(None)) and not is_causal else None, + dropout_p=0.0, + is_causal=is_causal, + backend=None, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states[:, :seq_len_q] + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + class QwenImage21Attention(torch.nn.Module, AttentionModuleMixin): r""" Attention module for [`QwenImage21TransformerBlock`]. Projection layout matches the legacy [`~models.attention_processor.Attention`] so Qwen-Image 2.x checkpoints load into it unchanged. """ - _default_processor_cls = QwenImage21AttnProcessor - _available_processors = [QwenImage21AttnProcessor] + _default_processor_cls = QwenImage21FlexAttnProcessor if _FLEX_AVAILABLE else QwenImage21SDPAAttnProcessor + _available_processors = [QwenImage21FlexAttnProcessor, QwenImage21SDPAAttnProcessor] def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-6, processor: Any | None = None): super().__init__() @@ -438,7 +556,8 @@ def forward( rotary_emb: torch.Tensor | None = None, attention_mask: Any | None = None, target_token_mask: torch.Tensor | None = None, - kv_cache: dict[str, torch.Tensor] | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, is_causal: bool = False, ) -> torch.Tensor: @@ -449,7 +568,8 @@ def forward( hidden_states=img_modulated, attention_mask=attention_mask, rotary_emb=rotary_emb, - kv_cache=kv_cache, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, cache_write_slice=cache_write_slice, is_causal=is_causal, ) @@ -535,10 +655,11 @@ class QwenImage21Transformer2DModel( positions the vision-language encoder reserved for them, and the target image's tokens are appended. A single shared `modulation` projection feeds every block, so blocks hold no modulation parameters of their own. - Two behaviours distinguish 2.1 from 2.0, both switched on by config and neither adding parameters: + Two behaviours distinguish 2.1: - - `causal_block` — attention follows `(q_idx >= kv_idx) or same_image_block`, so the sequence is causal while each - image block stays internally bidirectional. This requires the `flex` attention backend. + - **Block-causal attention** — attention follows `(q_idx >= kv_idx) or same_image_block`, so the sequence is causal + while each image block stays internally bidirectional. The `flex` attention backend gives efficient single-pass + attention; without it the model uses an exact multi-pass SDPA prefill that processes each block separately. - `causal_condition` — text and condition-image tokens are modulated from `t = 0` instead of the sampled timestep, which also makes their activations timestep-independent and so cacheable across denoising steps. @@ -565,8 +686,6 @@ class QwenImage21Transformer2DModel( Epsilon for the norm layers. causal_condition (`bool`, defaults to `True`): Modulate text and condition-image tokens from `t = 0`. Required for KV caching. - causal_block (`bool`, defaults to `True`): - Use block-causal attention. Requires the `flex` attention backend. """ _supports_gradient_checkpointing = True @@ -589,7 +708,6 @@ def __init__( axes_dims_rope: tuple[int, int, int] = (16, 56, 56), eps: float = 1e-6, causal_condition: bool = True, - causal_block: bool = True, ): super().__init__() self.out_channels = out_channels or in_channels @@ -660,6 +778,7 @@ def build_token_metadata( target_token_mask[image_positions[-block_lengths[-1] :]] = True return image_ids, target_token_mask + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -669,7 +788,8 @@ def forward( img_mask: torch.Tensor, encoder_hidden_states_mask: torch.Tensor | None = None, attention_kwargs: dict[str, Any] | None = None, - kv_cache: list[dict[str, torch.Tensor]] | None = None, + kv_cache: QwenImage21KVCache | None = None, + kv_cache_mode: str | None = None, return_dict: bool = True, ) -> torch.Tensor | Transformer2DModelOutput: r""" @@ -688,20 +808,12 @@ def forward( encoder_hidden_states_mask (`torch.Tensor`, *optional*): `(batch_size, text_sequence_length)` bool marking valid text tokens. Padded positions are excluded from attention. - kv_cache (`list[dict[str, torch.Tensor]]`, *optional*): - One dict per block. Empty dicts prefill the text and condition-image keys and values; populated dicts - switch to decode, where only the target image's tokens are recomputed. Requires `causal_condition`. + kv_cache (`QwenImage21KVCache`, *optional*): + Cache container. Pass together with `kv_cache_mode` to enable prefix KV caching. + kv_cache_mode (`str`, *optional*): + `"extract"` to prefill the cache (first denoising step), `"cached"` to decode from it (later steps). + Requires `causal_condition=True`. """ - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - scale_lora_layers(self, lora_scale) - elif attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.") batch_size = hidden_states.shape[0] hidden_states = self.img_in(hidden_states) @@ -709,7 +821,7 @@ def forward( # Each vision-language image slot stands for 2x2 latent tokens, so expand those positions four-fold and drop # the actual latents into them. Samples share a layout, hence the single row. - repeats = torch.where(img_mask, 4, 1)[0] + repeats = torch.where(img_mask, _IMG_TOKENS_PER_SLOT, 1)[0] image_pad_mask = torch.repeat_interleave(img_mask[0], repeats) target_tokens = math.prod(img_shapes[0][-1]) @@ -742,6 +854,10 @@ def forward( "kv_cache requires `causal_condition=True`. The cache is only valid because text and condition-image " "tokens modulate from t=0, which makes their activations independent of the denoising step." ) + if kv_cache is not None and kv_cache_mode not in ("extract", "cached"): + raise ValueError( + f"kv_cache_mode must be 'extract' or 'cached' when kv_cache is provided, got {kv_cache_mode!r}." + ) # Right-padded prompt positions must never be attended to, on any path. Text positions of the joint sequence # line up, in order, with the non-image positions of the vision-language sequence — the two are interleaved, @@ -756,8 +872,7 @@ def forward( joint_key_valid[:, text_positions] = encoder_hidden_states_mask.bool()[:, vlm_text_positions] prefix_len = int((~target_token_mask).sum()) - is_decode = kv_cache is not None and len(kv_cache[0]) > 0 - use_flex = _FLEX_AVAILABLE and self.config.causal_block + is_decode = kv_cache_mode == "cached" if is_decode: # Only the target image's queries are recomputed. The block-causal mask degenerates to full attention @@ -767,74 +882,107 @@ def forward( modulation_mask = modulation_mask[prefix_len:] attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] cache_write_slice = None - use_two_pass = False - elif use_flex: + use_multi_pass = False + elif _FLEX_AVAILABLE: # flex path: single-pass with a compiled BlockMask - cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None + cache_write_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None attention_mask = build_qwenimage21_block_causal_mask( image_ids, joint_key_valid, batch_size, hidden_states.device ) - use_two_pass = False - elif self.config.causal_block: - # No flex_attention available: approximate the block-causal mask with a two-pass prefill. - # Pass 1 runs the prefix (text + condition images) causally and caches its keys and values; - # pass 2 runs the target image attending fully over that prefix plus itself. - # - # This is exact for the target image and for text, but not for condition images: under the real - # block-causal mask a condition image attends within its own block bidirectionally, whereas a - # single causal pass only lets each of its tokens see earlier ones. Text-only prompts are - # therefore unaffected; prompts with condition images differ slightly. Install a PyTorch build - # with flex_attention for the exact mask. - use_two_pass = True - cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None + use_multi_pass = False else: - # causal_block disabled: full attention - cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None - attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] - use_two_pass = False + # Multi-pass SDPA: exact block-causal attention without flex_attention. + # The prefix is split into segments at image-block boundaries (using image_ids). Image-block + # segments get full (bidirectional) attention within themselves, and text segments get a causal + # mask within the segment. Both attend fully to all preceding segments via an accumulating KV + # cache ("extend" mode). This exactly matches the block-causal mask. + # The target image pass is unchanged: full attention over [cached prefix, target]. + use_multi_pass = True + cache_write_slice = None - if use_two_pass: - # Two-pass prefill: split into prefix and target, run prefix causally then target fully. + if use_multi_pass: prefix_hs = joint_hidden_states[:, :prefix_len] target_hs = joint_hidden_states[:, prefix_len:] prefix_rope = rotary_emb[:prefix_len] target_rope = rotary_emb[prefix_len:] prefix_mod_mask = modulation_mask[:prefix_len] if modulation_mask is not None else None target_mod_mask = modulation_mask[prefix_len:] if modulation_mask is not None else None + + # Build segment boundaries from image_ids in the prefix. Consecutive tokens with the same + # image_id form one segment (-1 = text, >=0 = image block). + prefix_ids = image_ids[:prefix_len] + segments = [] + if prefix_len > 0: + seg_start = 0 + for i in range(1, prefix_len): + if prefix_ids[i] != prefix_ids[i - 1]: + segments.append((seg_start, i)) + seg_start = i + segments.append((seg_start, prefix_len)) + for index_block, block in enumerate(self.transformer_blocks): - block_kv_cache = kv_cache[index_block] if kv_cache is not None else None - # Pass 1: prefix with causal attention. No padding mask — padded text tokens have zero embeddings - # from right-padding and is_causal handles the structural mask. - prefix_cache = {} if kv_cache is not None else None - prefix_hs = block( - hidden_states=prefix_hs, - modulation=modulation, - rotary_emb=prefix_rope, - attention_mask=None, - target_token_mask=prefix_mod_mask, - kv_cache=prefix_cache, - cache_write_slice=slice(0, prefix_len), - is_causal=True, + layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None + accumulated_cache = QwenImage21KVLayerCache() + segment_outputs = [] + + for seg_start, seg_end in segments: + seg_hs = prefix_hs[:, seg_start:seg_end] + seg_rope = prefix_rope[seg_start:seg_end] + seg_mod = prefix_mod_mask[seg_start:seg_end] if prefix_mod_mask is not None else None + + # Text segments (image_id == -1) get a causal mask within the segment so that + # each text token only sees earlier text tokens + the full cached prefix. Image + # blocks get None (full / bidirectional), which is exact for the block-causal mask. + seg_mask = None + seg_is_text = prefix_ids[seg_start].item() < 0 + seg_len = seg_end - seg_start + if seg_is_text and seg_len > 1: + cached_len = accumulated_cache.k.shape[1] if accumulated_cache.is_populated else 0 + prefix_visible = torch.ones(seg_len, cached_len, dtype=torch.bool, device=hidden_states.device) + causal_part = torch.tril( + torch.ones(seg_len, seg_len, dtype=torch.bool, device=hidden_states.device) + ) + seg_mask = torch.cat([prefix_visible, causal_part], dim=1)[None, None] + + seg_hs = block( + hidden_states=seg_hs, + modulation=modulation, + rotary_emb=seg_rope, + attention_mask=seg_mask, + target_token_mask=seg_mod, + layer_cache=accumulated_cache, + kv_cache_mode="extend", + cache_write_slice=None, + ) + segment_outputs.append(seg_hs) + + prefix_hs = torch.cat(segment_outputs, dim=1) if segment_outputs else prefix_hs + + # Store the accumulated prefix cache into the main cache for this layer. + if layer_cache is not None and accumulated_cache.is_populated: + layer_cache.store(*accumulated_cache.get()) + + # Target: full attention over [cached prefix, target]. + target_cache = ( + layer_cache + if layer_cache is not None + else (accumulated_cache if accumulated_cache.is_populated else None) ) - # Pass 2: target image with full attention over [cached prefix, target]. - # No attention mask needed: the prefix cache already excludes padded positions, and the target - # image should see everything (block-causal degenerates to full attention for target rows). - if prefix_cache is not None and block_kv_cache is not None: - block_kv_cache.update(prefix_cache) target_hs = block( hidden_states=target_hs, modulation=modulation, rotary_emb=target_rope, attention_mask=None, target_token_mask=target_mod_mask, - kv_cache=block_kv_cache, + layer_cache=target_cache, + kv_cache_mode="cached" if target_cache is not None else None, cache_write_slice=None, ) joint_hidden_states = torch.cat([prefix_hs, target_hs], dim=1) else: for index_block, block in enumerate(self.transformer_blocks): - block_kv_cache = kv_cache[index_block] if kv_cache is not None else None + layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None if torch.is_grad_enabled() and self.gradient_checkpointing: joint_hidden_states = self._gradient_checkpointing_func( block, @@ -843,7 +991,8 @@ def forward( rotary_emb, attention_mask, modulation_mask, - block_kv_cache, + layer_cache, + kv_cache_mode, cache_write_slice, ) else: @@ -853,16 +1002,14 @@ def forward( rotary_emb=rotary_emb, attention_mask=attention_mask, target_token_mask=modulation_mask, - kv_cache=block_kv_cache, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, cache_write_slice=cache_write_slice, ) joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) output = self.proj_out(joint_hidden_states) - if USE_PEFT_BACKEND: - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 985cf4cd70f1..a05a6e2b55b0 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -131,7 +131,10 @@ def retrieve_timesteps( return timesteps, num_inference_steps -def retrieve_latents(encoder_output, generator=None, sample_mode="sample"): +# Copied from diffusers.pipelines.flux.pipeline_flux_control_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": @@ -320,6 +323,8 @@ def _get_qwen_prompt_embeds( } if not is_t2i and hasattr(model_inputs, "pixel_values"): forward_kwargs.update(pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw) + if hasattr(model_inputs, "mm_token_type_ids"): + forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids outputs = self.text_encoder(**forward_kwargs) hidden_states = outputs.hidden_states[-1] @@ -333,11 +338,6 @@ def _get_qwen_prompt_embeds( ] image_pad_mask = [e[drop_idx:] for e in image_pad_mask] - if not is_t2i: - split_hidden_states, image_pad_mask = self._downsample_image_pad_tokens( - split_hidden_states, image_pad_mask - ) - attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max(e.size(0) for e in split_hidden_states) prompt_embeds = torch.stack( @@ -421,6 +421,7 @@ def _unpack_latents(latents, height, width, vae_scale_factor): latents = latents.transpose(1, 2).reshape(batch_size, channels, 1, height, width) return latents + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): if isinstance(generator, list): image_latents = [ @@ -430,7 +431,6 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): image_latents = torch.cat(image_latents, dim=0) else: image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") - latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.latent_channels, 1, 1, 1) @@ -441,7 +441,8 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): .view(1, self.latent_channels, 1, 1, 1) .to(image_latents.device, image_latents.dtype) ) - return (image_latents - latents_mean) / latents_std + image_latents = (image_latents - latents_mean) / latents_std + return image_latents def prepare_latents( self, images, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None @@ -616,6 +617,8 @@ def __call__( image = image if isinstance(image, list) else [image] input_images, vae_images = [], [] for img in image: + if hasattr(img, "mode") and img.mode != "RGBA": + img = img.convert("RGBA") image_width, image_height = img.size input_width, input_height = calculate_dimensions( output_resolution * output_resolution, image_width / image_height @@ -708,10 +711,12 @@ def __call__( # Text and condition-image keys and values are step-independent under `causal_condition`, so the first step # prefills them and later steps only recompute the target image's tokens. + from ...models.transformers.transformer_qwenimage21 import QwenImage21KVCache + num_blocks = len(self.transformer.transformer_blocks) cache_enabled = use_kv_cache and self.transformer.config.causal_condition - cond_cache = [{} for _ in range(num_blocks)] if cache_enabled else None - neg_cache = [{} for _ in range(num_blocks)] if cache_enabled and do_true_cfg else None + cond_cache = QwenImage21KVCache(num_blocks) if cache_enabled else None + neg_cache = QwenImage21KVCache(num_blocks) if cache_enabled and do_true_cfg else None # 5. Denoising loop self.scheduler.set_begin_index(0) @@ -721,6 +726,7 @@ def __call__( continue self._current_timestep = t + kv_mode = "extract" if (cache_enabled and i == 0) else ("cached" if cache_enabled else None) latent_model_input = latents if input_images_latents is not None: @@ -737,6 +743,7 @@ def __call__( img_mask=image_pad_mask, attention_kwargs=self.attention_kwargs, kv_cache=cond_cache, + kv_cache_mode=kv_mode, return_dict=False, )[0] noise_pred = noise_pred[:, -latents.size(1) :] @@ -752,6 +759,7 @@ def __call__( img_mask=negative_image_pad_mask, attention_kwargs=self.attention_kwargs, kv_cache=neg_cache, + kv_cache_mode=kv_mode, return_dict=False, )[0] neg_noise_pred = neg_noise_pred[:, -latents.size(1) :] diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index ae3a43fff77a..8626f943c8f6 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -104,22 +104,6 @@ class TestQwenImage21Transformer(QwenImage21TransformerTesterConfig, ModelTester def test_model_parallelism(self): pass - def test_causal_block_changes_output(self): - """Block-causal attention must actually change the result relative to full attention.""" - inputs = self.get_dummy_inputs() - - torch.manual_seed(0) - causal = self.model_class(**self.get_init_dict(), causal_block=True).to(torch_device).eval() - torch.manual_seed(0) - full = self.model_class(**self.get_init_dict(), causal_block=False).to(torch_device).eval() - - with torch.no_grad(): - causal_out = causal(**inputs, return_dict=False)[0] - full_out = full(**inputs, return_dict=False)[0] - - assert causal_out.shape == full_out.shape - assert not torch.allclose(causal_out, full_out, atol=1e-5) - def test_kv_cache_matches_full_forward(self): """ Decoding from a cache prefilled at a different timestep must match a full forward. This only holds because @@ -129,25 +113,31 @@ def test_kv_cache_matches_full_forward(self): inputs = self.get_dummy_inputs() target_tokens = inputs["hidden_states"].shape[1] + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + torch.manual_seed(0) - model = self.model_class(**self.get_init_dict()).to(torch_device).eval() - kv_cache = [{} for _ in range(self.get_init_dict()["num_layers"])] + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device).eval() + kv_cache = QwenImage21KVCache(init_dict["num_layers"]) prefill_inputs = dict(inputs, timestep=torch.tensor([0.9], device=torch_device)) decode_inputs = dict(inputs, timestep=torch.tensor([0.4], device=torch_device)) with torch.no_grad(): - model(**prefill_inputs, kv_cache=kv_cache, return_dict=False) - decoded = model(**decode_inputs, kv_cache=kv_cache, return_dict=False)[0] + model(**prefill_inputs, kv_cache=kv_cache, kv_cache_mode="extract", return_dict=False) + decoded = model(**decode_inputs, kv_cache=kv_cache, kv_cache_mode="cached", return_dict=False)[0] reference = model(**decode_inputs, return_dict=False)[0] assert decoded.shape[1] == target_tokens torch.testing.assert_close(decoded, reference[:, -target_tokens:], atol=2e-5, rtol=2e-5) def test_kv_cache_requires_causal_condition(self): + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + init_dict = dict(self.get_init_dict(), causal_condition=False) model = self.model_class(**init_dict).to(torch_device).eval() + kv_cache = QwenImage21KVCache(init_dict["num_layers"]) with pytest.raises(ValueError, match="causal_condition"): - model(**self.get_dummy_inputs(), kv_cache=[{} for _ in range(init_dict["num_layers"])]) + model(**self.get_dummy_inputs(), kv_cache=kv_cache, kv_cache_mode="extract") def test_non_flex_backend_rejected_when_causal(self): model = self.model_class(**self.get_init_dict()).to(torch_device).eval() From c17116b058fd6a33b18da5684060789323c5fd3d Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 18:43:54 +0800 Subject: [PATCH 04/20] fix: KV cache pinned the whole prefill sequence at batch size 1 The "extract" branch stored the prefix as `key[:, cache_write_slice].contiguous()`. At batch size 1 that slice already counts as contiguous, because PyTorch ignores size-1 dimensions in the check, so `contiguous()` returned the same view and the cache pinned the full prefill K/V for every step of the denoising loop: 8.0 GiB of resident memory at 2048x2048 across 32 layers. At batch size 2 and above the slice is not contiguous, the copy happens, and the leak disappears, so no test caught it. Store a `clone()` instead, and assert in the tests that the cached prefix owns its storage. Measured at 2048x2048 on one H100, bf16, batch 1, 20 steps: peak memory for a full pipeline call drops from 64.5 GiB to 56.5 GiB. --- .../transformers/transformer_qwenimage21.py | 7 +++++-- .../test_models_transformer_qwenimage21.py | 20 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 760117a0ae76..2bf2cad5477f 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -359,9 +359,12 @@ def _qwenimage21_prepare_qkv( if layer_cache is not None: if kv_cache_mode == "extract" and cache_write_slice is not None: + # `clone()`, not `contiguous()`: at batch size 1 the prefix slice already counts as contiguous + # (size-1 dims are ignored), so `contiguous()` returns the same view and the cache would pin the + # whole prefill K/V for every step of the denoising loop. layer_cache.store( - key[:, cache_write_slice].contiguous(), - value[:, cache_write_slice].contiguous(), + key[:, cache_write_slice].clone(), + value[:, cache_write_slice].clone(), ) elif kv_cache_mode == "cached": cached_k, cached_v = layer_cache.get() diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index 8626f943c8f6..ad22ddba23f0 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -130,6 +130,26 @@ def test_kv_cache_matches_full_forward(self): assert decoded.shape[1] == target_tokens torch.testing.assert_close(decoded, reference[:, -target_tokens:], atol=2e-5, rtol=2e-5) + def test_kv_cache_owns_its_memory(self): + """ + The cached prefix must own its storage. At batch size 1 the prefix slice already counts as contiguous, so + storing `key[:, :prefix].contiguous()` hands the cache a view that pins the whole prefill K/V — 8 GiB at + 2048² — for every step of the denoising loop. + """ + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device).eval() + kv_cache = QwenImage21KVCache(init_dict["num_layers"]) + with torch.no_grad(): + model(**self.get_dummy_inputs(batch_size=1), kv_cache=kv_cache, kv_cache_mode="extract") + + for index in range(init_dict["num_layers"]): + for cached in kv_cache.get_layer(index).get(): + assert cached.untyped_storage().nbytes() == cached.numel() * cached.element_size(), ( + f"layer {index} cached a view into the full prefill K/V instead of a copy" + ) + def test_kv_cache_requires_causal_condition(self): from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache From 9a1a603e2737124bc4139739f8e8ebdd589e1068 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 16 Sep 2026 04:13:57 +0200 Subject: [PATCH 05/20] move the block-causal segmentation into the SDPA processor Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WpU4nuugzCw8T1c4tL6N2A --- .../transformers/transformer_qwenimage21.py | 264 ++++++++---------- .../test_models_transformer_qwenimage21.py | 30 ++ 2 files changed, 147 insertions(+), 147 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 2bf2cad5477f..dfdb0b1027a0 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -370,14 +370,6 @@ def _qwenimage21_prepare_qkv( cached_k, cached_v = layer_cache.get() key = torch.cat([cached_k, key], dim=1) value = torch.cat([cached_v, value], dim=1) - elif kv_cache_mode == "extend": - # Read existing cache (if any), prepend to current KV for attention, then store the full - # concatenated KV back. Used by the multi-pass SDPA prefill to accumulate segment-by-segment. - if layer_cache.is_populated: - cached_k, cached_v = layer_cache.get() - key = torch.cat([cached_k, key], dim=1) - value = torch.cat([cached_v, value], dim=1) - layer_cache.store(key.contiguous(), value.contiguous()) seq_len_q = query.shape[1] return query, key, value, seq_len_q @@ -405,7 +397,9 @@ def __call__( layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, - is_causal: bool = False, + image_ids: torch.Tensor | None = None, + key_valid: torch.Tensor | None = None, + prefix_len: int | None = None, ) -> torch.Tensor: query, key, value, seq_len_q = _qwenimage21_prepare_qkv( attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice @@ -413,6 +407,7 @@ def __call__( seq_len_kv = key.shape[1] if isinstance(attention_mask, BlockMask): + # prefill: the BlockMask expresses the block-causal structure in one flex call pad_q = int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_q pad_kv = int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_kv if pad_q: @@ -428,14 +423,13 @@ def __call__( block_mask=attention_mask, ).transpose(1, 2) else: - # Decode path or no BlockMask: full attention via SDPA. + # decode: full attention over [cached prefix, target] via SDPA hidden_states = dispatch_attention_fn( query, key, value, - attn_mask=attention_mask if not isinstance(attention_mask, type(None)) and not is_causal else None, + attn_mask=attention_mask, dropout_p=0.0, - is_causal=is_causal, backend=None, parallel_config=self._parallel_config, ) @@ -468,22 +462,76 @@ def __call__( layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, - is_causal: bool = False, + image_ids: torch.Tensor | None = None, + key_valid: torch.Tensor | None = None, + prefix_len: int | None = None, ) -> torch.Tensor: query, key, value, seq_len_q = _qwenimage21_prepare_qkv( attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice ) - hidden_states = dispatch_attention_fn( - query, - key, - value, - attn_mask=attention_mask if not isinstance(attention_mask, type(None)) and not is_causal else None, - dropout_p=0.0, - is_causal=is_causal, - backend=None, - parallel_config=self._parallel_config, - ) + if image_ids is None: + # decode: full attention over [cached prefix, target] + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + backend=None, + parallel_config=self._parallel_config, + ) + else: + # prefill: the block-causal mask decomposes into one attention call per prefix segment plus one for the + # target image. Every segment attends to the keys `[0, end)` (everything before it plus its own block); + # text segments additionally get a causal triangle over their own keys; padded text keys are dropped. + # `attention_mask` is the flex BlockMask of the same structure, meant for `QwenImage21FlexAttnProcessor`; + # it is not used here. + prefix_ids = image_ids[:prefix_len].tolist() + segments = [] + start = 0 + for i in range(1, prefix_len + 1): + if i == prefix_len or prefix_ids[i] != prefix_ids[start]: + segments.append((start, i, prefix_ids[start] < 0)) + start = i + outputs = [] + for start, end, is_text in segments: + seg_mask = None + if is_text: + seg_len = end - start + seg_mask = torch.cat( + [ + torch.ones(seg_len, start, dtype=torch.bool, device=query.device), + torch.tril(torch.ones(seg_len, seg_len, dtype=torch.bool, device=query.device)), + ], + dim=1, + )[None, None] + if key_valid is not None: + seg_key_valid = key_valid[:, None, None, :end] + seg_mask = seg_key_valid if seg_mask is None else (seg_mask & seg_key_valid) + outputs.append( + dispatch_attention_fn( + query[:, start:end], + key[:, :end], + value[:, :end], + attn_mask=seg_mask, + dropout_p=0.0, + backend=None, + parallel_config=self._parallel_config, + ) + ) + outputs.append( + dispatch_attention_fn( + query[:, prefix_len:], + key, + value, + attn_mask=None if key_valid is None else key_valid[:, None, None, :], + dropout_p=0.0, + backend=None, + parallel_config=self._parallel_config, + ) + ) + hidden_states = torch.cat(outputs, dim=1) hidden_states = hidden_states[:, :seq_len_q] hidden_states = hidden_states.flatten(2, 3).type_as(query) @@ -562,7 +610,9 @@ def forward( layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, - is_causal: bool = False, + image_ids: torch.Tensor | None = None, + key_valid: torch.Tensor | None = None, + prefix_len: int | None = None, ) -> torch.Tensor: mod1, mod2 = modulation.chunk(2, dim=-1) @@ -574,7 +624,9 @@ def forward( layer_cache=layer_cache, kv_cache_mode=kv_cache_mode, cache_write_slice=cache_write_slice, - is_causal=is_causal, + image_ids=image_ids, + key_valid=key_valid, + prefix_len=prefix_len, ) hidden_states = hidden_states + img_gate1.tanh() * attn_output @@ -875,141 +927,59 @@ def forward( joint_key_valid[:, text_positions] = encoder_hidden_states_mask.bool()[:, vlm_text_positions] prefix_len = int((~target_token_mask).sum()) - is_decode = kv_cache_mode == "cached" - if is_decode: - # Only the target image's queries are recomputed. The block-causal mask degenerates to full attention - # for target rows (they can see the entire prefix + their own block), so no structural mask is needed. + if kv_cache_mode == "cached": + # decode: only the target image's queries are recomputed. The block-causal mask degenerates to full + # attention for target rows (they see the entire prefix + their own block), so only the padding mask is + # needed. joint_hidden_states = joint_hidden_states[:, prefix_len:] rotary_emb = rotary_emb[prefix_len:] modulation_mask = modulation_mask[prefix_len:] attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] cache_write_slice = None - use_multi_pass = False - elif _FLEX_AVAILABLE: - # flex path: single-pass with a compiled BlockMask - cache_write_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None + block_image_ids, block_key_valid, block_prefix_len = None, None, None + else: + # prefill: the whole joint sequence. The block-causal structure is passed down both as a flex + # `BlockMask` (used by `QwenImage21FlexAttnProcessor`) and as its ingredients (`image_ids`, + # `key_valid`, `prefix_len`, used by `QwenImage21SDPAAttnProcessor`); the processor picks. attention_mask = build_qwenimage21_block_causal_mask( image_ids, joint_key_valid, batch_size, hidden_states.device ) - use_multi_pass = False - else: - # Multi-pass SDPA: exact block-causal attention without flex_attention. - # The prefix is split into segments at image-block boundaries (using image_ids). Image-block - # segments get full (bidirectional) attention within themselves, and text segments get a causal - # mask within the segment. Both attend fully to all preceding segments via an accumulating KV - # cache ("extend" mode). This exactly matches the block-causal mask. - # The target image pass is unchanged: full attention over [cached prefix, target]. - use_multi_pass = True - cache_write_slice = None - - if use_multi_pass: - prefix_hs = joint_hidden_states[:, :prefix_len] - target_hs = joint_hidden_states[:, prefix_len:] - prefix_rope = rotary_emb[:prefix_len] - target_rope = rotary_emb[prefix_len:] - prefix_mod_mask = modulation_mask[:prefix_len] if modulation_mask is not None else None - target_mod_mask = modulation_mask[prefix_len:] if modulation_mask is not None else None - - # Build segment boundaries from image_ids in the prefix. Consecutive tokens with the same - # image_id form one segment (-1 = text, >=0 = image block). - prefix_ids = image_ids[:prefix_len] - segments = [] - if prefix_len > 0: - seg_start = 0 - for i in range(1, prefix_len): - if prefix_ids[i] != prefix_ids[i - 1]: - segments.append((seg_start, i)) - seg_start = i - segments.append((seg_start, prefix_len)) - - for index_block, block in enumerate(self.transformer_blocks): - layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None - accumulated_cache = QwenImage21KVLayerCache() - segment_outputs = [] - - for seg_start, seg_end in segments: - seg_hs = prefix_hs[:, seg_start:seg_end] - seg_rope = prefix_rope[seg_start:seg_end] - seg_mod = prefix_mod_mask[seg_start:seg_end] if prefix_mod_mask is not None else None - - # Text segments (image_id == -1) get a causal mask within the segment so that - # each text token only sees earlier text tokens + the full cached prefix. Image - # blocks get None (full / bidirectional), which is exact for the block-causal mask. - seg_mask = None - seg_is_text = prefix_ids[seg_start].item() < 0 - seg_len = seg_end - seg_start - if seg_is_text and seg_len > 1: - cached_len = accumulated_cache.k.shape[1] if accumulated_cache.is_populated else 0 - prefix_visible = torch.ones(seg_len, cached_len, dtype=torch.bool, device=hidden_states.device) - causal_part = torch.tril( - torch.ones(seg_len, seg_len, dtype=torch.bool, device=hidden_states.device) - ) - seg_mask = torch.cat([prefix_visible, causal_part], dim=1)[None, None] - - seg_hs = block( - hidden_states=seg_hs, - modulation=modulation, - rotary_emb=seg_rope, - attention_mask=seg_mask, - target_token_mask=seg_mod, - layer_cache=accumulated_cache, - kv_cache_mode="extend", - cache_write_slice=None, - ) - segment_outputs.append(seg_hs) - - prefix_hs = torch.cat(segment_outputs, dim=1) if segment_outputs else prefix_hs - - # Store the accumulated prefix cache into the main cache for this layer. - if layer_cache is not None and accumulated_cache.is_populated: - layer_cache.store(*accumulated_cache.get()) - - # Target: full attention over [cached prefix, target]. - target_cache = ( - layer_cache - if layer_cache is not None - else (accumulated_cache if accumulated_cache.is_populated else None) + cache_write_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None + block_image_ids, block_key_valid, block_prefix_len = image_ids, joint_key_valid, prefix_len + + for index_block, block in enumerate(self.transformer_blocks): + layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None + if torch.is_grad_enabled() and self.gradient_checkpointing: + joint_hidden_states = self._gradient_checkpointing_func( + block, + joint_hidden_states, + modulation, + rotary_emb, + attention_mask, + modulation_mask, + layer_cache, + kv_cache_mode, + cache_write_slice, + block_image_ids, + block_key_valid, + block_prefix_len, ) - target_hs = block( - hidden_states=target_hs, + else: + joint_hidden_states = block( + hidden_states=joint_hidden_states, modulation=modulation, - rotary_emb=target_rope, - attention_mask=None, - target_token_mask=target_mod_mask, - layer_cache=target_cache, - kv_cache_mode="cached" if target_cache is not None else None, - cache_write_slice=None, + rotary_emb=rotary_emb, + attention_mask=attention_mask, + target_token_mask=modulation_mask, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, + cache_write_slice=cache_write_slice, + image_ids=block_image_ids, + key_valid=block_key_valid, + prefix_len=block_prefix_len, ) - joint_hidden_states = torch.cat([prefix_hs, target_hs], dim=1) - else: - for index_block, block in enumerate(self.transformer_blocks): - layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None - if torch.is_grad_enabled() and self.gradient_checkpointing: - joint_hidden_states = self._gradient_checkpointing_func( - block, - joint_hidden_states, - modulation, - rotary_emb, - attention_mask, - modulation_mask, - layer_cache, - kv_cache_mode, - cache_write_slice, - ) - else: - joint_hidden_states = block( - hidden_states=joint_hidden_states, - modulation=modulation, - rotary_emb=rotary_emb, - attention_mask=attention_mask, - target_token_mask=modulation_mask, - layer_cache=layer_cache, - kv_cache_mode=kv_cache_mode, - cache_write_slice=cache_write_slice, - ) - joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) output = self.proj_out(joint_hidden_states) diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index ad22ddba23f0..0259f6648792 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -159,6 +159,36 @@ def test_kv_cache_requires_causal_condition(self): with pytest.raises(ValueError, match="causal_condition"): model(**self.get_dummy_inputs(), kv_cache=kv_cache, kv_cache_mode="extract") + @pytest.mark.parametrize("pad_prompt", [False, True]) + @pytest.mark.parametrize("batch_size", [1, 2]) + def test_sdpa_processor_matches_flex(self, pad_prompt, batch_size): + # both processors implement the same block-causal attention: prefill and decode must agree, with and + # without right-padded text + from diffusers.models.transformers.transformer_qwenimage21 import ( + QwenImage21FlexAttnProcessor, + QwenImage21KVCache, + QwenImage21SDPAAttnProcessor, + ) + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device).eval() + inputs = self.get_dummy_inputs(batch_size=batch_size) + if pad_prompt: + inputs["encoder_hidden_states_mask"][:, -1] = 0 + + outputs = {} + for processor_cls in (QwenImage21FlexAttnProcessor, QwenImage21SDPAAttnProcessor): + model.set_attn_processor(processor_cls()) + kv_cache = QwenImage21KVCache(init_dict["num_layers"]) + with torch.no_grad(): + prefill = model(**inputs, kv_cache=kv_cache, kv_cache_mode="extract", return_dict=False)[0] + decode = model(**inputs, kv_cache=kv_cache, kv_cache_mode="cached", return_dict=False)[0] + outputs[processor_cls] = (prefill[:, -decode.shape[1] :], decode) + + (flex_prefill, flex_decode), (sdpa_prefill, sdpa_decode) = outputs.values() + assert torch.allclose(flex_prefill, sdpa_prefill, atol=1e-5) + assert torch.allclose(flex_decode, sdpa_decode, atol=1e-5) + def test_non_flex_backend_rejected_when_causal(self): model = self.model_class(**self.get_init_dict()).to(torch_device).eval() model.set_attention_backend("native") From de43a12a8420a01cc020004f1573676af5e2ae12 Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 19:34:06 +0800 Subject: [PATCH 06/20] refactor: let the processor pick the block-causal path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds on the previous commit by @yiyixuxu, which moved the segmentation into the processor. - Drop the internal `torch.compile(flex_attention)`. The convention is that the caller compiles, so `QwenImage21FlexAttnProcessor` goes through `dispatch_attention_fn(..., backend="flex")` and warns once when it finds an uncompiled `flex_attention` — that falls back to a dense fp32 score matrix and runs out of memory at high resolution. - `_attention_backend` is `None` on the flex processor. It is what `set_attention_backend()` sets and only applies to the cached decode steps; the prefill needs the flex kernel for its `BlockMask` and is not configurable. The processor raises from `__init__` when flex_attention is unavailable. - Pad the sequence axis directly instead of transposing around `F.pad`, so the padded tensors stay contiguous, which the compiled flex kernel requires. - Rename `QwenImage21SDPAAttnProcessor` to `QwenImage21AttnProcessor` and make it the default. Once compiled, flex is 1.5% faster end to end at 2048x2048 (31.8s vs 32.3s over 20 steps) and 3.8% faster with two condition images, in exchange for 37s of compilation; uncompiled it cannot render 2048x2048 at all. A default that only works when the caller compiles is the wrong trade, so flex is documented as the opt-in path instead. - Derive the prefix segment boundaries once per forward rather than once per layer. They only depend on `image_ids` and `prefix_len`, so the per-layer version repeated the same `tolist()` device sync 32 times. `forward` passes down whichever representation the installed processors read, and builds neither for a processor that does not need it. - Replace `test_non_flex_backend_rejected_when_causal`, which no longer describes the intended behaviour, with a check that `set_attention_backend` only affects decode. Measured on one H100, bf16, batch 1, 20 steps at 2048x2048: the default path runs out of the box in 32.7s at 56.5 GiB peak, and `set_attn_processor(QwenImage21AttnProcessor())` now works at all — it used to hand the flex `BlockMask` to a non-flex kernel and raise. --- .../transformers/transformer_qwenimage21.py | 193 ++++++++++-------- .../test_models_transformer_qwenimage21.py | 33 ++- 2 files changed, 139 insertions(+), 87 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index dfdb0b1027a0..b5d89179dfdc 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -23,7 +23,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import logging from ...utils.peft_utils import apply_lora_scale -from ...utils.torch_utils import maybe_allow_in_graph +from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -41,31 +41,16 @@ # `create_block_mask` quantizes the mask to 128-token blocks. _FLEX_BLOCK_SIZE = 128 -# flex_attention is optional. When available we use a compiled flex_attention with a BlockMask for -# efficient single-pass block-causal attention. When unavailable, we fall back to an exact multi-pass -# SDPA prefill that processes each image block with bidirectional attention and text segments with -# causal attention, matching the block-causal mask exactly. +# flex_attention is optional: `QwenImage21FlexAttnProcessor` needs it, `QwenImage21AttnProcessor` does not. _FLEX_AVAILABLE = False -_compiled_flex_attention = None try: - from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention + import torch.nn.attention.flex_attention as flex_attention_module + from torch.nn.attention.flex_attention import BlockMask, create_block_mask _FLEX_AVAILABLE = True except ImportError: BlockMask = None - flex_attention = None - - -def _get_compiled_flex_attention(): - """Return a compiled flex_attention, compiling on first call. - - Compiling is required for the block-sparse kernel that avoids materializing the full Q@K^T matrix. Without it - flex_attention falls back to a dense fp32 math path that OOMs on long sequences. - """ - global _compiled_flex_attention - if _compiled_flex_attention is None: - _compiled_flex_attention = torch.compile(flex_attention) - return _compiled_flex_attention + flex_attention_module = None class QwenImage21KVLayerCache: @@ -333,6 +318,40 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): ) +@lru_cache_unless_export(maxsize=1) +def _warn_if_flex_attention_is_uncompiled(): + """Warn once per process when `flex_attention` has not been compiled. + + `dispatch_attention_fn` reaches `flex_attention` through its module, so a user who compiles it — directly or by + compiling the model — is picked up here. Uncompiled, flex_attention falls back to a dense fp32 score matrix, + which is far slower and runs out of memory at high resolution, so say so rather than let it happen quietly. + """ + if not hasattr(flex_attention_module.flex_attention, "_torchdynamo_orig_callable"): + logger.warning( + "`QwenImage21FlexAttnProcessor` is running an uncompiled `flex_attention`, which materializes the full " + "attention score matrix in fp32 and will run out of memory at high resolution. Compile the model with " + "`transformer.compile()`, or switch to `QwenImage21AttnProcessor`." + ) + + +def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]: + """Split the prefix into `(start, end, is_text)` runs of equal `image_ids`. + + This is the block-causal structure in the form [`QwenImage21AttnProcessor`] consumes it, the way + [`~build_qwenimage21_block_causal_mask`] is the form [`QwenImage21FlexAttnProcessor`] consumes. It only depends + on `image_ids` and `prefix_len`, so the model derives it once per forward rather than in every processor call — + `tolist()` is a device sync, and there is one processor call per layer. + """ + prefix_ids = image_ids[:prefix_len].tolist() + segments = [] + start = 0 + for index in range(1, prefix_len + 1): + if index == prefix_len or prefix_ids[index] != prefix_ids[start]: + segments.append((start, index, prefix_ids[start] < 0)) + start = index + return segments + + def _qwenimage21_prepare_qkv( attn: "QwenImage21Attention", hidden_states: torch.Tensor, @@ -377,17 +396,26 @@ def _qwenimage21_prepare_qkv( class QwenImage21FlexAttnProcessor: r""" - Attention processor for Qwen-Image 2.1 using compiled `flex_attention` with a `BlockMask` for exact block-causal - attention. This is the recommended path and is required for high resolutions (2048²+) where a dense score matrix - would OOM. + Attention processor for Qwen-Image 2.1 that runs the block-causal prefill as one `flex_attention` call driven by + a `BlockMask`, and the cached decode steps through the configured attention backend. - ``flex_attention`` is compiled on the first forward pass so the block-sparse kernel is used instead of the dense - fallback. The first call will be slower due to compilation. + Compile the model before using it, as the docs show. An uncompiled `flex_attention` falls back to a dense fp32 + score matrix, which is far slower and runs out of memory at high resolution. Use `QwenImage21AttnProcessor` when + you do not want to compile. """ - _attention_backend = "flex" + # Set by `set_attention_backend()` and only meaningful for the decode steps; the prefill needs the flex kernel + # for its `BlockMask` and is not configurable. + _attention_backend = None _parallel_config = None + def __init__(self): + if not _FLEX_AVAILABLE: + raise ImportError( + "`QwenImage21FlexAttnProcessor` requires `torch.nn.attention.flex_attention`, which needs " + "torch>=2.5. Use `QwenImage21AttnProcessor` instead." + ) + def __call__( self, attn: "QwenImage21Attention", @@ -397,9 +425,8 @@ def __call__( layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, - image_ids: torch.Tensor | None = None, + segments: list[tuple[int, int, bool]] | None = None, key_valid: torch.Tensor | None = None, - prefix_len: int | None = None, ) -> torch.Tensor: query, key, value, seq_len_q = _qwenimage21_prepare_qkv( attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice @@ -407,30 +434,38 @@ def __call__( seq_len_kv = key.shape[1] if isinstance(attention_mask, BlockMask): - # prefill: the BlockMask expresses the block-causal structure in one flex call + # prefill: the BlockMask expresses the block-causal structure in one flex call. Query and key are + # padded up to the mask's block-quantized length. + if not torch.compiler.is_compiling(): + _warn_if_flex_attention_is_uncompiled() pad_q = int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_q pad_kv = int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_kv + # Pad the sequence axis. `F.pad` counts from the last dimension, so the head and channel axes are + # padded by zero first. The result stays contiguous, which the compiled flex kernel requires. if pad_q: - query = F.pad(query.transpose(1, 3), (0, pad_q)).transpose(1, 3) + query = F.pad(query, (0, 0, 0, 0, 0, pad_q)) if pad_kv: - key = F.pad(key.transpose(1, 3), (0, pad_kv)).transpose(1, 3) - value = F.pad(value.transpose(1, 3), (0, pad_kv)).transpose(1, 3) - - hidden_states = _get_compiled_flex_attention()( - query.transpose(1, 2).contiguous(), - key.transpose(1, 2).contiguous(), - value.transpose(1, 2).contiguous(), - block_mask=attention_mask, - ).transpose(1, 2) + key = F.pad(key, (0, 0, 0, 0, 0, pad_kv)) + value = F.pad(value, (0, 0, 0, 0, 0, pad_kv)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + backend="flex", + parallel_config=self._parallel_config, + ) else: - # decode: full attention over [cached prefix, target] via SDPA + # decode: full attention over [cached prefix, target] hidden_states = dispatch_attention_fn( query, key, value, attn_mask=attention_mask, dropout_p=0.0, - backend=None, + backend=self._attention_backend, parallel_config=self._parallel_config, ) hidden_states = hidden_states[:, :seq_len_q] @@ -440,14 +475,13 @@ def __call__( return attn.to_out[1](hidden_states) -class QwenImage21SDPAAttnProcessor: +class QwenImage21AttnProcessor: r""" - SDPA attention processor for Qwen-Image 2.1. Use this when `flex_attention` is not available. + Attention processor for Qwen-Image 2.1 that needs neither `flex_attention` nor a compiled model. - The block-causal mask is implemented exactly via a multi-pass prefill orchestrated by the model's `forward`: each - image block in the prefix is processed with bidirectional attention within the block and full attention to all - preceding segments, and text segments get a causal mask. The target image attends fully to the cached prefix + - itself. This matches the block-causal mask without requiring `flex_attention`. + The prefill decomposes the block-causal mask into one attention call per prefix segment plus one for the target + image, which is exact but slower than [`QwenImage21FlexAttnProcessor`]. The segment boundaries are computed once + per forward by the model and passed in as `segments`. """ _attention_backend = None @@ -462,15 +496,14 @@ def __call__( layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, - image_ids: torch.Tensor | None = None, + segments: list[tuple[int, int, bool]] | None = None, key_valid: torch.Tensor | None = None, - prefix_len: int | None = None, ) -> torch.Tensor: query, key, value, seq_len_q = _qwenimage21_prepare_qkv( attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice ) - if image_ids is None: + if segments is None: # decode: full attention over [cached prefix, target] hidden_states = dispatch_attention_fn( query, @@ -478,22 +511,14 @@ def __call__( value, attn_mask=attention_mask, dropout_p=0.0, - backend=None, + backend=self._attention_backend, parallel_config=self._parallel_config, ) else: - # prefill: the block-causal mask decomposes into one attention call per prefix segment plus one for the - # target image. Every segment attends to the keys `[0, end)` (everything before it plus its own block); - # text segments additionally get a causal triangle over their own keys; padded text keys are dropped. - # `attention_mask` is the flex BlockMask of the same structure, meant for `QwenImage21FlexAttnProcessor`; - # it is not used here. - prefix_ids = image_ids[:prefix_len].tolist() - segments = [] - start = 0 - for i in range(1, prefix_len + 1): - if i == prefix_len or prefix_ids[i] != prefix_ids[start]: - segments.append((start, i, prefix_ids[start] < 0)) - start = i + # prefill: every segment attends to the keys `[0, end)` (everything before it plus its own block); text + # segments additionally get a causal triangle over their own keys; padded text keys are dropped. + # `attention_mask` may hold the flex `BlockMask` of the same structure, which is not used here. + prefix_len = segments[-1][1] if segments else 0 outputs = [] for start, end, is_text in segments: seg_mask = None @@ -545,8 +570,11 @@ class QwenImage21Attention(torch.nn.Module, AttentionModuleMixin): [`~models.attention_processor.Attention`] so Qwen-Image 2.x checkpoints load into it unchanged. """ - _default_processor_cls = QwenImage21FlexAttnProcessor if _FLEX_AVAILABLE else QwenImage21SDPAAttnProcessor - _available_processors = [QwenImage21FlexAttnProcessor, QwenImage21SDPAAttnProcessor] + # The default must not depend on the caller having compiled the model: an uncompiled `flex_attention` falls + # back to a dense fp32 score matrix and runs out of memory at high resolution. `QwenImage21FlexAttnProcessor` + # is the faster path once compiled, and the docs show how to opt into it. + _default_processor_cls = QwenImage21AttnProcessor + _available_processors = [QwenImage21AttnProcessor, QwenImage21FlexAttnProcessor] def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-6, processor: Any | None = None): super().__init__() @@ -610,9 +638,8 @@ def forward( layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, cache_write_slice: slice | None = None, - image_ids: torch.Tensor | None = None, + segments: list[tuple[int, int, bool]] | None = None, key_valid: torch.Tensor | None = None, - prefix_len: int | None = None, ) -> torch.Tensor: mod1, mod2 = modulation.chunk(2, dim=-1) @@ -624,9 +651,8 @@ def forward( layer_cache=layer_cache, kv_cache_mode=kv_cache_mode, cache_write_slice=cache_write_slice, - image_ids=image_ids, + segments=segments, key_valid=key_valid, - prefix_len=prefix_len, ) hidden_states = hidden_states + img_gate1.tanh() * attn_output @@ -937,16 +963,25 @@ def forward( modulation_mask = modulation_mask[prefix_len:] attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] cache_write_slice = None - block_image_ids, block_key_valid, block_prefix_len = None, None, None + block_segments, block_key_valid = None, None else: - # prefill: the whole joint sequence. The block-causal structure is passed down both as a flex - # `BlockMask` (used by `QwenImage21FlexAttnProcessor`) and as its ingredients (`image_ids`, - # `key_valid`, `prefix_len`, used by `QwenImage21SDPAAttnProcessor`); the processor picks. - attention_mask = build_qwenimage21_block_causal_mask( - image_ids, joint_key_valid, batch_size, hidden_states.device + # prefill: the whole joint sequence. The block-causal structure goes down in whichever form the + # installed processors read it — a flex `BlockMask`, per-segment boundaries, or both for a mixed set — + # so neither path pays for building the other's metadata. + processors = [block.attn.processor for block in self.transformer_blocks] + needs_block_mask = any(isinstance(processor, QwenImage21FlexAttnProcessor) for processor in processors) + attention_mask = ( + build_qwenimage21_block_causal_mask(image_ids, joint_key_valid, batch_size, hidden_states.device) + if needs_block_mask + else None + ) + block_segments = ( + None + if all(isinstance(processor, QwenImage21FlexAttnProcessor) for processor in processors) + else _qwenimage21_prefix_segments(image_ids, prefix_len) ) cache_write_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None - block_image_ids, block_key_valid, block_prefix_len = image_ids, joint_key_valid, prefix_len + block_key_valid = joint_key_valid for index_block, block in enumerate(self.transformer_blocks): layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None @@ -961,9 +996,8 @@ def forward( layer_cache, kv_cache_mode, cache_write_slice, - block_image_ids, + block_segments, block_key_valid, - block_prefix_len, ) else: joint_hidden_states = block( @@ -975,9 +1009,8 @@ def forward( layer_cache=layer_cache, kv_cache_mode=kv_cache_mode, cache_write_slice=cache_write_slice, - image_ids=block_image_ids, + segments=block_segments, key_valid=block_key_valid, - prefix_len=block_prefix_len, ) joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index 0259f6648792..e6d45c21779e 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -165,9 +165,9 @@ def test_sdpa_processor_matches_flex(self, pad_prompt, batch_size): # both processors implement the same block-causal attention: prefill and decode must agree, with and # without right-padded text from diffusers.models.transformers.transformer_qwenimage21 import ( + QwenImage21AttnProcessor, QwenImage21FlexAttnProcessor, QwenImage21KVCache, - QwenImage21SDPAAttnProcessor, ) init_dict = self.get_init_dict() @@ -177,7 +177,7 @@ def test_sdpa_processor_matches_flex(self, pad_prompt, batch_size): inputs["encoder_hidden_states_mask"][:, -1] = 0 outputs = {} - for processor_cls in (QwenImage21FlexAttnProcessor, QwenImage21SDPAAttnProcessor): + for processor_cls in (QwenImage21FlexAttnProcessor, QwenImage21AttnProcessor): model.set_attn_processor(processor_cls()) kv_cache = QwenImage21KVCache(init_dict["num_layers"]) with torch.no_grad(): @@ -189,11 +189,30 @@ def test_sdpa_processor_matches_flex(self, pad_prompt, batch_size): assert torch.allclose(flex_prefill, sdpa_prefill, atol=1e-5) assert torch.allclose(flex_decode, sdpa_decode, atol=1e-5) - def test_non_flex_backend_rejected_when_causal(self): - model = self.model_class(**self.get_init_dict()).to(torch_device).eval() - model.set_attention_backend("native") - with pytest.raises(ValueError, match="flex"): - model(**self.get_dummy_inputs()) + @pytest.mark.parametrize("processor_name", ["QwenImage21FlexAttnProcessor", "QwenImage21AttnProcessor"]) + def test_attention_backend_applies_to_decode_only(self, processor_name): + """ + `set_attention_backend` configures the cached decode steps. The prefill picks its own path from the installed + processor, so setting a non-flex backend must not change or reject it. + """ + import diffusers.models.transformers.transformer_qwenimage21 as transformer_module + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + init_dict = self.get_init_dict() + inputs = self.get_dummy_inputs() + outputs = [] + for backend in (None, "native"): + torch.manual_seed(0) + model = self.model_class(**init_dict).to(torch_device).eval() + model.set_attn_processor(getattr(transformer_module, processor_name)()) + if backend is not None: + model.set_attention_backend(backend) + kv_cache = QwenImage21KVCache(init_dict["num_layers"]) + with torch.no_grad(): + model(**inputs, kv_cache=kv_cache, kv_cache_mode="extract", return_dict=False) + outputs.append(model(**inputs, kv_cache=kv_cache, kv_cache_mode="cached", return_dict=False)[0]) + + torch.testing.assert_close(outputs[0], outputs[1]) class TestQwenImage21BlockCausalMask: From b3cb5d784ca700ef0d6fe9f39aed214372020ead Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 19:55:42 +0800 Subject: [PATCH 07/20] fix: VAE class defaults did not describe the released model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `scale_factor_spatial` was 8 while the encoder applies four spatial downsamples: `encode` takes a 1024x1024 image to a (1, 64, 1, 64, 64) latent, so the ratio is 16. Every tile-to-latent conversion divides by it, so tiling silently produced a wrong-shaped latent — a 2048x2048 encode came out as 168x168 instead of 128x128. The in/out channel defaults were 3 while this VAE takes four channels, so the class could not be instantiated from its own defaults. Add the model test file that was missing, covering the ratio against the architecture and the shape of a tiled encode. Tile values are not compared: each tile starts the causal convolution feature cache fresh, which is a property of the tiling implementation rather than of these defaults. --- .../autoencoder_kl_qwenimage21.py | 6 +- .../test_models_autoencoder_kl_qwenimage21.py | 110 ++++++++++++++++++ 2 files changed, 113 insertions(+), 3 deletions(-) create mode 100644 tests/models/autoencoders/test_models_autoencoder_kl_qwenimage21.py diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index 34525fae329a..49f401e07526 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -1124,11 +1124,11 @@ def __init__( 3.8161, ], is_residual: bool = True, - in_channels: int = 3, - out_channels: int = 3, + in_channels: int = 4, + out_channels: int = 4, patch_size: int | None = None, scale_factor_temporal: int | None = 8, - scale_factor_spatial: int | None = 8, + scale_factor_spatial: int | None = 16, ) -> None: super().__init__() diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_qwenimage21.py b/tests/models/autoencoders/test_models_autoencoder_kl_qwenimage21.py new file mode 100644 index 000000000000..a1db8eebb530 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_kl_qwenimage21.py @@ -0,0 +1,110 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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 torch + +from diffusers import AutoencoderKLQwenImage21 +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BaseModelTesterConfig, ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin + + +enable_full_determinism() + + +class AutoencoderKLQwenImage21TesterConfig(BaseModelTesterConfig): + # Four spatial downsampling stages, so the latents are 16x smaller per axis. The condition and target images + # carry an alpha channel, hence four input channels. + num_channels = 4 + spatial_compression_ratio = 16 + sizes = (96, 96) + + @property + def model_class(self): + return AutoencoderKLQwenImage21 + + @property + def output_shape(self): + return (self.num_channels, 1, *self.sizes) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self): + return { + "base_dim": 4, + "decoder_base_dim": 4, + "z_dim": 4, + "dim_mult": [1, 1, 1, 1, 1], + "num_res_blocks": 1, + "attn_scales": [], + "temperal_downsample": [False, True, True, True], + "latents_mean": [0.0] * 4, + "latents_std": [1.0] * 4, + } + + def get_dummy_inputs(self): + image = randn_tensor((1, self.num_channels, 1, *self.sizes), generator=self.generator, device=torch_device) + return {"sample": image} + + +class TestAutoencoderKLQwenImage21(AutoencoderKLQwenImage21TesterConfig, ModelTesterMixin): + base_precision = 1e-2 + + def test_spatial_compression_ratio_matches_architecture(self): + """ + `scale_factor_spatial` drives every tile-to-latent conversion, so it has to be the ratio the encoder + actually applies — one downsample per `dim_mult` stage after the first. + """ + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device).eval() + with torch.no_grad(): + latent = model.encode(self.get_dummy_inputs()["sample"]).latent_dist.mode() + + expected = 2 ** (len(init_dict["dim_mult"]) - 1) + assert model.spatial_compression_ratio == expected + assert latent.shape[-1] == self.sizes[-1] // expected + assert latent.shape[-2] == self.sizes[-2] // expected + + def test_tiled_encode_keeps_the_latent_shape(self): + """ + Every tile-to-latent conversion divides by `scale_factor_spatial`, so a wrong ratio silently changes the + shape of a tiled encode: 2048x2048 came out as a 168x168 latent instead of 128x128. The tile sizes are + lowered here so the tiled path actually runs on an input this small. + + Only the shape is asserted. Tile values differ from a single pass because the causal convolutions carry a + feature cache that each tile starts fresh, which is a property of the tiling implementation itself. + """ + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + sample = self.get_dummy_inputs()["sample"] + + with torch.no_grad(): + untiled = model.encode(sample).latent_dist.mode() + model.enable_tiling( + tile_sample_min_height=48, + tile_sample_min_width=48, + tile_sample_stride_height=32, + tile_sample_stride_width=32, + ) + tiled = model.encode(sample).latent_dist.mode() + + assert tiled.shape == untiled.shape + + +class TestAutoencoderKLQwenImage21SlicingTiling(AutoencoderKLQwenImage21TesterConfig, AutoencoderTesterMixin): + """Slicing and tiling tests for AutoencoderKLQwenImage21.""" From f374bef6db714bb2fe482e03c1a1826722229703 Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 20:11:06 +0800 Subject: [PATCH 08/20] address review feedback: copies, pipeline, docs VAE: - Mark the classes that are byte-identical to their Wan originals with `# Copied from` (`DupUp3D`, `WanUpsample`, `WanRMS_norm`, `WanAttentionBlock`, `patchify`, `unpatchify`). Adopting the upstream `RMS_norm.forward` in the process also picks up a fix we had missed: it normalizes in fp32 for fp16/bf16/fp8 inputs. - Drop the `non_linearity` argument, which was always "silu", from the blocks that take it, and validate `AvgDown3D`'s channel divisibility with a `ValueError` before the fields are assigned rather than with an `assert` after. Pipeline: - Take `calculate_dimensions` verbatim from the edit pipeline so it can carry a `# Copied from`, which is what fixes `check_repository_consistency`: the marker on `_encode_vae_image` was one blank line out of sync. - Move the `QwenImage21KVCache` import to the top of the module. Docs: - Apply @stevhliu's suggestions. `models.autoencoders.vae.AutoencoderKLOutput` does not exist, so that autodoc reference was broken; the module is `autoencoder_kl`. - Describe the two attention processors instead of the old "with and without the flex backend" split, and add the snippet for opting into flex, which has to be compiled. `make style` also reflowed a few docstrings from the previous commit. --- .../api/models/autoencoderkl_qwenimage21.md | 2 +- .../api/models/qwenimage21_transformer2d.md | 13 ++-- docs/source/en/api/pipelines/qwenimage21.md | 27 +++++--- .../autoencoder_kl_qwenimage21.py | 61 ++++++++++--------- .../transformers/transformer_qwenimage21.py | 12 ++-- .../qwenimage21/pipeline_qwenimage21.py | 15 +++-- 6 files changed, 77 insertions(+), 53 deletions(-) diff --git a/docs/source/en/api/models/autoencoderkl_qwenimage21.md b/docs/source/en/api/models/autoencoderkl_qwenimage21.md index eb81c7674470..fefb50b53c16 100644 --- a/docs/source/en/api/models/autoencoderkl_qwenimage21.md +++ b/docs/source/en/api/models/autoencoderkl_qwenimage21.md @@ -30,7 +30,7 @@ vae = AutoencoderKLQwenImage21.from_pretrained("Qwen/Qwen-Image-2.1", subfolder= ## AutoencoderKLOutput -[[autodoc]] models.autoencoders.vae.AutoencoderKLOutput +[[autodoc]] models.autoencoders.autoencoder_kl.AutoencoderKLOutput ## DecoderOutput diff --git a/docs/source/en/api/models/qwenimage21_transformer2d.md b/docs/source/en/api/models/qwenimage21_transformer2d.md index 278572747937..022f0710c173 100644 --- a/docs/source/en/api/models/qwenimage21_transformer2d.md +++ b/docs/source/en/api/models/qwenimage21_transformer2d.md @@ -17,14 +17,15 @@ The single-stream transformer used by Qwen-Image 2.1. Text and image latents sha Two behaviours distinguish 2.1 from earlier QwenImage transformers: - **Block-causal attention** — attention follows `(q_idx >= kv_idx) or same_image_block`, so the joint sequence is - causal while each image block stays internally bidirectional. The `flex` attention backend gives efficient - single-pass attention; without it the model uses an exact multi-pass SDPA prefill that processes each block - separately. Both paths produce the same results. + causal while each image block stays internally bidirectional. `QwenImage21AttnProcessor` implements it as one + attention call per prefix segment and is the default. `QwenImage21FlexAttnProcessor` implements it as a single + `flex_attention` call driven by a `BlockMask`, which is faster once the model is compiled. Both produce the same + results. - `causal_condition` — text and condition-image tokens are modulated from `t = 0` rather than the sampled timestep. - Their activations are therefore independent of the denoising step, which is what makes the keys and values of that - prefix cacheable across steps via the `kv_cache` argument. + Their activations are independent of the denoising step, so the keys and values of that prefix are cacheable + across steps via the `kv_cache` argument. -The model can be loaded with the following code snippet. +Load it with: ```python import torch diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md index a20da77ecb68..f8ef6e3f079d 100644 --- a/docs/source/en/api/pipelines/qwenimage21.md +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -13,18 +13,14 @@ specific language governing permissions and limitations under the License. --> Qwen-Image 2.1 encodes the prompt and any condition images together with a Qwen3-VL model, then denoises the target image with a single-stream block-causal transformer. See -[`QwenImage21Transformer2DModel`](../models/qwenimage21_transformer2d) for details on block-causal attention and -`causal_condition`. - -The `flex` attention backend (`torch.nn.attention.flex_attention`) gives efficient single-pass block-causal attention. -Without it, the model uses an exact multi-pass SDPA prefill that processes each image block with bidirectional -attention and text with causal attention, matching the block-causal mask exactly. Both paths produce the same results. +[`QwenImage21Transformer2DModel`](../models/qwenimage21_transformer2d) for block-causal attention, the attention +processors, and `causal_condition`. ```python import torch from diffusers import QwenImage21Pipeline -pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16).to("cuda") +pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16).to("cuda") # Text-to-image image = pipe("A capybara wearing a wizard hat, oil painting", num_inference_steps=40).images[0] @@ -35,6 +31,23 @@ edited = pipe("Move it to a snowy mountain top", image=image, num_inference_step edited.save("edit.png") ``` +## Faster attention with flex_attention + +The default `QwenImage21AttnProcessor` runs the block-causal prefill as one attention call per prefix segment. It +needs no compilation and works on any PyTorch build. `QwenImage21FlexAttnProcessor` expresses the same mask as a +single `flex_attention` call, which is faster once the model is compiled. + +> [!TIP] +> Compile the model when you switch to the flex processor. An uncompiled `flex_attention` materializes the full +> attention score matrix in fp32, which is much slower and runs out of memory at high resolution. + +```python +from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21FlexAttnProcessor + +pipe.transformer.set_attn_processor(QwenImage21FlexAttnProcessor()) +pipe.transformer.compile() +``` + ## QwenImage21Pipeline [[autodoc]] QwenImage21Pipeline diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index 49f401e07526..4ceb54f9800c 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -40,14 +40,19 @@ def __init__( factor_s=1, ): super().__init__() + factor = factor_t * factor_s * factor_s + if in_channels * factor % out_channels != 0: + raise ValueError( + f"`in_channels` ({in_channels}) times the downsampling factor ({factor}) must be divisible by " + f"`out_channels` ({out_channels})." + ) + self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s - self.factor = self.factor_t * self.factor_s * self.factor_s - - assert in_channels * self.factor % out_channels == 0 - self.group_size = in_channels * self.factor // out_channels + self.factor = factor + self.group_size = in_channels * factor // out_channels def forward(self, x: torch.Tensor) -> torch.Tensor: pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t @@ -84,6 +89,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.DupUp3D with DupUp3D->QwenImage21DupUp3D class QwenImage21DupUp3D(nn.Module): def __init__( self, @@ -173,6 +179,7 @@ def forward(self, x, cache_x=None): return x +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanRMS_norm with Wan->QwenImage21 class QwenImage21RMS_norm(nn.Module): r""" A custom RMS normalization layer. @@ -196,9 +203,17 @@ def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bi self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): - return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + t in str(x.dtype) for t in ("float4_", "float8_") + ) + normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=(1 if self.channel_first else -1)).to( + x.dtype + ) + + return normalized * self.scale * self.gamma + self.bias +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanUpsample with Wan->QwenImage21 class QwenImage21Upsample(nn.Upsample): r""" Perform upsampling while ensuring the output tensor has the same data type as the input. @@ -315,7 +330,6 @@ class QwenImage21ResidualBlock(nn.Module): in_dim (int): Number of input channels. out_dim (int): Number of output channels. dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. - non_linearity (str, optional): Type of non-linearity to use. Default is "silu". """ def __init__( @@ -323,12 +337,11 @@ def __init__( in_dim: int, out_dim: int, dropout: float = 0.0, - non_linearity: str = "silu", ) -> None: super().__init__() self.in_dim = in_dim self.out_dim = out_dim - self.nonlinearity = get_activation(non_linearity) + self.nonlinearity = get_activation("silu") # layers self.norm1 = QwenImage21RMS_norm(in_dim, images=False) @@ -383,6 +396,7 @@ def forward(self, x, feat_cache=None, feat_idx=None): return x + h +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanAttentionBlock with Wan->QwenImage21 class QwenImage21AttentionBlock(nn.Module): r""" Causal self-attention with a single head. @@ -435,19 +449,18 @@ class QwenImage21MidBlock(nn.Module): Args: dim (int): Number of input/output channels. dropout (float): Dropout rate. - non_linearity (str): Type of non-linearity to use. """ - def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): + def __init__(self, dim: int, dropout: float = 0.0, num_layers: int = 1): super().__init__() self.dim = dim # Create the components - resnets = [QwenImage21ResidualBlock(dim, dim, dropout, non_linearity)] + resnets = [QwenImage21ResidualBlock(dim, dim, dropout)] attentions = [] for _ in range(num_layers): attentions.append(QwenImage21AttentionBlock(dim)) - resnets.append(QwenImage21ResidualBlock(dim, dim, dropout, non_linearity)) + resnets.append(QwenImage21ResidualBlock(dim, dim, dropout)) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) @@ -519,7 +532,6 @@ class QwenImage21Encoder3d(nn.Module): attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_downsample (list of bool): Whether to downsample temporally in each block. dropout (float): Dropout rate for the dropout layers. - non_linearity (str): Type of non-linearity to use. """ def __init__( @@ -532,7 +544,6 @@ def __init__( attn_scales=[], temperal_downsample=[True, True, False], dropout=0.0, - non_linearity: str = "silu", is_residual: bool = False, ): super().__init__() @@ -542,7 +553,7 @@ def __init__( self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_downsample = temperal_downsample - self.nonlinearity = get_activation(non_linearity) + self.nonlinearity = get_activation("silu") # dimensions dims = [dim * u for u in [1] + dim_mult] @@ -580,7 +591,7 @@ def __init__( scale /= 2.0 # middle blocks - self.mid_block = QwenImage21MidBlock(out_dim, dropout, non_linearity, num_layers=1) + self.mid_block = QwenImage21MidBlock(out_dim, dropout, num_layers=1) # output blocks self.norm_out = QwenImage21RMS_norm(out_dim, images=False) @@ -642,7 +653,6 @@ class QwenImage21ResidualUpBlock(nn.Module): dropout (float): Dropout rate temperal_upsample (bool): Whether to upsample on temporal dimension up_flag (bool): Whether to upsample or not - non_linearity (str): Type of non-linearity to use """ def __init__( @@ -653,7 +663,6 @@ def __init__( dropout: float = 0.0, temperal_upsample: bool = False, up_flag: bool = False, - non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim @@ -673,7 +682,7 @@ def __init__( resnets = [] current_dim = in_dim for _ in range(num_res_blocks + 1): - resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout, non_linearity)) + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout)) current_dim = out_dim self.resnets = nn.ModuleList(resnets) @@ -731,7 +740,6 @@ class QwenImage21UpBlock(nn.Module): num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') - non_linearity (str): Type of non-linearity to use """ def __init__( @@ -741,7 +749,6 @@ def __init__( num_res_blocks: int, dropout: float = 0.0, upsample_mode: str | None = None, - non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim @@ -752,7 +759,7 @@ def __init__( # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): - resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout, non_linearity)) + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout)) current_dim = out_dim self.resnets = nn.ModuleList(resnets) @@ -804,7 +811,6 @@ class QwenImage21Decoder3d(nn.Module): attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_upsample (list of bool): Whether to upsample temporally in each block. dropout (float): Dropout rate for the dropout layers. - non_linearity (str): Type of non-linearity to use. """ def __init__( @@ -816,7 +822,6 @@ def __init__( attn_scales=[], temperal_upsample=[False, True, True], dropout=0.0, - non_linearity: str = "silu", out_channels: int = 3, is_residual: bool = False, ): @@ -828,7 +833,7 @@ def __init__( self.attn_scales = attn_scales self.temperal_upsample = temperal_upsample - self.nonlinearity = get_activation(non_linearity) + self.nonlinearity = get_activation("silu") # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] @@ -837,7 +842,7 @@ def __init__( self.conv_in = QwenImage21CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks - self.mid_block = QwenImage21MidBlock(dims[0], dropout, non_linearity, num_layers=1) + self.mid_block = QwenImage21MidBlock(dims[0], dropout, num_layers=1) # upsample blocks self.up_blocks = nn.ModuleList([]) @@ -863,7 +868,6 @@ def __init__( dropout=dropout, temperal_upsample=temperal_upsample[i] if up_flag else False, up_flag=up_flag, - non_linearity=non_linearity, ) else: up_block = QwenImage21UpBlock( @@ -872,7 +876,6 @@ def __init__( num_res_blocks=num_res_blocks, dropout=dropout, upsample_mode=upsample_mode, - non_linearity=non_linearity, ) self.up_blocks.append(up_block) @@ -922,6 +925,7 @@ def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): return x +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.patchify with patchify->_patchify def _patchify(x, patch_size): if patch_size == 1: return x @@ -945,6 +949,7 @@ def _patchify(x, patch_size): return x +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.unpatchify with unpatchify->_unpatchify def _unpatchify(x, patch_size): if patch_size == 1: return x diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index b5d89179dfdc..4ef344fe5a42 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -323,8 +323,8 @@ def _warn_if_flex_attention_is_uncompiled(): """Warn once per process when `flex_attention` has not been compiled. `dispatch_attention_fn` reaches `flex_attention` through its module, so a user who compiles it — directly or by - compiling the model — is picked up here. Uncompiled, flex_attention falls back to a dense fp32 score matrix, - which is far slower and runs out of memory at high resolution, so say so rather than let it happen quietly. + compiling the model — is picked up here. Uncompiled, flex_attention falls back to a dense fp32 score matrix, which + is far slower and runs out of memory at high resolution, so say so rather than let it happen quietly. """ if not hasattr(flex_attention_module.flex_attention, "_torchdynamo_orig_callable"): logger.warning( @@ -338,8 +338,8 @@ def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> li """Split the prefix into `(start, end, is_text)` runs of equal `image_ids`. This is the block-causal structure in the form [`QwenImage21AttnProcessor`] consumes it, the way - [`~build_qwenimage21_block_causal_mask`] is the form [`QwenImage21FlexAttnProcessor`] consumes. It only depends - on `image_ids` and `prefix_len`, so the model derives it once per forward rather than in every processor call — + [`~build_qwenimage21_block_causal_mask`] is the form [`QwenImage21FlexAttnProcessor`] consumes. It only depends on + `image_ids` and `prefix_len`, so the model derives it once per forward rather than in every processor call — `tolist()` is a device sync, and there is one processor call per layer. """ prefix_ids = image_ids[:prefix_len].tolist() @@ -396,8 +396,8 @@ def _qwenimage21_prepare_qkv( class QwenImage21FlexAttnProcessor: r""" - Attention processor for Qwen-Image 2.1 that runs the block-causal prefill as one `flex_attention` call driven by - a `BlockMask`, and the cached decode steps through the configured attention backend. + Attention processor for Qwen-Image 2.1 that runs the block-causal prefill as one `flex_attention` call driven by a + `BlockMask`, and the cached decode steps through the configured attention backend. Compile the model before using it, as the docs show. An uncompiled `flex_attention` falls back to a dense fp32 score matrix, which is far slower and runs out of memory at high resolution. Use `QwenImage21AttnProcessor` when diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index a05a6e2b55b0..74cd092f928c 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -25,6 +25,7 @@ from ...image_processor import PipelineImageInput, VaeImageProcessor from ...loaders import QwenImageLoraLoaderMixin from ...models import AutoencoderKLQwenImage21, QwenImage21Transformer2DModel +from ...models.transformers.transformer_qwenimage21 import QwenImage21KVCache from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor @@ -145,10 +146,15 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") +# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.calculate_dimensions def calculate_dimensions(target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio - return round(width / 32) * 32, round(height / 32) * 32 + + width = round(width / 32) * 32 + height = round(height / 32) * 32 + + return width, height, None class QwenImage21Pipeline(DiffusionPipeline, QwenImageLoraLoaderMixin): @@ -442,6 +448,7 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): .to(image_latents.device, image_latents.dtype) ) image_latents = (image_latents - latents_mean) / latents_std + return image_latents def prepare_latents( @@ -584,7 +591,7 @@ def __call__( """ if image is not None: image_size = image[-1].size if isinstance(image, list) else image.size - calculated_width, calculated_height = calculate_dimensions( + calculated_width, calculated_height, _ = calculate_dimensions( output_resolution * output_resolution, image_size[0] / image_size[1] ) height = height or calculated_height @@ -620,7 +627,7 @@ def __call__( if hasattr(img, "mode") and img.mode != "RGBA": img = img.convert("RGBA") image_width, image_height = img.size - input_width, input_height = calculate_dimensions( + input_width, input_height, _ = calculate_dimensions( output_resolution * output_resolution, image_width / image_height ) input_image_sizes.append((input_width, input_height)) @@ -711,8 +718,6 @@ def __call__( # Text and condition-image keys and values are step-independent under `causal_condition`, so the first step # prefills them and later steps only recompute the target image's tokens. - from ...models.transformers.transformer_qwenimage21 import QwenImage21KVCache - num_blocks = len(self.transformer.transformer_blocks) cache_enabled = use_kv_cache and self.transformer.config.causal_condition cond_cache = QwenImage21KVCache(num_blocks) if cache_enabled else None From c245e42d101bf27bb8ba5f6c313b43a329c74aae Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 20:37:48 +0800 Subject: [PATCH 09/20] address review feedback: validate before doing work The `kv_cache` checks in the transformer's `forward` ran after the input projections, the joint sequence build, the rotary embeddings and the modulation, so a bad `kv_cache_mode` was only reported once that work had been done. They now run first. Following the same point through the pipeline turned up a check that could never fire: `check_inputs` warns when `height` and `width` are not divisible by `vae_scale_factor * 2`, but the rounding happened before the call, so the values it saw were always divisible. It now runs before the rounding, and `output_resolution=1000` warns and yields 992x992. --- .../transformers/transformer_qwenimage21.py | 20 +++++++++---------- .../qwenimage21/pipeline_qwenimage21.py | 4 ++-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 4ef344fe5a42..37f29bbfbfbd 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -897,6 +897,16 @@ def forward( """ batch_size = hidden_states.shape[0] + if kv_cache is not None and not self.config.causal_condition: + raise ValueError( + "kv_cache requires `causal_condition=True`. The cache is only valid because text and condition-image " + "tokens modulate from t=0, which makes their activations independent of the denoising step." + ) + if kv_cache is not None and kv_cache_mode not in ("extract", "cached"): + raise ValueError( + f"kv_cache_mode must be 'extract' or 'cached' when kv_cache is provided, got {kv_cache_mode!r}." + ) + hidden_states = self.img_in(hidden_states) encoder_hidden_states = self.txt_in(encoder_hidden_states) @@ -930,16 +940,6 @@ def forward( temb = self.time_text_embed(timestep, hidden_states) modulation = self.modulation(temb) - if kv_cache is not None and not self.config.causal_condition: - raise ValueError( - "kv_cache requires `causal_condition=True`. The cache is only valid because text and condition-image " - "tokens modulate from t=0, which makes their activations independent of the denoising step." - ) - if kv_cache is not None and kv_cache_mode not in ("extract", "cached"): - raise ValueError( - f"kv_cache_mode must be 'extract' or 'cached' when kv_cache is provided, got {kv_cache_mode!r}." - ) - # Right-padded prompt positions must never be attended to, on any path. Text positions of the joint sequence # line up, in order, with the non-image positions of the vision-language sequence — the two are interleaved, # so the mask cannot be sliced off as a prefix. diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 74cd092f928c..8e41ff298e4e 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -599,12 +599,12 @@ def __call__( height = height or output_resolution width = width or output_resolution + self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs) + multiple_of = self.vae_scale_factor * 2 width = width // multiple_of * multiple_of height = height // multiple_of * multiple_of - self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs) - self._attention_kwargs = attention_kwargs or {} self._current_timestep = None self._interrupt = False From 711fe5e003e23d69bb64922b76cb8f41af467fc2 Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 21:01:11 +0800 Subject: [PATCH 10/20] fix: num_images_per_prompt > 1 raised in the pipeline `encode_prompt` expands `prompt_embeds` and its mask to `batch_size * num_images_per_prompt`, but returns `image_pad_mask` unexpanded, and the target slots appended to that mask were sized from `latents`, which is expanded. Any call with `num_images_per_prompt > 1` therefore died in the concatenation: RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 1 but got size 2 for tensor number 1 in the list. Size the slots from each mask's own batch instead. The transformer reads the layout from row 0 because samples share it, so the mask does not need expanding. Verified end to end: `num_images_per_prompt=2`, a list of two prompts, both together, and each of those with classifier-free guidance and with a condition image. --- .../pipelines/qwenimage21/pipeline_qwenimage21.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 8e41ff298e4e..9501776eabd6 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -709,12 +709,14 @@ def __call__( self._num_timesteps = len(timesteps) # The transformer's `img_mask` spans the joint sequence, so append one slot per 2x2 group of target latents. - target_slots = torch.ones( - [latents.shape[0], latents.shape[1] // 4], dtype=image_pad_mask.dtype, device=image_pad_mask.device - ) - image_pad_mask = torch.cat([image_pad_mask, target_slots], dim=1) + # The slots follow each mask's own batch size: `latents` is already expanded by `num_images_per_prompt` + # while the masks are not, and the transformer reads the layout from row 0 because samples share it. + def append_target_slots(mask): + return torch.cat([mask, mask.new_ones(mask.shape[0], latents.shape[1] // 4)], dim=1) + + image_pad_mask = append_target_slots(image_pad_mask) if do_true_cfg: - negative_image_pad_mask = torch.cat([negative_image_pad_mask, target_slots], dim=1) + negative_image_pad_mask = append_target_slots(negative_image_pad_mask) # Text and condition-image keys and values are step-independent under `causal_condition`, so the first step # prefills them and later steps only recompute the target image's tokens. From e18afdc1ea475c681084a6f6ec890f00c79c83c3 Mon Sep 17 00:00:00 2001 From: naykun Date: Wed, 16 Sep 2026 21:01:41 +0800 Subject: [PATCH 11/20] Use the recommended sampling defaults: 40 steps, no guidance Qwen-Image 2.1 is meant to be sampled in 40 steps without classifier-free guidance, so `num_inference_steps` defaults to 40 and `true_cfg_scale` to 1.0. The other QwenImage pipelines default to 50 and 4.0, which is why this differs from its siblings. It also removes a warning from every default call: `true_cfg_scale=4.0` with no negative prompt took the "guidance is not enabled" branch. Passing a `negative_prompt` without raising `true_cfg_scale` still warns, which is the case worth warning about. The docs and the example docstring rely on the defaults now instead of passing a step count, and the docs state the recommendation. --- docs/source/en/api/pipelines/qwenimage21.md | 7 +++++-- .../pipelines/qwenimage21/pipeline_qwenimage21.py | 13 +++++++------ 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md index f8ef6e3f079d..35d2d2ef4b48 100644 --- a/docs/source/en/api/pipelines/qwenimage21.md +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -16,6 +16,9 @@ image with a single-stream block-causal transformer. See [`QwenImage21Transformer2DModel`](../models/qwenimage21_transformer2d) for block-causal attention, the attention processors, and `causal_condition`. +The defaults are the values Qwen recommends: 40 steps and no guidance. Pass a `negative_prompt` together with +`true_cfg_scale > 1` to turn classifier-free guidance on, which doubles the work per step. + ```python import torch from diffusers import QwenImage21Pipeline @@ -23,11 +26,11 @@ from diffusers import QwenImage21Pipeline pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16).to("cuda") # Text-to-image -image = pipe("A capybara wearing a wizard hat, oil painting", num_inference_steps=40).images[0] +image = pipe("A capybara wearing a wizard hat, oil painting").images[0] image.save("t2i.png") # Image-conditioned editing -edited = pipe("Move it to a snowy mountain top", image=image, num_inference_steps=40).images[0] +edited = pipe("Move it to a snowy mountain top", image=image).images[0] edited.save("edit.png") ``` diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 9501776eabd6..01be5278e180 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -49,10 +49,10 @@ >>> import torch >>> from diffusers import QwenImage21Pipeline - >>> pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16) + >>> pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16) >>> pipe.to("cuda") >>> prompt = "A capybara wearing a wizard hat, reading a book by candlelight, oil painting" - >>> image = pipe(prompt, num_inference_steps=50).images[0] + >>> image = pipe(prompt).images[0] >>> image.save("qwenimage21.png") ``` """ @@ -515,10 +515,10 @@ def __call__( prompt: str | list[str] = None, image: PipelineImageInput | None = None, negative_prompt: str | list[str] = None, - true_cfg_scale: float = 4.0, + true_cfg_scale: float = 1.0, height: int | None = None, width: int | None = None, - num_inference_steps: int = 50, + num_inference_steps: int = 40, sigmas: list[float] | None = None, num_images_per_prompt: int = 1, generator: torch.Generator | list[torch.Generator] | None = None, @@ -546,13 +546,14 @@ def __call__( into latent tokens prepended to the noise. negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide image generation. Ignored when `true_cfg_scale` is not greater than 1. - true_cfg_scale (`float`, *optional*, defaults to 4.0): + true_cfg_scale (`float`, *optional*, defaults to 1.0): Classifier-free guidance scale. Enabled by `true_cfg_scale > 1` together with a negative prompt. + Qwen-Image 2.1 is meant to be sampled without guidance, hence the default of 1.0. height (`int`, *optional*): Height in pixels of the generated image. Derived from the condition image's aspect ratio if omitted. width (`int`, *optional*): Width in pixels of the generated image. Derived from the condition image's aspect ratio if omitted. - num_inference_steps (`int`, *optional*, defaults to 50): + num_inference_steps (`int`, *optional*, defaults to 40): Number of denoising steps. sigmas (`list[float]`, *optional*): Custom sigmas for the denoising schedule. From d6dfd676aefd4fcf694eb180f0525d9ccec54fc8 Mon Sep 17 00:00:00 2001 From: naykun Date: Thu, 17 Sep 2026 15:45:43 +0800 Subject: [PATCH 12/20] docs: complete the forward and __call__ docstrings `utils/check_forward_call_docstrings.py` on main checks that every argument in a forward/__call__ signature has a docstring entry and that a non-None return type has a Returns section. It landed after this branch's base, so it only started running here once the copy check stopped failing ahead of it. Add the missing entries: `sample_posterior` and `generator` plus Returns on the VAE's forward, `attention_kwargs` and `return_dict` plus Returns on the transformer's, and the four embedding arguments and `callback_on_step_end_tensor_inputs` on the pipeline's __call__. --- .../autoencoders/autoencoder_kl_qwenimage21.py | 8 ++++++++ .../models/transformers/transformer_qwenimage21.py | 10 ++++++++++ .../pipelines/qwenimage21/pipeline_qwenimage21.py | 12 ++++++++++++ 3 files changed, 30 insertions(+) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index 4ceb54f9800c..4b93b39e0085 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -1525,8 +1525,16 @@ def forward( """ Args: sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior instead of taking its mode. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + Generator used when `sample_posterior` is `True`. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + [`~models.autoencoders.vae.DecoderOutput`] if `return_dict` is True, otherwise a plain `tuple`. """ x = sample posterior = self.encode(x).latent_dist diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 37f29bbfbfbd..1f42338127ef 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -894,6 +894,16 @@ def forward( kv_cache_mode (`str`, *optional*): `"extract"` to prefill the cache (first denoising step), `"cached"` to decode from it (later steps). Requires `causal_condition=True`. + attention_kwargs (`dict`, *optional*): + Forwarded to the attention processors, and carries `scale` for the LoRA layers. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: + [`~models.modeling_outputs.Transformer2DModelOutput`] if `return_dict` is True, otherwise a plain + `tuple` whose first element is the denoised latents. """ batch_size = hidden_states.shape[0] diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 01be5278e180..6e11b1e0d4a9 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -563,6 +563,15 @@ def __call__( Generator(s) to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings, which skip prompt encoding. Pass `prompt_embeds_mask` with them. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Bool mask marking the valid positions of `prompt_embeds`. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings, used in place of `negative_prompt`. Pass + `negative_prompt_embeds_mask` with them. + negative_prompt_embeds_mask (`torch.Tensor`, *optional*): + Bool mask marking the valid positions of `negative_prompt_embeds`. output_type (`str`, *optional*, defaults to `"pil"`): Output format, `"pil"`, `"np"`, `"pt"` or `"latent"`. return_dict (`bool`, *optional*, defaults to `True`): @@ -571,6 +580,9 @@ def __call__( Passed through to the attention processor. callback_on_step_end (`Callable`, *optional*): Called at the end of each denoising step. + callback_on_step_end_tensor_inputs (`list[str]`, *optional*, defaults to `["latents"]`): + Tensors from the denoising loop to hand to `callback_on_step_end`. They must be listed in the + pipeline's `_callback_tensor_inputs`. output_resolution (`int`, *optional*, defaults to 1024): Target side length used to derive `height`/`width` and to resize condition images. use_kv_cache (`bool`, *optional*, defaults to `True`): From b027736e6ccb172abd76c7b4c51ed180a503673f Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 17 Sep 2026 15:43:02 +0300 Subject: [PATCH 13/20] add pipeline tests for qwenimage 2.1 (#6) --- .../test_models_transformer_qwenimage21.py | 4 +- tests/pipelines/qwenimage21/__init__.py | 0 .../pipelines/qwenimage21/test_qwenimage21.py | 223 ++++++++++++++++++ 3 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 tests/pipelines/qwenimage21/__init__.py create mode 100644 tests/pipelines/qwenimage21/test_qwenimage21.py diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index e6d45c21779e..cb23fc4ad1d2 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -287,7 +287,9 @@ class TestQwenImage21TransformerMemory(QwenImage21TransformerTesterConfig, Memor class TestQwenImage21TransformerTraining(QwenImage21TransformerTesterConfig, TrainingTesterMixin): - pass + def test_gradient_checkpointing_is_applied(self): + expected_set = {"QwenImage21Transformer2DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, AttentionTesterMixin): diff --git a/tests/pipelines/qwenimage21/__init__.py b/tests/pipelines/qwenimage21/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/qwenimage21/test_qwenimage21.py b/tests/pipelines/qwenimage21/test_qwenimage21.py new file mode 100644 index 000000000000..1fd501950d9c --- /dev/null +++ b/tests/pipelines/qwenimage21/test_qwenimage21.py @@ -0,0 +1,223 @@ +# Copyright 2026 The HuggingFace Team. +# +# 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 numpy as np +import pytest +import torch +from PIL import Image +from transformers import ( + AutoTokenizer, + Qwen2VLImageProcessor, + Qwen3VLConfig, + Qwen3VLForConditionalGeneration, + Qwen3VLProcessor, + Qwen3VLVideoProcessor, +) + +from diffusers import ( + AutoencoderKLQwenImage21, + FlowMatchEulerDiscreteScheduler, + QwenImage21Pipeline, + QwenImage21Transformer2DModel, +) + +from ...testing_utils import assert_tensors_close +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, +) + + +# The pipeline hardcodes `vae_scale_factor = 16` and rounds height/width down to a multiple of 32, so 32 is the +# smallest resolution that survives: a 2x2 latent, which is exactly one vision slot's worth of target tokens. +IMAGE_SIZE = 32 + +# `QwenImage21Pipeline.__init__` reads the chat template and the `<|image_pad|>` id off `processor` eagerly, so the +# pipeline cannot be constructed with `processor=None`. `test_encode_prompt_works_in_isolation` builds exactly that +# — a denoiser-only pipeline with the text stack removed — to check that the `encode_prompt` outputs it was handed +# are enough to finish a call. Deferring that derivation is a `src/` change and out of scope for adding tests, so +# the test is marked `xfail`: whoever makes it lazy will see it XPASS and can drop this marker. +PROCESSOR_REQUIRED_AT_INIT = pytest.mark.xfail( + reason="`QwenImage21Pipeline.__init__` derives the system-token count from `processor`, so it raises on " + "`processor=None` and a pipeline without the text stack cannot be built.", + strict=True, +) + + +class QwenImage21PipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = QwenImage21Pipeline + required_input_params_in_call_signature = frozenset( + ["prompt", "image", "negative_prompt", "true_cfg_scale", "height", "width", "prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + # The VAE reconstructs RGBA, so the generated image carries four channels rather than three. + output_shape = (4, IMAGE_SIZE, IMAGE_SIZE) + # `encode_prompt` builds the chat template and tokenizes through `processor`, which the default + # ("text", "tokenizer") filter would drop from the text-encoder-only pipeline. + text_stack_component_names = ("text", "tokenizer", "processor") + + def get_dummy_components(self, num_layers: int = 2): + # The transformer consumes VAE latents directly, so `in_channels` has to be the VAE's `z_dim`. Keep `z_dim` + # away from 4: `prepare_latents` treats a condition image whose channel count already equals + # `latent_channels` as pre-encoded latents, and an RGBA image has 4 channels. + z_dim = 8 + text_hidden_size = 16 + + torch.manual_seed(0) + transformer = QwenImage21Transformer2DModel( + patch_size=1, + in_channels=z_dim, + out_channels=z_dim, + num_layers=num_layers, + # flex_attention needs a head dim of at least 16, and `axes_dims_rope` must sum to it. + attention_head_dim=16, + num_attention_heads=2, + context_in_dim=text_hidden_size, + mlp_ratio=2, + axes_dims_rope=(4, 6, 6), + ) + + torch.manual_seed(0) + # Five `dim_mult` stages means four spatial downsamples, i.e. the 16x compression the pipeline assumes. + vae = AutoencoderKLQwenImage21( + base_dim=4, + decoder_base_dim=4, + z_dim=z_dim, + dim_mult=[1, 1, 1, 1, 1], + num_res_blocks=1, + attn_scales=[], + temperal_downsample=[False, True, True, True], + latents_mean=[0.0] * z_dim, + latents_std=[1.0] * z_dim, + ) + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler() + + torch.manual_seed(0) + config = Qwen3VLConfig( + text_config={ + "hidden_size": text_hidden_size, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "head_dim": 8, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 1000000.0, + "mrope_section": [1, 1, 2], + }, + }, + vision_config={ + "depth": 2, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": text_hidden_size, + # One merged vision token has to cover 32 pixels, i.e. the 2x2 group of 16x-compressed latents that + # the transformer expands each vision slot into. Shrinking these would hand the transformer more + # slots than there are latent tokens. + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "num_position_embeddings": 64, + # Defaults to (8, 16, 24), which is out of range for a 2-layer vision tower. + "deepstack_visual_indexes": [0], + }, + ) + text_encoder = Qwen3VLForConditionalGeneration(config).eval() + + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + processor = Qwen3VLProcessor( + image_processor=Qwen2VLImageProcessor( + patch_size=16, merge_size=2, temporal_patch_size=2, min_pixels=32 * 32, max_pixels=64 * 64 + ), + tokenizer=tokenizer, + video_processor=Qwen3VLVideoProcessor(patch_size=16, merge_size=2, temporal_patch_size=2), + # The pipeline derives how many system tokens to drop by running this template, and matches it against + # the `<|im_start|>system ...` prefix it formats by hand, so the two have to agree. + chat_template=tokenizer.chat_template, + ) + + return { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "processor": processor, + } + + def get_dummy_condition_image(self): + array = np.random.RandomState(0).randint(0, 255, (IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8) + return Image.fromarray(array).convert("RGBA") + + def get_dummy_inputs(self): + return { + "prompt": "dance monkey", + "negative_prompt": "bad quality", + "generator": self.get_generator(0), + "num_inference_steps": 2, + "true_cfg_scale": 1.0, + "height": IMAGE_SIZE, + "width": IMAGE_SIZE, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", + } + + +class TestQwenImage21Pipeline(QwenImage21PipelineTesterConfig, PipelineTesterMixin): + @PROCESSOR_REQUIRED_AT_INIT + def test_encode_prompt_works_in_isolation(self): + super().test_encode_prompt_works_in_isolation() + + def test_inference_batch_single_identical(self): + # The shared test batches prompts of different lengths, so the short ones are padded. `QwenImage21Rope` + # walks the joint sequence and lets every non-image token advance the shared frame position, padding + # included, which shifts where the image block lands and moves the output by ~1e-3. Batching prompts of + # equal length reproduces a single call to within float noise (~3e-7), so this is the padding, not the + # batching. + super().test_inference_batch_single_identical(expected_max_diff=2e-3) + + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() + + image = pipe(**self.get_dummy_inputs()).images + generated_image = image[0] + assert generated_image.shape == self.output_shape + + # fmt: off + expected_slice = torch.tensor([0.5222, 0.6035, 0.6467, 0.6340, 0.6303, 0.6155, 0.6152, 0.6231, 0.4488, 0.4447, 0.4433, 0.4788, 0.4240, 0.4476, 0.4906, 0.3529]) + # fmt: on + + generated_slice = generated_image.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert_tensors_close(generated_slice, expected_slice, atol=5e-3) + + def test_inference_with_condition_image(self): + pipe = self.get_pipeline() + + inputs = self.get_dummy_inputs() + inputs["image"] = self.get_dummy_condition_image() + inputs["output_resolution"] = IMAGE_SIZE + + image = pipe(**inputs).images + assert image[0].shape == self.output_shape + + +class TestQwenImage21PipelineMemory(QwenImage21PipelineTesterConfig, MemoryTesterMixin): + pass From 3e52c4f59c64097b787985a95adc3c1c0c137fc8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 17 Sep 2026 15:46:14 +0300 Subject: [PATCH 14/20] add more copied froms (#5) --- .../models/autoencoders/autoencoder_kl_qwenimage21.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index 4b93b39e0085..412d64e083c1 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -31,6 +31,7 @@ CACHE_T = 2 +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AvgDown3D with AvgDown3D->QwenImage21AvgDown3D class QwenImage21AvgDown3D(nn.Module): def __init__( self, @@ -1199,6 +1200,7 @@ def __init__( else 0, } + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.enable_tiling def enable_tiling( self, tile_sample_min_height: int | None = None, @@ -1229,6 +1231,7 @@ def enable_tiling( self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.clear_cache def clear_cache(self): # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call self._conv_num = self._cached_conv_counts["decoder"] @@ -1239,6 +1242,7 @@ def clear_cache(self): self._enc_conv_idx = [0] self._enc_feat_map = [None] * self._enc_conv_num + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan._encode with patchify->_patchify def _encode(self, x: torch.Tensor): _, _, num_frame, height, width = x.shape @@ -1267,6 +1271,7 @@ def _encode(self, x: torch.Tensor): return enc @apply_forward_hook + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.encode def encode( self, x: torch.Tensor, return_dict: bool = True ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: @@ -1293,6 +1298,7 @@ def encode( return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan._decode with unpatchify->_unpatchify def _decode(self, z: torch.Tensor, return_dict: bool = True): _, _, num_frame, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio @@ -1325,6 +1331,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True): return DecoderOutput(sample=out) @apply_forward_hook + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.decode def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -1349,6 +1356,7 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t return (decoded,) return DecoderOutput(sample=decoded) + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.blend_v def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) for y in range(blend_extent): @@ -1357,6 +1365,7 @@ def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch. ) return b + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.blend_h def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) for x in range(blend_extent): @@ -1365,6 +1374,7 @@ def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch. ) return b + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_encode def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: r"""Encode a batch of images using a tiled encoder. @@ -1437,6 +1447,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_decode with unpatchify->_unpatchify def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. From 9827e152b35265914bb8a49d20a507d0cdbd7e6c Mon Sep 17 00:00:00 2001 From: naykun Date: Thu, 17 Sep 2026 20:42:03 +0800 Subject: [PATCH 15/20] docs: describe multiple condition images, and inline the flex warning The pipeline page now has a section on passing several condition images, which is where @sayakpaul asked for it, and the flex section emphasises that the processor wants a compiled model. `_warn_if_flex_attention_is_uncompiled()` is inlined at its only call site, as requested. A class-level flag keeps it to one warning per process: the default processor is constructed per attention module, so 32 instances would otherwise each warn, and the logger has no `warning_once`. --- docs/source/en/api/pipelines/qwenimage21.md | 13 ++++++- .../transformers/transformer_qwenimage21.py | 34 ++++++++----------- 2 files changed, 27 insertions(+), 20 deletions(-) diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md index 35d2d2ef4b48..64461ed8c0d0 100644 --- a/docs/source/en/api/pipelines/qwenimage21.md +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -34,11 +34,22 @@ edited = pipe("Move it to a snowy mountain top", image=image).images[0] edited.save("edit.png") ``` +## Multiple condition images + +Pass a list to `image` and every entry becomes its own block in the joint sequence: the Qwen3-VL encoder sees them as +vision context and the VAE contributes their latent tokens. Block-causal attention keeps each block internally +bidirectional while letting later blocks and the target image attend to the earlier ones, so the order you pass them +in is the order the model reads them. + +```python +edited = pipe("Put the flowers from the first image into the second scene", image=[flowers, scene]).images[0] +``` + ## Faster attention with flex_attention The default `QwenImage21AttnProcessor` runs the block-causal prefill as one attention call per prefix segment. It needs no compilation and works on any PyTorch build. `QwenImage21FlexAttnProcessor` expresses the same mask as a -single `flex_attention` call, which is faster once the model is compiled. +single `flex_attention` call, which is faster once the model is **_compiled_**. > [!TIP] > Compile the model when you switch to the flex processor. An uncompiled `flex_attention` materializes the full diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 1f42338127ef..7ef5975ac19d 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -23,7 +23,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import logging from ...utils.peft_utils import apply_lora_scale -from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph +from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -318,22 +318,6 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): ) -@lru_cache_unless_export(maxsize=1) -def _warn_if_flex_attention_is_uncompiled(): - """Warn once per process when `flex_attention` has not been compiled. - - `dispatch_attention_fn` reaches `flex_attention` through its module, so a user who compiles it — directly or by - compiling the model — is picked up here. Uncompiled, flex_attention falls back to a dense fp32 score matrix, which - is far slower and runs out of memory at high resolution, so say so rather than let it happen quietly. - """ - if not hasattr(flex_attention_module.flex_attention, "_torchdynamo_orig_callable"): - logger.warning( - "`QwenImage21FlexAttnProcessor` is running an uncompiled `flex_attention`, which materializes the full " - "attention score matrix in fp32 and will run out of memory at high resolution. Compile the model with " - "`transformer.compile()`, or switch to `QwenImage21AttnProcessor`." - ) - - def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]: """Split the prefix into `(start, end, is_text)` runs of equal `image_ids`. @@ -408,6 +392,7 @@ class QwenImage21FlexAttnProcessor: # for its `BlockMask` and is not configurable. _attention_backend = None _parallel_config = None + _warned_uncompiled = False def __init__(self): if not _FLEX_AVAILABLE: @@ -436,8 +421,19 @@ def __call__( if isinstance(attention_mask, BlockMask): # prefill: the BlockMask expresses the block-causal structure in one flex call. Query and key are # padded up to the mask's block-quantized length. - if not torch.compiler.is_compiling(): - _warn_if_flex_attention_is_uncompiled() + # `dispatch_attention_fn` reaches flex_attention through its module, so a user who compiled it — directly + # or by compiling the model — is picked up here. Tracing means the model is compiled. + if ( + not self._warned_uncompiled + and not torch.compiler.is_compiling() + and not hasattr(flex_attention_module.flex_attention, "_torchdynamo_orig_callable") + ): + logger.warning( + "`QwenImage21FlexAttnProcessor` is running an uncompiled `flex_attention`, which materializes the " + "full attention score matrix in fp32 and will run out of memory at high resolution. Compile the " + "model with `transformer.compile()`, or switch to `QwenImage21AttnProcessor`." + ) + QwenImage21FlexAttnProcessor._warned_uncompiled = True pad_q = int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_q pad_kv = int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_kv # Pad the sequence axis. `F.pad` counts from the last dimension, so the head and channel axes are From 450abe2a46560fd5724d6f5facafad2651a333b1 Mon Sep 17 00:00:00 2001 From: naykun Date: Thu, 17 Sep 2026 20:50:57 +0800 Subject: [PATCH 16/20] fix-copies: drop the AvgDown3D marker `AvgDown3D` validates its channel divisibility with a `ValueError` before assigning its fields, where Wan still asserts after, so the copy is not consistent and `check_copies` fails on it. The other ten markers from #5 are fine and stay. --- src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index 412d64e083c1..dd32268903d6 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -31,7 +31,6 @@ CACHE_T = 2 -# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AvgDown3D with AvgDown3D->QwenImage21AvgDown3D class QwenImage21AvgDown3D(nn.Module): def __init__( self, From 0a33bcb288beea7bdc3ba03581fc56d5ef8628ab Mon Sep 17 00:00:00 2001 From: naykun Date: Thu, 17 Sep 2026 22:54:23 +0800 Subject: [PATCH 17/20] fix: feed the transformer the pre-norm text hidden state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The transformer was trained on the last decoder layer's output of the text encoder, before the encoder's final RMSNorm. Up to transformers 4.x that is what `hidden_states[-1]` holds. From transformers 5.0 the output capturing ties that entry to `last_hidden_state`, so it comes back normalized instead, and the transformer reads something a third of the way off — visible first as garbled text in the rendered image. Neutralize the final norm for the encoder call with a forward hook that returns the module's input, so `hidden_states[-1]` is the layer output on either version. Nothing else is touched: the weights stay untouched, which matters because they are on `meta` under offloading, and no version check is needed. On the released checkpoint the prompt embeddings now match the pre-norm value exactly and the rendered image is pixel-identical to it, where before the whole image shifted by 5.35/255 on average. --- .../qwenimage21/pipeline_qwenimage21.py | 20 +++++++++- .../pipelines/qwenimage21/test_qwenimage21.py | 39 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 6e11b1e0d4a9..e95156ad63de 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -284,6 +284,24 @@ def _downsample_image_pad_tokens(hidden_states_list, image_pad_mask_list): out_mask.append(result_mask) return out_hs, out_mask + def _encode_text(self, forward_kwargs: dict) -> Any: + """ + Run the text encoder so that `hidden_states[-1]` is the last decoder layer's output, before the encoder's + final RMSNorm. That is what the transformer was trained on. + + Up to transformers 4.x it is what `hidden_states[-1]` already holds. From transformers 5.0 the output + capturing ties that entry to `last_hidden_state`, so it comes back normalized instead — a third of the signal + the transformer reads, which shows up first in rendered text. A forward hook returning the module's input + replaces its output, which neutralizes the norm for this call and leaves the behaviour the same on either + version. + """ + text_model = getattr(self.text_encoder.model, "language_model", self.text_encoder.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + return self.text_encoder(**forward_kwargs) + finally: + handle.remove() + def _get_qwen_prompt_embeds( self, prompt: str | list[str] = None, @@ -332,7 +350,7 @@ def _get_qwen_prompt_embeds( if hasattr(model_inputs, "mm_token_type_ids"): forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids - outputs = self.text_encoder(**forward_kwargs) + outputs = self._encode_text(forward_kwargs) hidden_states = outputs.hidden_states[-1] split_hidden_states = list(self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)) diff --git a/tests/pipelines/qwenimage21/test_qwenimage21.py b/tests/pipelines/qwenimage21/test_qwenimage21.py index 1fd501950d9c..b3e3989cee79 100644 --- a/tests/pipelines/qwenimage21/test_qwenimage21.py +++ b/tests/pipelines/qwenimage21/test_qwenimage21.py @@ -192,6 +192,45 @@ def test_inference_batch_single_identical(self): # batching. super().test_inference_batch_single_identical(expected_max_diff=2e-3) + def test_prompt_embeds_are_pre_norm(self): + """ + The transformer was trained on the last decoder layer's output before the text encoder's final RMSNorm, and + `hidden_states[-1]` stopped being that value in transformers 5.0. What the transformer's text projection reads + has to match the pre-norm value on either version. + + The comparison is after `txt_in.text_norm`, which is where the two forms become equivalent: it cancels the + per-token scale the text encoder's norm applied, leaving only that norm's weight to undo. They agree to a + fraction of a percent rather than exactly, because neither RMSNorm's epsilon cancels. + """ + pipe = self.get_pipeline() + text_model = getattr(pipe.text_encoder.model, "language_model", pipe.text_encoder.model) + # A freshly initialized RMSNorm weight is all ones, which is exactly the case where the two forms agree by + # accident. The released text encoder's weight is not, so give the dummy one some spread. + with torch.no_grad(): + text_model.norm.weight.copy_(torch.linspace(0.5, 2.0, text_model.norm.weight.numel())) + + pre_norm = {} + handle = text_model.layers[-1].register_forward_hook( + lambda module, args, output: pre_norm.__setitem__( + "value", output[0] if isinstance(output, tuple) else output + ) + ) + try: + prompt_embeds, _, _ = pipe.encode_prompt(prompt="a cat") + finally: + handle.remove() + + text_norm = pipe.transformer.txt_in.text_norm + with torch.no_grad(): + expected = text_norm(pre_norm["value"][:, pipe._drop_idx_t2i :]) + fixed = text_norm(prompt_embeds) + # what the pipeline would read if the post-norm hidden state went through unchanged + unfixed = text_norm(prompt_embeds * text_model.norm.weight) + + scale = expected.abs().mean() + assert (fixed - expected).abs().mean() / scale < 0.01 + assert (unfixed - expected).abs().mean() / scale > 0.1 + def test_inference(self): # Run on CPU: the expected slice below is CPU-specific. pipe = self.get_pipeline() From 21b3a41d7f82acbe69fc6f00c1caf27d37c31f9a Mon Sep 17 00:00:00 2001 From: naykun Date: Fri, 18 Sep 2026 01:39:24 +0800 Subject: [PATCH 18/20] fix: match the checkpoint's text conditioning, and repair the unexercised paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aligned with the text encoder the checkpoint was trained with: - The image marker is ``, ``, … as in training, not `Picture 1: `. The two tokenize to different lengths (4 tokens against 5), so the conditioning the transformer read was a sequence that never occurred in training, and the rendered image moves 4.02/255 on average. With the marker corrected the output is pixel-identical to the training template. This also retires the `random.choice` over four spellings of that word, which ran on the global RNG and left image-conditioned generation irreproducible from `generator`. - Condition images reach the vision encoder with their alpha composited over white, as in training. The VAE still reads all four channels. - The processor pads on the left, as in training. The joint sequence is re-padded on the right either way, so this only changes the positions the encoder itself sees for a batch of prompts of different lengths. - An empty prompt becomes a space. Qwen has no bos token, so the encoder would otherwise have nothing to read. Prompt embeddings: - The 2D prompt mask was repeated with `repeat(1, n, 1)`, which prepends an axis and tiles the rows where the 3D embeddings interleave theirs. With more than one prompt and more than one image per prompt, each sample was denoised against another prompt's padding. - Supplying `prompt_embeds` raised: `image_pad_mask` only comes from `encode_prompt`'s own encoding, and the target slots appended to it were sized from `latents`. It is synthesized for text embeddings now, and required when the embeddings cover condition images. Supplying embeddings without a mask raised too. - `has_neg_prompt` no longer requires `negative_prompt_embeds_mask`, so a caller who passes `encode_prompt`'s own output back in keeps guidance. An unpadded prompt returns `None` for the mask, the pattern `pipeline_qwenimage.py` uses, since a mask that carries no information costs the backends that reject one. Condition images: - A tensor or ndarray `image` raised on `image.size`. They are normalized to PIL up front; a latents tensor is rejected with a message, because the text encoder has to see the image, and a per-prompt nested list with another, because one flat set applies to the whole batch. That also removes the half-wired path where a latents tensor was silently dropped from both the prompt and the latents. - A list of prompts with a condition image raised a bare `StopIteration`: every prompt's template repeats the placeholders, but the processor was handed one set of images. Denoising and validation: - Interrupting on the first step used to `continue`, skipping the step that prefills the KV cache and leaving the next one to decode from an empty one. It breaks out now. - `kv_cache_mode` without a `kv_cache` is rejected instead of failing later on a shape. - `prepare_latents` checks the generator list before spending a VAE encode per image. - The VAE's image convolution names its limitation instead of asserting: it folds the single frame away and has no temporal context, so it cannot take a feature cache. Removed: `_downsample_image_pad_tokens`, `_max_length`, the `_drop_idx_ti2i` alias, and the KV cache's `is_populated` and `clear()`, none of which anything reached. Training collapses each run of `<|image_pad|>` to one token; skipping that is equivalent here, since those positions are overwritten by the VAE latents either way. The text encoder hook now points at huggingface/transformers#48087, which lets the config untie `hidden_states[-1]` from 5.18 and will make the hook unnecessary. --- .../autoencoder_kl_qwenimage21.py | 6 +- .../transformers/transformer_qwenimage21.py | 14 +- .../qwenimage21/pipeline_qwenimage21.py | 211 +++++++++--------- .../pipelines/qwenimage21/test_qwenimage21.py | 2 +- 4 files changed, 112 insertions(+), 121 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py index dd32268903d6..15575a1c5907 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -171,7 +171,11 @@ def __init__( def forward(self, x, cache_x=None): padding = list(self._padding) - assert cache_x is None + if cache_x is not None: + raise ValueError( + "This convolution is the image specialization of Wan's causal 3D one: it folds the single frame away " + "and has no temporal context to prepend, so it cannot take a feature cache." + ) x = x.squeeze(2) # Remove the temporal dimension x = F.pad(x, padding) x = super().forward(x) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 7ef5975ac19d..b6eabfd584bb 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -73,14 +73,6 @@ def get(self) -> tuple[torch.Tensor, torch.Tensor]: raise RuntimeError("KV cache has not been populated yet.") return self.k, self.v - @property - def is_populated(self) -> bool: - return self.k is not None - - def clear(self): - self.k = None - self.v = None - class QwenImage21KVCache: """Container for all transformer blocks' prefix KV caches.""" @@ -91,10 +83,6 @@ def __init__(self, num_layers: int): def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache: return self.layer_caches[layer_idx] - def clear(self): - for cache in self.layer_caches: - cache.clear() - # Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen def apply_rotary_emb_qwen( @@ -912,6 +900,8 @@ def forward( raise ValueError( f"kv_cache_mode must be 'extract' or 'cached' when kv_cache is provided, got {kv_cache_mode!r}." ) + if kv_cache is None and kv_cache_mode is not None: + raise ValueError(f"kv_cache_mode is {kv_cache_mode!r} but no kv_cache was passed to hold the prefix.") hidden_states = self.img_in(hidden_states) encoder_hidden_states = self.txt_in(encoder_hidden_states) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index e95156ad63de..15669058a35d 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -14,7 +14,6 @@ import inspect import math -import random from typing import Any, Callable import numpy as np @@ -216,18 +215,15 @@ def __init__( ) self.prompt_template_ti2i = ( f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" - f"<|im_start|>user\nPicture 1: <|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n" + f"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n" f"<|im_start|>assistant\n" ) - self.ref_token_list = ["Picture ", "Image ", "图 ", "图片 "] # Number of leading system-role tokens to drop from the hidden states. Derived from the # tokenized system message rather than hardcoded, so it tracks the processor's template. sys_message = [{"role": "system", "content": [{"type": "text", "text": self.sys_prompt}]}] sys_tokens = self.processor.apply_chat_template(sys_message, tokenize=True, return_dict=False) - self._drop_idx_t2i = len(sys_tokens[0]) - self._drop_idx_ti2i = self._drop_idx_t2i + self._drop_idx = len(sys_tokens[0]) self._img_token_id = self.processor.tokenizer.encode("<|image_pad|>")[0] - self._max_length = 8192 def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -235,73 +231,6 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor selected = hidden_states[bool_mask] return torch.split(selected, valid_lengths.tolist(), dim=0) - @staticmethod - def _downsample_image_pad_tokens(hidden_states_list, image_pad_mask_list): - """Collapse consecutive `<|image_pad|>` tokens into one per contiguous region. - - The vision-language processor expands each condition image into many vision tokens, but the transformer expects - one token per image slot, which it then expands 4x. Keep the first token of each contiguous image-pad region - and drop the rest. - """ - out_hs, out_mask = [], [] - for hidden_state, pad_mask in zip(hidden_states_list, image_pad_mask_list): - non_pad = ~pad_mask - non_pad_tokens = hidden_state[non_pad] - non_pad_mask = pad_mask[non_pad] - - pad_indices = torch.where(pad_mask)[0] - insert_tokens, insert_positions = [], [] - if len(pad_indices) > 0: - region_starts = [pad_indices[0].item()] - if len(pad_indices) > 1: - diff = torch.diff(pad_indices) - for j, d in enumerate(diff): - if d > 1: - region_starts.append(pad_indices[j + 1].item()) - - for start_idx in region_starts: - insert_pos = non_pad[:start_idx].sum().item() - insert_tokens.append(hidden_state[start_idx]) - insert_positions.append(insert_pos) - - result_tokens = non_pad_tokens - result_mask = non_pad_mask - if insert_positions: - for idx in sorted(range(len(insert_positions)), key=lambda i: insert_positions[i], reverse=True): - pos = insert_positions[idx] - result_tokens = torch.cat( - [result_tokens[:pos], insert_tokens[idx].unsqueeze(0), result_tokens[pos:]] - ) - result_mask = torch.cat( - [ - result_mask[:pos], - torch.tensor([True], dtype=torch.bool, device=result_mask.device), - result_mask[pos:], - ] - ) - - out_hs.append(result_tokens) - out_mask.append(result_mask) - return out_hs, out_mask - - def _encode_text(self, forward_kwargs: dict) -> Any: - """ - Run the text encoder so that `hidden_states[-1]` is the last decoder layer's output, before the encoder's - final RMSNorm. That is what the transformer was trained on. - - Up to transformers 4.x it is what `hidden_states[-1]` already holds. From transformers 5.0 the output - capturing ties that entry to `last_hidden_state`, so it comes back normalized instead — a third of the signal - the transformer reads, which shows up first in rendered text. A forward hook returning the module's input - replaces its output, which neutralizes the norm for this call and leaves the behaviour the same on either - version. - """ - text_model = getattr(self.text_encoder.model, "language_model", self.text_encoder.model) - handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) - try: - return self.text_encoder(**forward_kwargs) - finally: - handle.remove() - def _get_qwen_prompt_embeds( self, prompt: str | list[str] = None, @@ -310,31 +239,46 @@ def _get_qwen_prompt_embeds( ): device = device or self._execution_device prompt = [prompt] if isinstance(prompt, str) else prompt + # Qwen has no bos token, so an empty string leaves the encoder with nothing to read. + prompt = [" " if not p else p for p in prompt] is_t2i = image is None if is_t2i: prompts = [self.prompt_template_t2i.format(t) for t in prompt] - drop_idx = self._drop_idx_t2i else: prompts = [] condition_pil_list = [] for t in prompt: n_imgs = len(image) - replace = "Picture 1: <|vision_start|><|image_pad|><|vision_end|>" + replace = "<|vision_start|><|image_pad|><|vision_end|>" for i in range(2, n_imgs + 1): - replace += f" Picture {i}: <|vision_start|><|image_pad|><|vision_end|>" + replace += f" <|vision_start|><|image_pad|><|vision_end|>" template = self.prompt_template_ti2i.replace( - "Picture 1: <|vision_start|><|image_pad|><|vision_end|>", - replace.replace("Picture ", random.choice(self.ref_token_list)), + "<|vision_start|><|image_pad|><|vision_end|>", replace ) prompts.append(template.format(t)) - for img in image: - if not isinstance(img, PILImage.Image): - img = PILImage.fromarray(img) - condition_pil_list.append(img) - drop_idx = self._drop_idx_ti2i - - processor_kwargs = {"text": prompts, "padding": True, "return_tensors": "pt"} + # Each prompt's template repeats the `<|image_pad|>` placeholders, so hand the processor one set of + # images per prompt, in the order the placeholders appear. + for _ in prompt: + for img in image: + if not isinstance(img, PILImage.Image): + img = PILImage.fromarray(img) + if img.mode == "RGBA": + # The checkpoint was trained with the alpha composited over white for the vision encoder. + # Only this copy is flattened; the VAE still reads all four channels. + white = PILImage.new("RGB", img.size, (255, 255, 255)) + white.paste(img, mask=img.getchannel("A")) + img = white + condition_pil_list.append(img) + + # Left padding, as the checkpoint was trained with. `_extract_masked_hidden` drops the padding either way, + # but the side decides the positions the encoder sees for a batch of prompts of different lengths. + processor_kwargs = { + "text": prompts, + "padding": True, + "padding_side": "left", + "return_tensors": "pt", + } if not is_t2i: processor_kwargs["images"] = condition_pil_list @@ -350,17 +294,29 @@ def _get_qwen_prompt_embeds( if hasattr(model_inputs, "mm_token_type_ids"): forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids - outputs = self._encode_text(forward_kwargs) + # `hidden_states[-1]` has to be the last decoder layer's output, before the text encoder's final RMSNorm: + # that is what the transformer was trained on. It is what transformers 4.x returns there, but from + # transformers 5.0 the output capturing ties that entry to `last_hidden_state`, so it comes back normalized + # instead — a third of the signal the transformer reads, which shows up first in rendered text. A forward hook + # returning the module's input replaces its output, which neutralizes the norm for this call on either version. + # transformers 5.18 will accept `tie_last_hidden_states=False` in the text encoder's config + # (huggingface/transformers#48087); this can go once that is the floor. + text_model = getattr(self.text_encoder.model, "language_model", self.text_encoder.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + outputs = self.text_encoder(**forward_kwargs) + finally: + handle.remove() hidden_states = outputs.hidden_states[-1] split_hidden_states = list(self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)) - split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + split_hidden_states = [e[self._drop_idx :] for e in split_hidden_states] image_pad_mask = [ (sample_ids[sample_mask.bool()] == self._img_token_id) for sample_ids, sample_mask in zip(model_inputs.input_ids, model_inputs.attention_mask) ] - image_pad_mask = [e[drop_idx:] for e in image_pad_mask] + image_pad_mask = [e[self._drop_idx :] for e in image_pad_mask] attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max(e.size(0) for e in split_hidden_states) @@ -404,12 +360,30 @@ def encode_prompt( if prompt_embeds is None: prompt_embeds, prompt_embeds_mask, image_pad_mask = self._get_qwen_prompt_embeds(prompt, image, device) + elif image_pad_mask is None: + if image is not None: + raise ValueError( + "Pass `image_pad_mask` alongside `prompt_embeds` when the embeddings cover condition images, so " + "the transformer knows which positions hold image tokens." + ) + # Embeddings supplied without a mask can only be text, so no position holds an image token. + image_pad_mask = prompt_embeds.new_zeros(prompt_embeds.shape[:2], dtype=torch.bool) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) - prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + # `repeat(1, n)` on the 2D mask, so its rows interleave the same way the 3D embeddings' do. With + # `repeat(1, n, 1)` the mask picks up a leading axis and the rows come out tiled instead, which pairs each + # sample with another prompt's padding. + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + # Without padding there is nothing to mask, and a mask that carries no information costs the attention + # backends that reject one outright. + if prompt_embeds_mask is not None and prompt_embeds_mask.all(): + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask, image_pad_mask def check_inputs(self, prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs): @@ -475,12 +449,18 @@ def prepare_latents( height = 2 * (int(height) // (self.vae_scale_factor * 2)) width = 2 * (int(width) // (self.vae_scale_factor * 2)) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + image_latents = None if images is not None: all_image_latents = [] for image in images: image = image.to(device=device, dtype=dtype) - encoded = image if image.shape[1] == self.latent_channels else self._encode_vae_image(image, generator) + encoded = self._encode_vae_image(image, generator) if batch_size > encoded.shape[0]: if batch_size % encoded.shape[0] != 0: raise ValueError( @@ -495,12 +475,6 @@ def prepare_latents( ) image_latents = torch.cat(all_image_latents, dim=1) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - if latents is None: shape = (batch_size, 1, num_channels_latents, height, width) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) @@ -560,8 +534,9 @@ def __call__( prompt (`str` or `list[str]`, *optional*): The prompt to guide image generation. Pass `prompt_embeds` instead to supply embeddings directly. image (`PipelineImageInput`, *optional*): - One or more condition images. They are encoded by the text encoder as vision context and by the VAE - into latent tokens prepended to the noise. + One or more condition images, as a PIL image or a numpy array. They are encoded by the text encoder as + vision context and by the VAE into latent tokens prepended to the noise. A list is one set of images + shared by every prompt in the batch, not one entry per prompt. negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide image generation. Ignored when `true_cfg_scale` is not greater than 1. true_cfg_scale (`float`, *optional*, defaults to 1.0): @@ -621,9 +596,30 @@ def __call__( element is a list with the generated images. """ if image is not None: - image_size = image[-1].size if isinstance(image, list) else image.size + # The text encoder reads each condition image as vision context, so the pixels have to be there. Normalize + # to PIL up front, and everything downstream — the aspect ratio below, the resize, the VAE — sees one type. + image = image if isinstance(image, list) else [image] + condition_images = [] + for img in image: + if isinstance(img, PILImage.Image): + condition_images.append(img) + elif isinstance(img, np.ndarray): + condition_images.append(PILImage.fromarray(img)) + elif isinstance(img, (list, tuple)): + raise ValueError( + "`image` is one flat set of condition images that applies to every prompt in the batch, so it " + "cannot be nested per prompt. Call the pipeline once per prompt when they need different " + "condition images." + ) + else: + raise ValueError( + f"`image` accepts a PIL image or a numpy array, or a list of either, but got " + f"{type(img).__name__}. Latents cannot stand in for a condition image here, because the text " + f"encoder has to see the image itself." + ) + image = condition_images calculated_width, calculated_height, _ = calculate_dimensions( - output_resolution * output_resolution, image_size[0] / image_size[1] + output_resolution * output_resolution, image[-1].size[0] / image[-1].size[1] ) height = height or calculated_height width = width or calculated_width @@ -651,8 +647,7 @@ def __call__( # 1. Preprocess condition images: one resize feeds both the text encoder and the VAE. input_image_sizes, input_images, vae_images = [], None, None - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = image if isinstance(image, list) else [image] + if image is not None: input_images, vae_images = [], [] for img in image: if hasattr(img, "mode") and img.mode != "RGBA": @@ -668,9 +663,9 @@ def __call__( ) # 2. Encode prompt - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + # The mask is not part of the condition: `encode_prompt` returns `None` for it when nothing is padded, so + # requiring it here would turn guidance off for a caller who passes that output straight back in. + has_neg_prompt = negative_prompt is not None or negative_prompt_embeds is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt if true_cfg_scale > 1 and not has_neg_prompt: logger.warning( @@ -761,7 +756,9 @@ def append_target_slots(mask): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: - continue + # `continue` would skip the step that prefills the cache and leave the next one decoding from an + # empty one, so stop the loop instead. + break self._current_timestep = t kv_mode = "extract" if (cache_enabled and i == 0) else ("cached" if cache_enabled else None) diff --git a/tests/pipelines/qwenimage21/test_qwenimage21.py b/tests/pipelines/qwenimage21/test_qwenimage21.py index b3e3989cee79..1d79334cdfc9 100644 --- a/tests/pipelines/qwenimage21/test_qwenimage21.py +++ b/tests/pipelines/qwenimage21/test_qwenimage21.py @@ -222,7 +222,7 @@ def test_prompt_embeds_are_pre_norm(self): text_norm = pipe.transformer.txt_in.text_norm with torch.no_grad(): - expected = text_norm(pre_norm["value"][:, pipe._drop_idx_t2i :]) + expected = text_norm(pre_norm["value"][:, pipe._drop_idx :]) fixed = text_norm(prompt_embeds) # what the pipeline would read if the post-norm hidden state went through unchanged unfixed = text_norm(prompt_embeds * text_model.norm.weight) From 24c36a6cff9364a46077d4f67fc9c91ea8869b7a Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 18 Sep 2026 07:30:10 +0300 Subject: [PATCH 19/20] fix callback test (#2) --- src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 15669058a35d..7233cef43bec 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -808,7 +808,9 @@ def append_target_slots(mask): latents = latents.to(latents_dtype) if callback_on_step_end is not None: - callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs} + 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) From 8d3c30bfda9b511c00992f40cff4170a5502814d Mon Sep 17 00:00:00 2001 From: naykun Date: Fri, 18 Sep 2026 12:34:33 +0800 Subject: [PATCH 20/20] Add a TODO at the text encoder hook It can be replaced with `tie_last_hidden_states=False` in the text encoder's config once huggingface/transformers#48087 ships in a stable release. --- src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py index 7233cef43bec..786b09b4e3cd 100644 --- a/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -299,8 +299,8 @@ def _get_qwen_prompt_embeds( # transformers 5.0 the output capturing ties that entry to `last_hidden_state`, so it comes back normalized # instead — a third of the signal the transformer reads, which shows up first in rendered text. A forward hook # returning the module's input replaces its output, which neutralizes the norm for this call on either version. - # transformers 5.18 will accept `tie_last_hidden_states=False` in the text encoder's config - # (huggingface/transformers#48087); this can go once that is the floor. + # TODO: replace this with `tie_last_hidden_states=False` in the text encoder's config, which + # huggingface/transformers#48087 adds, once that ships in a stable transformers release (5.18). text_model = getattr(self.text_encoder.model, "language_model", self.text_encoder.model) handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) try: