diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f2801a0cd6a6..525734a2e4cb 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 @@ -637,6 +641,8 @@ title: PRX - local: api/pipelines/prx_pixel title: PRX Pixel + - local: api/pipelines/qwenimage21 + title: Qwen-Image 2.1 - local: api/pipelines/qwenimage title: QwenImage - local: api/pipelines/sana 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..fefb50b53c16 --- /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.autoencoder_kl.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..022f0710c173 --- /dev/null +++ b/docs/source/en/api/models/qwenimage21_transformer2d.md @@ -0,0 +1,45 @@ + + +# 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 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. `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 independent of the denoising step, so the keys and values of that prefix are cacheable + across steps via the `kv_cache` argument. + +Load it with: + +```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..64461ed8c0d0 --- /dev/null +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -0,0 +1,73 @@ + + +# 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 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 + +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").images[0] +image.save("t2i.png") + +# Image-conditioned editing +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_**. + +> [!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 + - all + - __call__ + +## QwenImagePipelineOutput + +[[autodoc]] pipelines.qwenimage.pipeline_output.QwenImagePipelineOutput diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 3ed956c75e49..cee61551a5c8 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -253,6 +253,7 @@ "AutoencoderKLMiniMaxH3Audio", "AutoencoderKLMochi", "AutoencoderKLQwenImage", + "AutoencoderKLQwenImage21", "AutoencoderKLTemporalDecoder", "AutoencoderKLWan", "AutoencoderOobleck", @@ -332,6 +333,7 @@ "PixArtTransformer2DModel", "PriorTransformer", "PRXTransformer2DModel", + "QwenImage21Transformer2DModel", "QwenImageControlNetModel", "QwenImageMultiControlNetModel", "QwenImageTransformer2DModel", @@ -764,6 +766,7 @@ "PixArtSigmaPipeline", "PRXPipeline", "PRXPixelPipeline", + "QwenImage21Pipeline", "QwenImageControlNetInpaintPipeline", "QwenImageControlNetPipeline", "QwenImageEditInpaintPipeline", @@ -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..15575a1c5907 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py @@ -0,0 +1,1561 @@ +# 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__() + 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 = 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 + 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 + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.DupUp3D with DupUp3D->QwenImage21DupUp3D +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) + 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) + x = x.unsqueeze(2) # Add the temporal dimension back + 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. + + 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): + 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. + + 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=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: + 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. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = get_activation("silu") + + # 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=None): + if feat_idx is 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 + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanAttentionBlock with Wan->QwenImage21 +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. + """ + + 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)] + attentions = [] + for _ in range(num_layers): + attentions.append(QwenImage21AttentionBlock(dim)) + resnets.append(QwenImage21ResidualBlock(dim, dim, dropout)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + 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) + + # 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=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) + 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. + """ + + 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, + 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("silu") + + # 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, 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=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() + 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 + """ + + 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, + ): + 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)) + 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=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] + """ + 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') + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: str | None = None, + ): + 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)) + 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=None, first_chunk=None): + if feat_idx is None: + feat_idx = [0] + """ + 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. + """ + + 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, + 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("silu") + + # 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, 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, + ) + else: + up_block = QwenImage21UpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + ) + 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=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] + ## 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 + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.patchify with patchify->_patchify +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 + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.unpatchify with unpatchify->_unpatchify +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 = 4, + out_channels: int = 4, + patch_size: int | None = None, + scale_factor_temporal: int | None = 8, + scale_factor_spatial: int | None = 16, + ) -> 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, + } + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.enable_tiling + 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 + + # 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"] + 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 + + # 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 + + 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 + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.encode + 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) + + # 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 + 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 + # 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. + + 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) + + # 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): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + 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): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + 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. + + 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 + + # 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. + + 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. + 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 + + 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..b6eabfd584bb --- /dev/null +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -0,0 +1,1018 @@ +# 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 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 +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 + +# 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: `QwenImage21FlexAttnProcessor` needs it, `QwenImage21AttnProcessor` does not. +_FLEX_AVAILABLE = False +try: + 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_module = None + + +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 + + +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] + + +# 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, + ) + + +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, + 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: + # `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].clone(), + value[:, cache_write_slice].clone(), + ) + 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) + + seq_len_q = query.shape[1] + return query, key, value, seq_len_q + + +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. + + 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. + """ + + # 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 + _warned_uncompiled = False + + 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", + 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, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | 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 + ) + + seq_len_kv = key.shape[1] + 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. + # `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 + # padded by zero first. The result stays contiguous, which the compiled flex kernel requires. + if pad_q: + query = F.pad(query, (0, 0, 0, 0, 0, pad_q)) + if pad_kv: + 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] + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + backend=self._attention_backend, + 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 QwenImage21AttnProcessor: + r""" + Attention processor for Qwen-Image 2.1 that needs neither `flex_attention` nor a compiled model. + + 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 + _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, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | 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 segments 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=self._attention_backend, + parallel_config=self._parallel_config, + ) + else: + # 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 + 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) + + 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. + """ + + # 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__() + 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, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, + ) -> 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, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, + cache_write_slice=cache_write_slice, + segments=segments, + key_valid=key_valid, + ) + 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: + + - **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. + + 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. + """ + + _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, + ): + 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 + + @apply_lora_scale("attention_kwargs") + 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: QwenImage21KVCache | None = None, + kv_cache_mode: str | 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 (`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`. + 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] + 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}." + ) + 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) + + # 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, _IMG_TOKENS_PER_SLOT, 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) + + # 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()) + + 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 + block_segments, block_key_valid = None, None + else: + # 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_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 + 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_segments, + block_key_valid, + ) + 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, + segments=block_segments, + key_valid=block_key_valid, + ) + + joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) + output = self.proj_out(joint_hidden_states) + + 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..786b09b4e3cd --- /dev/null +++ b/src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py @@ -0,0 +1,849 @@ +# 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 +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 ...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 +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", dtype=torch.bfloat16) + >>> pipe.to("cuda") + >>> prompt = "A capybara wearing a wizard hat, reading a book by candlelight, oil painting" + >>> image = pipe(prompt).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 + + +# 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": + 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") + + +# 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 + + width = round(width / 32) * 32 + height = round(height / 32) * 32 + + return width, height, None + + +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\n<|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + # 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 = len(sys_tokens[0]) + self._img_token_id = self.processor.tokenizer.encode("<|image_pad|>")[0] + + 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) + + 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 + # 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] + else: + prompts = [] + condition_pil_list = [] + for t in prompt: + n_imgs = len(image) + replace = "<|vision_start|><|image_pad|><|vision_end|>" + for i in range(2, n_imgs + 1): + replace += f" <|vision_start|><|image_pad|><|vision_end|>" + template = self.prompt_template_ti2i.replace( + "<|vision_start|><|image_pad|><|vision_end|>", replace + ) + prompts.append(template.format(t)) + # 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 + + 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) + if hasattr(model_inputs, "mm_token_type_ids"): + forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids + + # `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. + # 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: + 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[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[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) + 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) + 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) + # `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): + 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 + + # 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 = [ + 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) + ) + 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 + ): + 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 = 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 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 = 1.0, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 40, + 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, 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): + 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 40): + 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. + 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`): + 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. + 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`): + 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: + # 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[-1].size[0] / image[-1].size[1] + ) + height = height or calculated_height + width = width or calculated_width + 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._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: + 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 + ) + 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 + # 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( + 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. + # 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 = 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. + 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 + neg_cache = QwenImage21KVCache(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` 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) + + 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, + kv_cache_mode=kv_mode, + 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, + kv_cache_mode=kv_mode, + 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 = {} + 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) + + 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/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.""" 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..cb23fc4ad1d2 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -0,0 +1,296 @@ +# 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_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] + + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + torch.manual_seed(0) + 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, 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_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 + + 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=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 ( + QwenImage21AttnProcessor, + QwenImage21FlexAttnProcessor, + QwenImage21KVCache, + ) + + 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, QwenImage21AttnProcessor): + 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) + + @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: + """ + 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): + def test_gradient_checkpointing_is_applied(self): + expected_set = {"QwenImage21Transformer2DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, AttentionTesterMixin): + pass 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..1d79334cdfc9 --- /dev/null +++ b/tests/pipelines/qwenimage21/test_qwenimage21.py @@ -0,0 +1,262 @@ +# 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_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 :]) + 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() + + 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