From e0d8014dc65279eed1cc0ebf4693e16c4437df17 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Tue, 6 Oct 2026 19:38:27 -0700 Subject: [PATCH] add modular pipeline for qwen-image 2.1 --- docs/source/en/api/pipelines/qwenimage21.md | 30 + src/diffusers/__init__.py | 4 + src/diffusers/modular_pipelines/__init__.py | 2 + .../modular_pipelines/modular_pipeline.py | 1 + .../modular_pipelines/qwenimage21/__init__.py | 47 ++ .../qwenimage21/before_denoise.py | 571 ++++++++++++++ .../modular_pipelines/qwenimage21/decoders.py | 186 +++++ .../modular_pipelines/qwenimage21/denoise.py | 474 ++++++++++++ .../modular_pipelines/qwenimage21/encoders.py | 711 ++++++++++++++++++ .../modular_pipelines/qwenimage21/inputs.py | 376 +++++++++ .../qwenimage21/modular_blocks_qwenimage21.py | 586 +++++++++++++++ .../qwenimage21/modular_pipeline.py | 51 ++ src/diffusers/pipelines/auto_pipeline.py | 2 + .../dummy_torch_and_transformers_objects.py | 30 + .../modular_pipelines/qwenimage21/__init__.py | 0 .../test_modular_pipeline_qwenimage21.py | 261 +++++++ .../modular_pipelines/testing_utils/utils.py | 2 +- 17 files changed, 3333 insertions(+), 1 deletion(-) create mode 100644 src/diffusers/modular_pipelines/qwenimage21/__init__.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/decoders.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/denoise.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/encoders.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/inputs.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py create mode 100644 src/diffusers/modular_pipelines/qwenimage21/modular_pipeline.py create mode 100644 tests/modular_pipelines/qwenimage21/__init__.py create mode 100644 tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md index 4962cf70bb7b..93198d25b2c2 100644 --- a/docs/source/en/api/pipelines/qwenimage21.md +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -97,6 +97,28 @@ the terminal sigma, which the scheduler appends. Explicit `sigmas` override the list determines the number of steps instead of `num_inference_steps`. If neither is provided, the pipeline uses `num_inference_steps` to generate the schedule. The scheduler applies its configured processing to either grid. +## Modular + +Qwen-Image 2.1 is also available as a [modular pipeline](../../modular_diffusers/overview). The same blocks run +text-to-image and image-conditioned generation: pass `image` to condition on one or more images. Classifier-free +guidance is configured through the `guider` component rather than a `true_cfg_scale` call argument, and the prefix KV +cache is on by default (`use_kv_cache=False` turns it off). + +```python +import torch +from diffusers import ClassifierFreeGuidance, ModularPipeline + +pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image-2.1") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +image = pipe(prompt="A capybara wearing a wizard hat, oil painting").images[0] + +# Turn classifier-free guidance on for an edit. +pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=4.0)) +edited = pipe(prompt="Move it to a snowy mountain top", negative_prompt="blurry", image=image).images[0] +``` + ## QwenImage21Pipeline [[autodoc]] QwenImage21Pipeline @@ -106,3 +128,11 @@ list determines the number of steps instead of `num_inference_steps`. If neither ## QwenImagePipelineOutput [[autodoc]] pipelines.qwenimage.pipeline_output.QwenImagePipelineOutput + +## QwenImage21ModularPipeline + +[[autodoc]] QwenImage21ModularPipeline + +## QwenImage21AutoBlocks + +[[autodoc]] QwenImage21AutoBlocks diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 599237a6f2a0..339e3bec7f38 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -566,6 +566,8 @@ "MiniMaxH3ModularPipeline", "MiniMaxMusic3Blocks", "MiniMaxMusic3ModularPipeline", + "QwenImage21AutoBlocks", + "QwenImage21ModularPipeline", "QwenImageAutoBlocks", "QwenImageEditAutoBlocks", "QwenImageEditModularPipeline", @@ -1438,6 +1440,8 @@ MiniMaxH3ModularPipeline, MiniMaxMusic3Blocks, MiniMaxMusic3ModularPipeline, + QwenImage21AutoBlocks, + QwenImage21ModularPipeline, QwenImageAutoBlocks, QwenImageEditAutoBlocks, QwenImageEditModularPipeline, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index aaa37918956a..0afc7788d133 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -101,6 +101,7 @@ "QwenImageLayeredModularPipeline", "QwenImageLayeredAutoBlocks", ] + _import_structure["qwenimage21"] = ["QwenImage21AutoBlocks", "QwenImage21ModularPipeline"] _import_structure["anima"] = [ "AnimaAutoBlocks", "AnimaModularPipeline", @@ -226,6 +227,7 @@ QwenImageLayeredModularPipeline, QwenImageModularPipeline, ) + from .qwenimage21 import QwenImage21AutoBlocks, QwenImage21ModularPipeline from .stable_diffusion_3 import StableDiffusion3AutoBlocks, StableDiffusion3ModularPipeline from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline from .wan import ( diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index c501d7b1fc42..f4374d1d1af9 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -190,6 +190,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("qwenimage-edit", _create_default_map_fn("QwenImageEditModularPipeline")), ("qwenimage-edit-plus", _create_default_map_fn("QwenImageEditPlusModularPipeline")), ("qwenimage-layered", _create_default_map_fn("QwenImageLayeredModularPipeline")), + ("qwenimage21", _create_default_map_fn("QwenImage21ModularPipeline")), ("anima", _create_default_map_fn("AnimaModularPipeline")), ("z-image", _create_default_map_fn("ZImageModularPipeline")), ("cosmos3-omni", _create_default_map_fn("Cosmos3OmniModularPipeline")), diff --git a/src/diffusers/modular_pipelines/qwenimage21/__init__.py b/src/diffusers/modular_pipelines/qwenimage21/__init__.py new file mode 100644 index 000000000000..39bc4ec56652 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_qwenimage21"] = ["QwenImage21AutoBlocks"] + _import_structure["modular_pipeline"] = ["QwenImage21ModularPipeline"] + +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 .modular_blocks_qwenimage21 import QwenImage21AutoBlocks + from .modular_pipeline import QwenImage21ModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/qwenimage21/before_denoise.py b/src/diffusers/modular_pipelines/qwenimage21/before_denoise.py new file mode 100644 index 000000000000..3d9310bbf13a --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/before_denoise.py @@ -0,0 +1,571 @@ +# Copyright 2026 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 numpy as np +import torch + +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from .inputs import pack_latents +from .modular_pipeline import QwenImage21ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# 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.stable_diffusion.pipeline_stable_diffusion.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 + + +# auto_docstring +class QwenImage21PrepareLatentsStep(ModularPipelineBlocks): + """ + Prepare the initial random noise for the generation process. `height` and `width` default to `output_resolution` + and are rounded down to a multiple of 32, as the transformer groups the 16x compressed latents in 2x2 blocks. + + Inputs: + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + height (`int`): + if not set, updated to the default value + width (`int`): + if not set, updated to the default value + latents (`Tensor`): + The initial latents to use for the denoising process, packed to (B, H * W, C) + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Prepare the initial random noise for the generation process. `height` and `width` default to " + "`output_resolution` and are rounded down to a multiple of 32, as the transformer groups the 16x " + "compressed latents in 2x2 blocks." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam.template("height"), + InputParam.template("width"), + InputParam( + name="output_resolution", + type_hint=int, + default=1024, + description="Target side length used to derive the output size and to resize condition images.", + ), + InputParam.template("num_images_per_prompt"), + InputParam.template("generator"), + InputParam.template("batch_size"), + InputParam.template("dtype"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="height", type_hint=int, description="if not set, updated to the default value"), + OutputParam(name="width", type_hint=int, description="if not set, updated to the default value"), + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process, packed to (B, H * W, C)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + device = components._execution_device + batch_size = block_state.batch_size * block_state.num_images_per_prompt + + multiple_of = components.vae_scale_factor * 2 + height = block_state.height or block_state.output_resolution + width = block_state.width or block_state.output_resolution + block_state.height = height // multiple_of * multiple_of + block_state.width = width // multiple_of * multiple_of + + latent_height = block_state.height // components.vae_scale_factor + latent_width = block_state.width // components.vae_scale_factor + shape = (batch_size, components.num_channels_latents, 1, latent_height, latent_width) + + if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if block_state.latents is None: + latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=block_state.dtype) + block_state.latents = pack_latents(latents) + else: + block_state.latents = block_state.latents.to(device=device, dtype=block_state.dtype) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21SetTimestepsStep(ModularPipelineBlocks): + """ + Step that sets the scheduler's timesteps. The sampling grid comes from `sigmas`, then from the `sample_sigmas` + pipeline config, then from `num_inference_steps`. Should be run after the prepare latents step. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Configs: + sample_sigmas (default: None): Default sampling grid of the checkpoint, used when `sigmas` is not passed. + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + latents (`Tensor`): + The initial random noised latents for the denoising process. Can be generated in prepare latents step. + + Outputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process + num_inference_steps (`int`): + The number of denoising steps, updated when the grid comes from sigmas + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Step that sets the scheduler's timesteps. The sampling grid comes from `sigmas`, then from the " + "`sample_sigmas` pipeline config, then from `num_inference_steps`. Should be run after the prepare " + "latents step." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec( + name="sample_sigmas", + default=None, + description="Default sampling grid of the checkpoint, used when `sigmas` is not passed.", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", default=40), + InputParam.template("sigmas"), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The initial random noised latents for the denoising process. Can be generated in prepare latents step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" + ), + OutputParam( + name="num_inference_steps", + type_hint=int, + description="The number of denoising steps, updated when the grid comes from sigmas", + ), + ] + + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + device = components._execution_device + + sigmas = block_state.sigmas + if sigmas is None: + sigmas = components.config.sample_sigmas + if sigmas is None: + sigmas = np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps) + + mu = calculate_shift( + image_seq_len=block_state.latents.shape[1], + base_seq_len=components.scheduler.config.get("base_image_seq_len", 256), + max_seq_len=components.scheduler.config.get("max_image_seq_len", 4096), + base_shift=components.scheduler.config.get("base_shift", 0.5), + max_shift=components.scheduler.config.get("max_shift", 1.15), + ) + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + scheduler=components.scheduler, + num_inference_steps=block_state.num_inference_steps, + device=device, + sigmas=sigmas, + mu=mu, + ) + components.scheduler.set_begin_index(0) + + self.set_block_state(state, block_state) + return components, state + + +def joint_vision_mask(image_pad_mask: torch.Tensor, num_target_tokens: int) -> torch.Tensor: + # The transformer expands every vision slot into a 2x2 group of latent tokens, so the target image takes one + # slot per four latents, appended after the prompt's vision positions. + target_slots = num_target_tokens // 4 + return torch.cat([image_pad_mask, image_pad_mask.new_ones(image_pad_mask.shape[0], target_slots)], dim=1) + + +# auto_docstring +class QwenImage21RoPEInputsStep(ModularPipelineBlocks): + """ + Step that prepares the layout inputs of the transformer for text-to-image generation: the latent grid of the target + image and the joint vision mask that appends one slot per 2x2 group of target latents to the prompt positions. + Should be placed after the prepare latents step. + + Inputs: + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`): + The packed target latents. Can be generated in prepare latents step. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings. Can be generated in the text input step. + negative_image_pad_mask (`Tensor`, *optional*): + Bool mask marking the vision positions of the negative prompt embeddings. Can be generated in the text + input step. + + Outputs: + img_shapes (`list`): + Per-sample (frame, height, width) of the target image in latent tokens + img_mask (`Tensor`): + Joint vision mask over the prompt positions and the target image slots + negative_img_mask (`Tensor`): + Joint vision mask over the negative prompt positions and the target image slots + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Step that prepares the layout inputs of the transformer for text-to-image generation: the latent grid " + "of the target image and the joint vision mask that appends one slot per 2x2 group of target latents to " + "the prompt positions. Should be placed after the prepare latents step." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("batch_size"), + InputParam.template("height"), + InputParam.template("width"), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The packed target latents. Can be generated in prepare latents step.", + ), + InputParam( + name="image_pad_mask", + required=True, + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the prompt embeddings. Can be generated in the text input step.", + ), + InputParam( + name="negative_image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the negative prompt embeddings. Can be generated in the text input step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="img_shapes", + kwargs_type="denoiser_input_fields", + type_hint=list[list[tuple[int, int, int]]], + description="Per-sample (frame, height, width) of the target image in latent tokens", + ), + OutputParam( + name="img_mask", + type_hint=torch.Tensor, + description="Joint vision mask over the prompt positions and the target image slots", + ), + OutputParam( + name="negative_img_mask", + type_hint=torch.Tensor, + description="Joint vision mask over the negative prompt positions and the target image slots", + ), + ] + + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + vae_scale_factor = components.vae_scale_factor + block_state.img_shapes = [ + [(1, block_state.height // vae_scale_factor, block_state.width // vae_scale_factor)] + ] * block_state.batch_size + + num_target_tokens = block_state.latents.shape[1] + block_state.img_mask = joint_vision_mask(block_state.image_pad_mask, num_target_tokens) + block_state.negative_img_mask = None + if block_state.negative_image_pad_mask is not None: + block_state.negative_img_mask = joint_vision_mask(block_state.negative_image_pad_mask, num_target_tokens) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21ImageConditionedRoPEInputsStep(ModularPipelineBlocks): + """ + Step that prepares the layout inputs of the transformer for image-conditioned generation: the latent grid of every + condition image followed by the target image, and the joint vision mask that marks the condition-image slots in the + prompt and appends one slot per 2x2 group of target latents. Should be placed after the prepare latents step. + + Inputs: + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`): + The packed target latents. Can be generated in prepare latents step. + image_height (`list`): + The pixel heights of the condition images. Can be generated in the additional inputs step. + image_width (`list`): + The pixel widths of the condition images. Can be generated in the additional inputs step. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings. Can be generated in the text input step. + negative_image_pad_mask (`Tensor`, *optional*): + Bool mask marking the vision positions of the negative prompt embeddings. Can be generated in the text + input step. + + Outputs: + img_shapes (`list`): + Per-sample (frame, height, width) of each image in latent tokens, condition images first and the target + image last + img_mask (`Tensor`): + Joint vision mask over the prompt positions and the target image slots + negative_img_mask (`Tensor`): + Joint vision mask over the negative prompt positions and the target image slots + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Step that prepares the layout inputs of the transformer for image-conditioned generation: the latent " + "grid of every condition image followed by the target image, and the joint vision mask that marks the " + "condition-image slots in the prompt and appends one slot per 2x2 group of target latents. Should be " + "placed after the prepare latents step." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("batch_size"), + InputParam.template("height"), + InputParam.template("width"), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The packed target latents. Can be generated in prepare latents step.", + ), + InputParam( + name="image_height", + required=True, + type_hint=list[int], + description="The pixel heights of the condition images. Can be generated in the additional inputs step.", + ), + InputParam( + name="image_width", + required=True, + type_hint=list[int], + description="The pixel widths of the condition images. Can be generated in the additional inputs step.", + ), + InputParam( + name="image_pad_mask", + required=True, + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the prompt embeddings. Can be generated in the text input step.", + ), + InputParam( + name="negative_image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the negative prompt embeddings. Can be generated in the text input step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="img_shapes", + kwargs_type="denoiser_input_fields", + type_hint=list[list[tuple[int, int, int]]], + description="Per-sample (frame, height, width) of each image in latent tokens, condition images first and the target image last", + ), + OutputParam( + name="img_mask", + type_hint=torch.Tensor, + description="Joint vision mask over the prompt positions and the target image slots", + ), + OutputParam( + name="negative_img_mask", + type_hint=torch.Tensor, + description="Joint vision mask over the negative prompt positions and the target image slots", + ), + ] + + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + vae_scale_factor = components.vae_scale_factor + block_state.img_shapes = [ + [ + *[ + (1, image_height // vae_scale_factor, image_width // vae_scale_factor) + for image_height, image_width in zip(block_state.image_height, block_state.image_width) + ], + (1, block_state.height // vae_scale_factor, block_state.width // vae_scale_factor), + ] + ] * block_state.batch_size + + num_target_tokens = block_state.latents.shape[1] + block_state.img_mask = joint_vision_mask(block_state.image_pad_mask, num_target_tokens) + block_state.negative_img_mask = None + if block_state.negative_image_pad_mask is not None: + block_state.negative_img_mask = joint_vision_mask(block_state.negative_image_pad_mask, num_target_tokens) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/qwenimage21/decoders.py b/src/diffusers/modular_pipelines/qwenimage21/decoders.py new file mode 100644 index 000000000000..e49eb23f20b4 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/decoders.py @@ -0,0 +1,186 @@ +# Copyright 2026 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 + +from ...configuration_utils import FrozenDict +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKLQwenImage21 +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import QwenImage21ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def unpack_latents(latents: torch.Tensor, height: int, width: int, vae_scale_factor: int) -> torch.Tensor: + # (batch_size, height * width, channels) -> (batch_size, channels, 1, height, width) + batch_size, _, channels = latents.shape + latent_height = 2 * (int(height) // (vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (vae_scale_factor * 2)) + return latents.transpose(1, 2).reshape(batch_size, channels, 1, latent_height, latent_width) + + +# auto_docstring +class QwenImage21UnpackLatentsStep(ModularPipelineBlocks): + """ + Step that unpacks the latents from (batch_size, sequence_length, channels) into (batch_size, channels, 1, height, + width) + + Inputs: + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`): + The packed latents to unpack, can be generated in the denoise step. + + Outputs: + latents (`Tensor`): + The denoised latents unpacked to B, C, 1, H, W + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return "Step that unpacks the latents from (batch_size, sequence_length, channels) into (batch_size, channels, 1, height, width)" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height"), + InputParam.template("width"), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The packed latents to unpack, can be generated in the denoise step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="latents", type_hint=torch.Tensor, description="The denoised latents unpacked to B, C, 1, H, W" + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + block_state.latents = unpack_latents( + block_state.latents, block_state.height, block_state.width, components.vae_scale_factor + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21DecodeStep(ModularPipelineBlocks): + """ + Step that decodes the latents to RGBA images and postprocesses them + + Components: + vae (`AutoencoderKLQwenImage21`) image_processor (`VaeImageProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the unpack latents + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return "Step that decodes the latents to RGBA images and postprocesses them" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLQwenImage21), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents to decode, can be generated in the denoise step and unpacked in the unpack latents step.", + ), + InputParam.template("output_type"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + + @staticmethod + def check_inputs(output_type): + if output_type not in ["pil", "np", "pt"]: + raise ValueError(f"Invalid output_type: {output_type}") + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + self.check_inputs(block_state.output_type) + + if block_state.latents.ndim == 4: + block_state.latents = block_state.latents.unsqueeze(dim=2) + elif block_state.latents.ndim != 5: + raise ValueError( + f"expect latents to be a 4D or 5D tensor but got: {block_state.latents.shape}. Please make sure the latents are unpacked before decode step." + ) + + latents = block_state.latents.to(components.vae.dtype) + latents_mean = ( + torch.tensor(components.vae.config.latents_mean) + .view(1, components.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(components.vae.config.latents_std) + .view(1, components.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + images = components.vae.decode(latents, return_dict=False)[0][:, :, 0] + block_state.images = components.image_processor.postprocess(images, output_type=block_state.output_type) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/qwenimage21/denoise.py b/src/diffusers/modular_pipelines/qwenimage21/denoise.py new file mode 100644 index 000000000000..68ec686f000c --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/denoise.py @@ -0,0 +1,474 @@ +# Copyright 2026 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 torch + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import QwenImage21Transformer2DModel +from ...models.transformers.transformer_qwenimage21 import QwenImage21KVCache +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging +from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import QwenImage21ModularPipeline + + +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 + + +# ==================== +# 1. LOOP STEPS (run at each denoising step) +# ==================== + + +# loop step: before denoiser +class QwenImage21LoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "step within the denoising loop that prepares the latent input for the denoiser. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `QwenImage21DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> tuple[QwenImage21ModularPipeline, BlockState]: + block_state.latent_model_input = block_state.latents + block_state.timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype) + return components, block_state + + +class QwenImage21ImageConditionedLoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "step within the denoising loop that prepends the condition image tokens to the target latents. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `QwenImage21DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", + ), + InputParam( + name="image_latents", + required=True, + type_hint=torch.Tensor, + description="Packed condition image latents. Can be generated in the additional inputs step.", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> tuple[QwenImage21ModularPipeline, BlockState]: + # Condition images come first in the joint sequence, the target image last. + block_state.latent_model_input = torch.cat([block_state.image_latents, block_state.latents], dim=1) + block_state.timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype) + return components, block_state + + +# loop step: denoiser +class QwenImage21LoopDenoiser(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "step within the denoising loop that denoise the latent input for the denoiser. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `QwenImage21DenoiseLoopWrapper`)" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", QwenImage21Transformer2DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("attention_kwargs"), + InputParam.template("denoiser_input_fields"), + InputParam( + name="img_shapes", + required=True, + type_hint=list[list[tuple[int, int, int]]], + description="Per-sample latent grids, condition images first and the target image last. Can be generated in prepare_rope_inputs step.", + ), + InputParam( + name="img_mask", + required=True, + type_hint=torch.Tensor, + description="Joint vision mask over the prompt positions and the target image slots. Can be generated in prepare_rope_inputs step.", + ), + InputParam( + name="negative_img_mask", + type_hint=torch.Tensor, + description="Joint vision mask for the negative prompt. Can be generated in prepare_rope_inputs step.", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> tuple[QwenImage21ModularPipeline, BlockState]: + guider_inputs = { + "encoder_hidden_states": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "encoder_hidden_states_mask": ( + getattr(block_state, "prompt_embeds_mask", None), + getattr(block_state, "negative_prompt_embeds_mask", None), + ), + "img_mask": (block_state.img_mask, block_state.negative_img_mask), + } + + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + additional_cond_kwargs = {} + for field_name, field_value in block_state.denoiser_input_fields.items(): + if field_name in transformer_args and field_name not in guider_inputs: + additional_cond_kwargs[field_name] = field_value + block_state.additional_cond_kwargs.update(additional_cond_kwargs) + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + num_target_tokens = block_state.latents.shape[1] + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} + context_name = getattr(guider_state_batch, components.guider._identifier_key) + + # Text and condition-image keys and values are step-independent under `causal_condition`, so the first + # call of each guidance branch prefills its cache and the later calls only recompute the target tokens. + kv_cache, kv_cache_mode = None, None + if block_state.cache_enabled: + kv_cache = block_state.kv_caches.get(context_name) + kv_cache_mode = "cached" if kv_cache is not None else "extract" + if kv_cache is None: + kv_cache = QwenImage21KVCache(len(components.transformer.transformer_blocks)) + block_state.kv_caches[context_name] = kv_cache + + with components.transformer.cache_context(context_name): + noise_pred = components.transformer( + hidden_states=block_state.latent_model_input, + timestep=block_state.timestep / 1000, + attention_kwargs=block_state.attention_kwargs, + kv_cache=kv_cache, + kv_cache_mode=kv_cache_mode, + return_dict=False, + **cond_kwargs, + **block_state.additional_cond_kwargs, + )[0] + guider_state_batch.noise_pred = noise_pred[:, -num_target_tokens:] + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + +# loop step: after denoiser +class QwenImage21LoopAfterDenoiser(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "step within the denoising loop that updates the latents. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `QwenImage21DenoiseLoopWrapper`)" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents"), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> tuple[QwenImage21ModularPipeline, BlockState]: + latents_dtype = block_state.latents.dtype + block_state.latents = components.scheduler.step( + block_state.noise_pred, + t, + block_state.latents, + return_dict=False, + )[0] + + if block_state.latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + block_state.latents = block_state.latents.to(latents_dtype) + + return components, block_state + + +# ==================== +# 2. LOOP WRAPPER (the denoising loop) +# ==================== + + +class QwenImage21DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoise the latents over `timesteps`. " + "The specific steps with each iteration can be customized with `sub_blocks` attributes" + ) + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec("transformer", QwenImage21Transformer2DModel), + ] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + name="timesteps", + required=True, + type_hint=torch.Tensor, + description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", + ), + InputParam.template("num_inference_steps", required=True), + InputParam( + name="use_kv_cache", + type_hint=bool, + default=True, + description=( + "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 it does not reproduce the same image bit-for-bit in reduced precision." + ), + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + block_state.num_warmup_steps = max( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 + ) + block_state.additional_cond_kwargs = {} + # One prefix cache per guidance branch, filled on that branch's first call and discarded with the loop. + block_state.cache_enabled = block_state.use_kv_cache and components.transformer.config.causal_condition + block_state.kv_caches = {} + + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + if XLA_AVAILABLE: + xm.mark_step() + + self.set_block_state(state, block_state) + return components, state + + +# ==================== +# 3. DENOISE STEPS: compose the denoising loop with loop wrapper + loop steps +# ==================== + + +# auto_docstring +class QwenImage21DenoiseStep(QwenImage21DenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImage21DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImage21LoopBeforeDenoiser` + - `QwenImage21LoopDenoiser` + - `QwenImage21LoopAfterDenoiser` + This block supports text-to-image generation. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImage21Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + 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 it does not + reproduce the same image bit-for-bit in reduced precision. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + Per-sample latent grids, condition images first and the target image last. Can be generated in + prepare_rope_inputs step. + img_mask (`Tensor`): + Joint vision mask over the prompt positions and the target image slots. Can be generated in + prepare_rope_inputs step. + negative_img_mask (`Tensor`, *optional*): + Joint vision mask for the negative prompt. Can be generated in prepare_rope_inputs step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = [ + QwenImage21LoopBeforeDenoiser, + QwenImage21LoopDenoiser, + QwenImage21LoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents. \n" + "Its loop logic is defined in `QwenImage21DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n" + " - `QwenImage21LoopBeforeDenoiser`\n" + " - `QwenImage21LoopDenoiser`\n" + " - `QwenImage21LoopAfterDenoiser`\n" + "This block supports text-to-image generation." + ) + + +# auto_docstring +class QwenImage21ImageConditionedDenoiseStep(QwenImage21DenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImage21DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImage21ImageConditionedLoopBeforeDenoiser` + - `QwenImage21LoopDenoiser` + - `QwenImage21LoopAfterDenoiser` + This block supports image-conditioned generation. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImage21Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + 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 it does not + reproduce the same image bit-for-bit in reduced precision. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + image_latents (`Tensor`): + Packed condition image latents. Can be generated in the additional inputs step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + Per-sample latent grids, condition images first and the target image last. Can be generated in + prepare_rope_inputs step. + img_mask (`Tensor`): + Joint vision mask over the prompt positions and the target image slots. Can be generated in + prepare_rope_inputs step. + negative_img_mask (`Tensor`, *optional*): + Joint vision mask for the negative prompt. Can be generated in prepare_rope_inputs step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = [ + QwenImage21ImageConditionedLoopBeforeDenoiser, + QwenImage21LoopDenoiser, + QwenImage21LoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents. \n" + "Its loop logic is defined in `QwenImage21DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n" + " - `QwenImage21ImageConditionedLoopBeforeDenoiser`\n" + " - `QwenImage21LoopDenoiser`\n" + " - `QwenImage21LoopAfterDenoiser`\n" + "This block supports image-conditioned generation." + ) diff --git a/src/diffusers/modular_pipelines/qwenimage21/encoders.py b/src/diffusers/modular_pipelines/qwenimage21/encoders.py new file mode 100644 index 000000000000..1d305c94de4d --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/encoders.py @@ -0,0 +1,711 @@ +# Copyright 2026 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 + +import numpy as np +import PIL +import torch +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...image_processor import VaeImageProcessor, is_valid_image, is_valid_image_imagelist +from ...models import AutoencoderKLQwenImage21 +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import QwenImage21ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +QWENIMAGE21_SYSTEM_PROMPT = "Comprehend and analyze the provided prompt." +QWENIMAGE21_IMAGE_TEMPLATE = "<|vision_start|><|image_pad|><|vision_end|>" + + +# 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 + + +def get_qwenimage21_prompt_embeds( + text_encoder: Qwen3VLForConditionalGeneration, + processor: Qwen3VLProcessor, + prompt: str | list[str], + image: list[PIL.Image.Image] | None = None, + device: torch.device | None = None, +): + 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] + + # The prompt is built as a raw template string and passed straight to `processor(text=..., images=...)` rather + # than going through `apply_chat_template`: the two tokenize differently and the checkpoint expects this one. The + # number of leading system-role tokens to drop is derived from the tokenized system message so it tracks the + # processor's template. + sys_message = [{"role": "system", "content": [{"type": "text", "text": QWENIMAGE21_SYSTEM_PROMPT}]}] + drop_idx = len(processor.apply_chat_template(sys_message, tokenize=True, return_dict=False)[0]) + img_token_id = processor.tokenizer.encode("<|image_pad|>")[0] + + image_prompt = "" + condition_images = None + if image is not None: + image_prompt = " ".join(QWENIMAGE21_IMAGE_TEMPLATE.format(i + 1) for i in range(len(image))) + # 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. + condition_images = [] + for _ in prompt: + for img in image: + 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 = PIL.Image.new("RGB", img.size, (255, 255, 255)) + white.paste(img, mask=img.getchannel("A")) + img = white + condition_images.append(img) + template = ( + f"<|im_start|>system\n{QWENIMAGE21_SYSTEM_PROMPT}<|im_end|>\n" + f"<|im_start|>user\n{image_prompt}{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + txt = [template.format(p) for p in prompt] + + # Left padding, as the checkpoint was trained with. The padding is dropped either way, but the side decides the + # positions the encoder sees for a batch of prompts of different lengths. + processor_kwargs = {"text": txt, "padding": True, "padding_side": "left", "return_tensors": "pt"} + if condition_images is not None: + processor_kwargs["images"] = condition_images + model_inputs = processor(**processor_kwargs).to(device) + + forward_kwargs = { + "input_ids": model_inputs.input_ids, + "attention_mask": model_inputs.attention_mask, + "output_hidden_states": True, + } + if condition_images is not None 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. From transformers 5.0 that entry comes back normalized, so a forward hook + # returning the module's input neutralizes the norm for this call on either version. + text_model = getattr(text_encoder.model, "language_model", text_encoder.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + outputs = text_encoder(**forward_kwargs) + finally: + handle.remove() + hidden_states = outputs.hidden_states[-1] + + attention_mask = model_inputs.attention_mask + split_hidden_states = torch.split(hidden_states[attention_mask.bool()], attention_mask.sum(dim=1).tolist(), dim=0) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + + image_pad_mask = [ + (sample_ids[sample_mask.bool()] == img_token_id)[drop_idx:] + for sample_ids, sample_mask in zip(model_inputs.input_ids, attention_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 + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_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.modular_pipelines.qwenimage.encoders.encode_vae_image with AutoencoderKLQwenImage->AutoencoderKLQwenImage21 +def encode_vae_image( + image: torch.Tensor, + vae: AutoencoderKLQwenImage21, + generator: torch.Generator, + device: torch.device, + dtype: torch.dtype, + latent_channels: int = 16, + sample_mode: str = "argmax", +): + if not isinstance(image, torch.Tensor): + raise ValueError(f"Expected image to be a tensor, got {type(image)}.") + + # preprocessed image should be a 4D tensor: batch_size, num_channels, height, width + if image.dim() == 4: + image = image.unsqueeze(2) + elif image.dim() != 5: + raise ValueError(f"Expected image dims 4 or 5, got {image.dim()}.") + + image = image.to(device=device, dtype=dtype) + + if isinstance(generator, list): + image_latents = [ + retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode) + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + image_latents = (image_latents - latents_mean) / latents_std + + return image_latents + + +# auto_docstring +class QwenImage21ResizeStep(ModularPipelineBlocks): + """ + Resize condition images for Qwen-Image 2.1. Each image is converted to RGBA and resized to the `output_resolution` + target area while keeping its aspect ratio. The same resized images feed both the vision-language text encoder and + the VAE encoder. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + + Outputs: + resized_image (`list`): + RGBA condition images resized to the `output_resolution` target area + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Resize condition images for Qwen-Image 2.1. Each image is converted to RGBA and resized to the " + "`output_resolution` target area while keeping its aspect ratio. The same resized images feed both the " + "vision-language text encoder and the VAE encoder." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image"), + InputParam( + name="output_resolution", + type_hint=int, + default=1024, + description="Target side length used to derive the output size and to resize condition images.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="resized_image", + type_hint=list[PIL.Image.Image], + description="RGBA condition images resized to the `output_resolution` target area", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + images = block_state.image + if not is_valid_image_imagelist(images): + raise ValueError(f"Images must be image or list of images but are {type(images)}") + if is_valid_image(images): + images = [images] + + # 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 sees one type. + resized_images = [] + for image in images: + if isinstance(image, np.ndarray): + image = PIL.Image.fromarray(image) + elif not isinstance(image, PIL.Image.Image): + raise ValueError( + f"`image` accepts a PIL image or a numpy array, or a list of either, but got " + f"{type(image).__name__}. Latents cannot stand in for a condition image here, because the text " + f"encoder has to see the image itself." + ) + if image.mode != "RGBA": + image = image.convert("RGBA") + image_width, image_height = image.size + width, height, _ = calculate_dimensions( + block_state.output_resolution * block_state.output_resolution, image_width / image_height + ) + resized_images.append(components.image_processor.resize(image, height=height, width=width)) + + block_state.resized_image = resized_images + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21TextEncoderStep(ModularPipelineBlocks): + """ + Text Encoder step that generates text embeddings with Qwen3-VL to guide text-to-image generation. + + Components: + text_encoder (`Qwen3VLForConditionalGeneration`) processor (`Qwen3VLProcessor`) guider + (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings + negative_image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the negative prompt embeddings + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings with Qwen3-VL to guide text-to-image generation." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration), + ComponentSpec("processor", Qwen3VLProcessor), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), + OutputParam( + name="image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the prompt embeddings", + ), + OutputParam( + name="negative_image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the negative prompt embeddings", + ), + ] + + @staticmethod + def check_inputs(prompt, negative_prompt): + if not isinstance(prompt, str) and not isinstance(prompt, list): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if ( + negative_prompt is not None + and not isinstance(negative_prompt, str) + and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + self.check_inputs(block_state.prompt, block_state.negative_prompt) + + device = components._execution_device + + block_state.prompt_embeds, block_state.prompt_embeds_mask, block_state.image_pad_mask = ( + get_qwenimage21_prompt_embeds( + components.text_encoder, + components.processor, + prompt=block_state.prompt, + device=device, + ) + ) + + block_state.negative_prompt_embeds = None + block_state.negative_prompt_embeds_mask = None + block_state.negative_image_pad_mask = None + if components.requires_unconditional_embeds: + negative_prompt = block_state.negative_prompt or "" + ( + block_state.negative_prompt_embeds, + block_state.negative_prompt_embeds_mask, + block_state.negative_image_pad_mask, + ) = get_qwenimage21_prompt_embeds( + components.text_encoder, + components.processor, + prompt=negative_prompt, + device=device, + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21VLTextEncoderStep(ModularPipelineBlocks): + """ + Text Encoder step that encodes the prompt together with the condition images with Qwen3-VL to guide + image-conditioned generation. + + Components: + text_encoder (`Qwen3VLForConditionalGeneration`) processor (`Qwen3VLProcessor`) guider + (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + resized_image (`list`): + RGBA condition images resized to the output resolution. Can be generated in the resize step. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings + negative_image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the negative prompt embeddings + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Text Encoder step that encodes the prompt together with the condition images with Qwen3-VL to guide " + "image-conditioned generation." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration), + ComponentSpec("processor", Qwen3VLProcessor), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + InputParam( + name="resized_image", + required=True, + type_hint=list[PIL.Image.Image], + description="RGBA condition images resized to the output resolution. Can be generated in the resize step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), + OutputParam( + name="image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the prompt embeddings", + ), + OutputParam( + name="negative_image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the negative prompt embeddings", + ), + ] + + @staticmethod + def check_inputs(prompt, negative_prompt): + if not isinstance(prompt, str) and not isinstance(prompt, list): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if ( + negative_prompt is not None + and not isinstance(negative_prompt, str) + and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + self.check_inputs(block_state.prompt, block_state.negative_prompt) + + device = components._execution_device + + block_state.prompt_embeds, block_state.prompt_embeds_mask, block_state.image_pad_mask = ( + get_qwenimage21_prompt_embeds( + components.text_encoder, + components.processor, + prompt=block_state.prompt, + image=block_state.resized_image, + device=device, + ) + ) + + block_state.negative_prompt_embeds = None + block_state.negative_prompt_embeds_mask = None + block_state.negative_image_pad_mask = None + if components.requires_unconditional_embeds: + negative_prompt = block_state.negative_prompt or "" + ( + block_state.negative_prompt_embeds, + block_state.negative_prompt_embeds_mask, + block_state.negative_image_pad_mask, + ) = get_qwenimage21_prompt_embeds( + components.text_encoder, + components.processor, + prompt=negative_prompt, + image=block_state.resized_image, + device=device, + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21ProcessImagesInputStep(ModularPipelineBlocks): + """ + Image preprocess step that turns the resized RGBA condition images into normalized tensors for the VAE. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + resized_image (`list`): + RGBA condition images resized to the output resolution. Can be generated in the resize step. + + Outputs: + processed_image (`list`): + Normalized RGBA image tensors, one per condition image + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Image preprocess step that turns the resized RGBA condition images into normalized tensors for the VAE." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="resized_image", + required=True, + type_hint=list[PIL.Image.Image], + description="RGBA condition images resized to the output resolution. Can be generated in the resize step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_image", + type_hint=list[torch.Tensor], + description="Normalized RGBA image tensors, one per condition image", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + block_state.processed_image = [ + components.image_processor.preprocess(image, height=image.height, width=image.width) + for image in block_state.resized_image + ] + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21VaeEncoderStep(ModularPipelineBlocks): + """ + VAE Encoder step that encodes each processed condition image into normalized latents. Images can have different + resolutions, so the latents are returned as a list. + + Components: + vae (`AutoencoderKLQwenImage21`) + + Inputs: + processed_image (`list`): + Normalized RGBA image tensors to encode. Can be generated in the preprocess step. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`list`): + Normalized latents of each condition image, each of shape (1, C, 1, H, W) + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "VAE Encoder step that encodes each processed condition image into normalized latents. Images can have " + "different resolutions, so the latents are returned as a list." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLQwenImage21)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_image", + required=True, + type_hint=list[torch.Tensor], + description="Normalized RGBA image tensors to encode. Can be generated in the preprocess step.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="image_latents", + type_hint=list[torch.Tensor], + description="Normalized latents of each condition image, each of shape (1, C, 1, H, W)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + device = components._execution_device + dtype = components.vae.dtype + + block_state.image_latents = [ + encode_vae_image( + image=image, + vae=components.vae, + generator=block_state.generator, + device=device, + dtype=dtype, + latent_channels=components.vae.config.z_dim, + ) + for image in block_state.processed_image + ] + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/qwenimage21/inputs.py b/src/diffusers/modular_pipelines/qwenimage21/inputs.py new file mode 100644 index 000000000000..c8efe42d2ce7 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/inputs.py @@ -0,0 +1,376 @@ +# Copyright 2026 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 + +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam +from .modular_pipeline import QwenImage21ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def pack_latents(latents: torch.Tensor) -> torch.Tensor: + # Qwen-Image 2.1 consumes latents unpatched, so packing is a plain spatial flatten: + # (batch_size, channels, 1, height, width) -> (batch_size, height * width, channels) + batch_size, channels, _, height, width = latents.shape + return latents.view(batch_size, channels, height * width).transpose(1, 2) + + +# Copied from diffusers.modular_pipelines.qwenimage.inputs.repeat_tensor_to_batch_size +def repeat_tensor_to_batch_size( + input_name: str, + input_tensor: torch.Tensor, + batch_size: int, + num_images_per_prompt: int = 1, +) -> torch.Tensor: + """Repeat tensor elements to match the final batch size. + + This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt) + by repeating each element along dimension 0. + + The input tensor must have batch size 1 or batch_size. The function will: + - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times + - If batch size equals batch_size: repeat each element num_images_per_prompt times + + Args: + input_name (str): Name of the input tensor (used for error messages) + input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size. + batch_size (int): The base batch size (number of prompts) + num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1. + + Returns: + torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt) + + Raises: + ValueError: If input_tensor is not a torch.Tensor or has invalid batch size + + Examples: + tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor, + batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape: + [4, 3] + + tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image", + tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]]) + - shape: [4, 3] + """ + # make sure input is a tensor + if not isinstance(input_tensor, torch.Tensor): + raise ValueError(f"`{input_name}` must be a tensor") + + # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts + if input_tensor.shape[0] == 1: + repeat_by = batch_size * num_images_per_prompt + elif input_tensor.shape[0] == batch_size: + repeat_by = num_images_per_prompt + else: + raise ValueError(f"`{input_name}` must have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}") + + # expand the tensor to match the batch_size * num_images_per_prompt + input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0) + + return input_tensor + + +# auto_docstring +class QwenImage21TextInputsStep(ModularPipelineBlocks): + """ + Text input processing step that standardizes text embeddings for the pipeline. + This step: + 1. Determines `batch_size` and `dtype` based on `prompt_embeds` + 2. Expands all text embeddings and masks to batch_size * num_images_per_prompt + 3. Drops an attention mask that carries no padding, as the attention backends then skip masking + + This block should be placed after all encoder steps to process the text embeddings before they are used in + subsequent pipeline steps. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings. Can be generated from text_encoder step. + negative_image_pad_mask (`Tensor`, *optional*): + Bool mask marking the vision positions of the negative prompt embeddings. Can be generated from + text_encoder step. + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded, None when nothing is padded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded, None when nothing is padded) + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings (batch-expanded) + negative_image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the negative prompt embeddings (batch-expanded) + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Text input processing step that standardizes text embeddings for the pipeline.\n" + "This step:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Expands all text embeddings and masks to batch_size * num_images_per_prompt\n" + " 3. Drops an attention mask that carries no padding, as the attention backends then skip masking\n\n" + "This block should be placed after all encoder steps to process the text embeddings before they are used " + "in subsequent pipeline steps." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt"), + InputParam.template("prompt_embeds"), + InputParam.template("prompt_embeds_mask"), + InputParam.template("negative_prompt_embeds"), + InputParam.template("negative_prompt_embeds_mask"), + InputParam( + name="image_pad_mask", + required=True, + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the prompt embeddings. Can be generated from text_encoder step.", + ), + InputParam( + name="negative_image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the negative prompt embeddings. Can be generated from text_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="batch_size", type_hint=int, description="The batch size of the prompt embeddings"), + OutputParam(name="dtype", type_hint=torch.dtype, description="The data type of the prompt embeddings"), + OutputParam.template("prompt_embeds", note="batch-expanded"), + OutputParam.template("prompt_embeds_mask", note="batch-expanded, None when nothing is padded"), + OutputParam.template("negative_prompt_embeds", note="batch-expanded"), + OutputParam.template("negative_prompt_embeds_mask", note="batch-expanded, None when nothing is padded"), + OutputParam( + name="image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the prompt embeddings (batch-expanded)", + ), + OutputParam( + name="negative_image_pad_mask", + type_hint=torch.Tensor, + description="Bool mask marking the vision positions of the negative prompt embeddings (batch-expanded)", + ), + ] + + @staticmethod + def check_inputs( + prompt_embeds, + prompt_embeds_mask, + negative_prompt_embeds, + negative_prompt_embeds_mask, + image_pad_mask, + negative_image_pad_mask, + ): + if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: + raise ValueError("`negative_prompt_embeds_mask` is required when `negative_prompt_embeds` is not None") + if negative_prompt_embeds is None and negative_prompt_embeds_mask is not None: + raise ValueError("cannot pass `negative_prompt_embeds_mask` without `negative_prompt_embeds`") + if negative_prompt_embeds is not None and negative_image_pad_mask is None: + raise ValueError("`negative_image_pad_mask` is required when `negative_prompt_embeds` is not None") + + batch_size = prompt_embeds.shape[0] + for name, value in ( + ("prompt_embeds_mask", prompt_embeds_mask), + ("image_pad_mask", image_pad_mask), + ("negative_prompt_embeds", negative_prompt_embeds), + ("negative_prompt_embeds_mask", negative_prompt_embeds_mask), + ("negative_image_pad_mask", negative_image_pad_mask), + ): + if value is not None and value.shape[0] != batch_size: + raise ValueError(f"`{name}` must have the same batch size as `prompt_embeds`") + + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + self.check_inputs( + prompt_embeds=block_state.prompt_embeds, + prompt_embeds_mask=block_state.prompt_embeds_mask, + negative_prompt_embeds=block_state.negative_prompt_embeds, + negative_prompt_embeds_mask=block_state.negative_prompt_embeds_mask, + image_pad_mask=block_state.image_pad_mask, + negative_image_pad_mask=block_state.negative_image_pad_mask, + ) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + + for name in ( + "prompt_embeds", + "prompt_embeds_mask", + "negative_prompt_embeds", + "negative_prompt_embeds_mask", + "image_pad_mask", + "negative_image_pad_mask", + ): + value = getattr(block_state, name) + if value is not None: + value = value.repeat_interleave(block_state.num_images_per_prompt, dim=0) + setattr(block_state, name, value) + + # Without padding there is nothing to mask, and a mask that carries no information costs the attention + # backends that reject one outright. + if block_state.prompt_embeds_mask.all(): + block_state.prompt_embeds_mask = None + if block_state.negative_prompt_embeds_mask is not None and block_state.negative_prompt_embeds_mask.all(): + block_state.negative_prompt_embeds_mask = None + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21AdditionalInputsStep(ModularPipelineBlocks): + """ + Input processing step for image-conditioned generation that: + 1. Records the pixel height/width of each condition image from its latents + 2. Packs each image latent, concatenates them along the sequence dimension and expands the batch + 3. Defaults `height`/`width` to the size of the last condition image + + This block should be placed after the encoder steps and the text input step. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`list`): + Normalized latents of each condition image. Can be generated from vae_encoder step. + + Outputs: + image_height (`list`): + The pixel heights of the condition images, calculated from the image latents + image_width (`list`): + The pixel widths of the condition images, calculated from the image latents + height (`int`): + if not provided, updated to the last image height + width (`int`): + if not provided, updated to the last image width + image_latents (`Tensor`): + Condition image latents packed, concatenated along the sequence dimension and batch-expanded + """ + + model_name = "qwenimage21" + + @property + def description(self) -> str: + return ( + "Input processing step for image-conditioned generation that:\n" + " 1. Records the pixel height/width of each condition image from its latents\n" + " 2. Packs each image latent, concatenates them along the sequence dimension and expands the batch\n" + " 3. Defaults `height`/`width` to the size of the last condition image\n\n" + "This block should be placed after the encoder steps and the text input step." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt"), + InputParam.template("batch_size"), + InputParam.template("height"), + InputParam.template("width"), + InputParam( + name="image_latents", + required=True, + type_hint=list[torch.Tensor], + description="Normalized latents of each condition image. Can be generated from vae_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="image_height", + type_hint=list[int], + description="The pixel heights of the condition images, calculated from the image latents", + ), + OutputParam( + name="image_width", + type_hint=list[int], + description="The pixel widths of the condition images, calculated from the image latents", + ), + OutputParam(name="height", type_hint=int, description="if not provided, updated to the last image height"), + OutputParam(name="width", type_hint=int, description="if not provided, updated to the last image width"), + OutputParam( + name="image_latents", + type_hint=torch.Tensor, + description="Condition image latents packed, concatenated along the sequence dimension and batch-expanded", + ), + ] + + def __call__( + self, components: QwenImage21ModularPipeline, state: PipelineState + ) -> tuple[QwenImage21ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + + image_latents = block_state.image_latents + if not isinstance(image_latents, list): + image_latents = [image_latents] + + image_heights = [] + image_widths = [] + packed_image_latents = [] + for i, image_latent in enumerate(image_latents): + latent_height, latent_width = image_latent.shape[-2:] + image_heights.append(latent_height * components.vae_scale_factor) + image_widths.append(latent_width * components.vae_scale_factor) + packed_image_latents.append( + repeat_tensor_to_batch_size( + input_name=f"image_latents[{i}]", + input_tensor=pack_latents(image_latent), + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + ) + + block_state.image_latents = torch.cat(packed_image_latents, dim=1) + block_state.image_height = image_heights + block_state.image_width = image_widths + block_state.height = block_state.height or image_heights[-1] + block_state.width = block_state.width or image_widths[-1] + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py new file mode 100644 index 000000000000..d2db9e80a1b6 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py @@ -0,0 +1,586 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + QwenImage21ImageConditionedRoPEInputsStep, + QwenImage21PrepareLatentsStep, + QwenImage21RoPEInputsStep, + QwenImage21SetTimestepsStep, +) +from .decoders import QwenImage21DecodeStep, QwenImage21UnpackLatentsStep +from .denoise import QwenImage21DenoiseStep, QwenImage21ImageConditionedDenoiseStep +from .encoders import ( + QwenImage21ProcessImagesInputStep, + QwenImage21ResizeStep, + QwenImage21TextEncoderStep, + QwenImage21VaeEncoderStep, + QwenImage21VLTextEncoderStep, +) +from .inputs import QwenImage21AdditionalInputsStep, QwenImage21TextInputsStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. TEXT ENCODER +# ==================== + +QwenImage21VLEncoderBlocks = InsertableDict( + [ + ("resize", QwenImage21ResizeStep()), + ("encode", QwenImage21VLTextEncoderStep()), + ] +) + + +# auto_docstring +class QwenImage21VLEncoderStep(SequentialPipelineBlocks): + """ + Vision-language encoder step that resizes the condition images and encodes them together with the prompt. + + Components: + image_processor (`VaeImageProcessor`) text_encoder (`Qwen3VLForConditionalGeneration`) processor + (`Qwen3VLProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + + Outputs: + resized_image (`list`): + RGBA condition images resized to the `output_resolution` target area + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings + negative_image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the negative prompt embeddings + """ + + model_name = "qwenimage21" + block_classes = QwenImage21VLEncoderBlocks.values() + block_names = QwenImage21VLEncoderBlocks.keys() + + @property + def description(self) -> str: + return ( + "Vision-language encoder step that resizes the condition images and encodes them together with the prompt." + ) + + +# auto_docstring +class QwenImage21AutoTextEncoderStep(AutoPipelineBlocks): + """ + Text encoder step that encodes the prompt, together with the condition images when there are any. + This is an auto pipeline block that works for text-to-image and image-conditioned generation. + - `QwenImage21VLEncoderStep` is used when `image` is provided. + - `QwenImage21TextEncoderStep` is used otherwise. + + Components: + image_processor (`VaeImageProcessor`) text_encoder (`Qwen3VLForConditionalGeneration`) processor + (`Qwen3VLProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + + Outputs: + resized_image (`list`): + RGBA condition images resized to the `output_resolution` target area + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings + negative_image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the negative prompt embeddings + """ + + model_name = "qwenimage21" + block_classes = [QwenImage21VLEncoderStep, QwenImage21TextEncoderStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image", None] + + @property + def description(self) -> str: + return ( + "Text encoder step that encodes the prompt, together with the condition images when there are any.\n" + "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n" + " - `QwenImage21VLEncoderStep` is used when `image` is provided.\n" + " - `QwenImage21TextEncoderStep` is used otherwise." + ) + + +# ==================== +# 2. VAE ENCODER +# ==================== + +QwenImage21VaeEncoderBlocks = InsertableDict( + [ + ("resize", QwenImage21ResizeStep()), + ("preprocess", QwenImage21ProcessImagesInputStep()), + ("encode", QwenImage21VaeEncoderStep()), + ] +) + + +# auto_docstring +class QwenImage21VaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that resizes, preprocesses and encodes the condition images into latents. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage21`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + resized_image (`list`): + RGBA condition images resized to the `output_resolution` target area + processed_image (`list`): + Normalized RGBA image tensors, one per condition image + image_latents (`list`): + Normalized latents of each condition image, each of shape (1, C, 1, H, W) + """ + + model_name = "qwenimage21" + block_classes = QwenImage21VaeEncoderBlocks.values() + block_names = QwenImage21VaeEncoderBlocks.keys() + + @property + def description(self) -> str: + return "VAE encoder step that resizes, preprocesses and encodes the condition images into latents." + + +# auto_docstring +class QwenImage21AutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the condition images into their latent representations. + This is an auto pipeline block that works for image-conditioned generation. + - `QwenImage21VaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage21`) + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + resized_image (`list`): + RGBA condition images resized to the `output_resolution` target area + processed_image (`list`): + Normalized RGBA image tensors, one per condition image + image_latents (`list`): + Normalized latents of each condition image, each of shape (1, C, 1, H, W) + """ + + model_name = "qwenimage21" + block_classes = [QwenImage21VaeEncoderSequentialStep] + block_names = ["image_conditioned"] + block_trigger_inputs = ["image"] + + @property + def description(self) -> str: + return ( + "VAE encoder step that encodes the condition images into their latent representations.\n" + "This is an auto pipeline block that works for image-conditioned generation.\n" + " - `QwenImage21VaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +# ==================== +# 3. DENOISE +# ==================== + +QwenImage21CoreDenoiseBlocks = InsertableDict( + [ + ("input", QwenImage21TextInputsStep()), + ("prepare_latents", QwenImage21PrepareLatentsStep()), + ("set_timesteps", QwenImage21SetTimestepsStep()), + ("prepare_rope_inputs", QwenImage21RoPEInputsStep()), + ("denoise", QwenImage21DenoiseStep()), + ("unpack_latents", QwenImage21UnpackLatentsStep()), + ] +) + + +# auto_docstring +class QwenImage21CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for text-to-image generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`QwenImage21Transformer2DModel`) + + Configs: + sample_sigmas (default: None): Default sampling grid of the checkpoint, used when `sigmas` is not passed. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings. Can be generated from text_encoder step. + negative_image_pad_mask (`Tensor`, *optional*): + Bool mask marking the vision positions of the negative prompt embeddings. Can be generated from + text_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + 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 it does not + reproduce the same image bit-for-bit in reduced precision. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = QwenImage21CoreDenoiseBlocks.values() + block_names = QwenImage21CoreDenoiseBlocks.keys() + + @property + def description(self) -> str: + return "Core denoise step that performs the denoising process for text-to-image generation." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +QwenImage21ImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", QwenImage21TextInputsStep()), + ("additional_inputs", QwenImage21AdditionalInputsStep()), + ("prepare_latents", QwenImage21PrepareLatentsStep()), + ("set_timesteps", QwenImage21SetTimestepsStep()), + ("prepare_rope_inputs", QwenImage21ImageConditionedRoPEInputsStep()), + ("denoise", QwenImage21ImageConditionedDenoiseStep()), + ("unpack_latents", QwenImage21UnpackLatentsStep()), + ] +) + + +# auto_docstring +class QwenImage21ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for image-conditioned generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`QwenImage21Transformer2DModel`) + + Configs: + sample_sigmas (default: None): Default sampling grid of the checkpoint, used when `sigmas` is not passed. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings. Can be generated from text_encoder step. + negative_image_pad_mask (`Tensor`, *optional*): + Bool mask marking the vision positions of the negative prompt embeddings. Can be generated from + text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`list`): + Normalized latents of each condition image. Can be generated from vae_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + 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 it does not + reproduce the same image bit-for-bit in reduced precision. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = QwenImage21ImageConditionedCoreDenoiseBlocks.values() + block_names = QwenImage21ImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self) -> str: + return "Core denoise step that performs the denoising process for image-conditioned generation." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# auto_docstring +class QwenImage21AutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto core denoise step that performs the denoising process. + This is an auto pipeline block that works for text-to-image and image-conditioned generation. + - `QwenImage21ImageConditionedCoreDenoiseStep` is used when `image_latents` is provided. + - `QwenImage21CoreDenoiseStep` is used otherwise. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`QwenImage21Transformer2DModel`) + + Configs: + sample_sigmas (default: None): Default sampling grid of the checkpoint, used when `sigmas` is not passed. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + image_pad_mask (`Tensor`): + Bool mask marking the vision positions of the prompt embeddings. Can be generated from text_encoder step. + negative_image_pad_mask (`Tensor`, *optional*): + Bool mask marking the vision positions of the negative prompt embeddings. Can be generated from + text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`list`, *optional*): + Normalized latents of each condition image. Can be generated from vae_encoder step. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + 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 it does not + reproduce the same image bit-for-bit in reduced precision. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = [QwenImage21ImageConditionedCoreDenoiseStep, QwenImage21CoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self) -> str: + return ( + "Auto core denoise step that performs the denoising process.\n" + "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n" + " - `QwenImage21ImageConditionedCoreDenoiseStep` is used when `image_latents` is provided.\n" + " - `QwenImage21CoreDenoiseStep` is used otherwise." + ) + + +# ==================== +# 4. AUTO BLOCKS +# ==================== + +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImage21AutoTextEncoderStep()), + ("vae_encoder", QwenImage21AutoVaeEncoderStep()), + ("denoise", QwenImage21AutoCoreDenoiseStep()), + ("decode", QwenImage21DecodeStep()), + ] +) + + +# auto_docstring +class QwenImage21AutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image and image-conditioned generation using Qwen-Image 2.1. + - for text-to-image generation, all you need to provide is `prompt` + - for image-conditioned generation, you need to provide `prompt` and `image` (one image or a list) + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `prompt`, `image` + + Components: + image_processor (`VaeImageProcessor`) text_encoder (`Qwen3VLForConditionalGeneration`) processor + (`Qwen3VLProcessor`) guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLQwenImage21`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImage21Transformer2DModel`) + + Configs: + sample_sigmas (default: None): Default sampling grid of the checkpoint, used when `sigmas` is not passed. + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + output_resolution (`int`, *optional*, defaults to 1024): + Target side length used to derive the output size and to resize condition images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`list`, *optional*): + Normalized latents of each condition image. Can be generated from vae_encoder step. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + 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 it does not + reproduce the same image bit-for-bit in reduced precision. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage21" + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"prompt": True, "image": True}, + } + + @property + def description(self) -> str: + return ( + "Auto Modular pipeline for text-to-image and image-conditioned generation using Qwen-Image 2.1.\n" + "- for text-to-image generation, all you need to provide is `prompt`\n" + "- for image-conditioned generation, you need to provide `prompt` and `image` (one image or a list)" + ) + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/src/diffusers/modular_pipelines/qwenimage21/modular_pipeline.py b/src/diffusers/modular_pipelines/qwenimage21/modular_pipeline.py new file mode 100644 index 000000000000..8795108a4077 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/modular_pipeline.py @@ -0,0 +1,51 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...loaders import QwenImageLoraLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class QwenImage21ModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): + """ + A ModularPipeline for Qwen-Image 2.1. + """ + + default_blocks_name = "QwenImage21AutoBlocks" + + @property + def vae_scale_factor(self): + # The VAE compresses 16x spatially and the transformer consumes latents unpatched, so one token covers a + # 16x16 pixel tile. + vae_scale_factor = 16 + if hasattr(self, "vae") and self.vae is not None: + vae_scale_factor = 2 ** len(self.vae.config.temperal_downsample) + return vae_scale_factor + + @property + def num_channels_latents(self): + num_channels_latents = 64 + if hasattr(self, "transformer") and self.transformer is not None: + num_channels_latents = self.transformer.config.in_channels + return num_channels_latents + + @property + def requires_unconditional_embeds(self): + requires_unconditional_embeds = False + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + return requires_unconditional_embeds diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index a03f41412bcb..bd84ced2120e 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -115,6 +115,7 @@ QwenImageLayeredPipeline, QwenImagePipeline, ) +from .qwenimage21 import QwenImage21Pipeline from .sana import SanaPipeline from .stable_audio import StableAudioPipeline from .stable_audio_3 import StableAudio3Pipeline @@ -193,6 +194,7 @@ ("cogview4-control", CogView4ControlPipeline), ("nucleusmoe-image", NucleusMoEImagePipeline), ("qwenimage", QwenImagePipeline), + ("qwenimage21", QwenImage21Pipeline), ("qwenimage-controlnet", QwenImageControlNetPipeline), ("z-image", ZImagePipeline), ("z-image-controlnet", ZImageControlNetPipeline), diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 9ede7a7543f6..b2b567dbc9f2 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -662,6 +662,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class QwenImage21AutoBlocks(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 QwenImage21ModularPipeline(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 QwenImageAutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/modular_pipelines/qwenimage21/__init__.py b/tests/modular_pipelines/qwenimage21/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py new file mode 100644 index 000000000000..248aa0bbd6ff --- /dev/null +++ b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py @@ -0,0 +1,261 @@ +# 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 numpy as np +import PIL + +from diffusers.modular_pipelines import QwenImage21AutoBlocks, QwenImage21ModularPipeline + +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularGuiderTesterMixin, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) + + +# The tiny VAE compresses 16x and the transformer groups latents in 2x2 blocks, so 32 is the smallest resolution: +# a 2x2 latent, which is exactly one vision slot's worth of target tokens. +IMAGE_SIZE = 32 + +QWENIMAGE21_TEXT2IMAGE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "QwenImage21TextEncoderStep"), + ("denoise.input", "QwenImage21TextInputsStep"), + ("denoise.prepare_latents", "QwenImage21PrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImage21SetTimestepsStep"), + ("denoise.prepare_rope_inputs", "QwenImage21RoPEInputsStep"), + ("denoise.denoise", "QwenImage21DenoiseStep"), + ("denoise.unpack_latents", "QwenImage21UnpackLatentsStep"), + ("decode", "QwenImage21DecodeStep"), + ], +} + +QWENIMAGE21_IMAGE_CONDITIONED_WORKFLOWS = { + "image_conditioned": [ + ("text_encoder.resize", "QwenImage21ResizeStep"), + ("text_encoder.encode", "QwenImage21VLTextEncoderStep"), + ("vae_encoder.resize", "QwenImage21ResizeStep"), + ("vae_encoder.preprocess", "QwenImage21ProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImage21VaeEncoderStep"), + ("denoise.input", "QwenImage21TextInputsStep"), + ("denoise.additional_inputs", "QwenImage21AdditionalInputsStep"), + ("denoise.prepare_latents", "QwenImage21PrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImage21SetTimestepsStep"), + ("denoise.prepare_rope_inputs", "QwenImage21ImageConditionedRoPEInputsStep"), + ("denoise.denoise", "QwenImage21ImageConditionedDenoiseStep"), + ("denoise.unpack_latents", "QwenImage21UnpackLatentsStep"), + ("decode", "QwenImage21DecodeStep"), + ], +} + + +QWENIMAGE21_WORKFLOW_DEFAULTS = { + "text2image": { + "components": { + "text_encoder": "Qwen3VLForConditionalGeneration", + "processor": "Qwen3VLProcessor", + "guider": "ClassifierFreeGuidance", + "scheduler": "FlowMatchEulerDiscreteScheduler", + "transformer": "QwenImage21Transformer2DModel", + "vae": "AutoencoderKLQwenImage21", + "image_processor": "VaeImageProcessor", + }, + "configs": {"sample_sigmas": None}, + "required_inputs": ["prompt"], + "inputs": { + "negative_prompt": None, + "num_images_per_prompt": 1, + "latents": None, + "height": None, + "width": None, + "output_resolution": 1024, + "generator": None, + "num_inference_steps": 40, + "sigmas": None, + "use_kv_cache": True, + "attention_kwargs": None, + "output_type": "pil", + }, + "component_configs": {"guider": {"guidance_scale": 1.0}, "image_processor": {"vae_scale_factor": 16}}, + }, + "image_conditioned": { + "components": { + "image_processor": "VaeImageProcessor", + "text_encoder": "Qwen3VLForConditionalGeneration", + "processor": "Qwen3VLProcessor", + "guider": "ClassifierFreeGuidance", + "vae": "AutoencoderKLQwenImage21", + "scheduler": "FlowMatchEulerDiscreteScheduler", + "transformer": "QwenImage21Transformer2DModel", + }, + "configs": {"sample_sigmas": None}, + "required_inputs": ["image", "prompt"], + "inputs": { + "output_resolution": 1024, + "negative_prompt": None, + "generator": None, + "num_images_per_prompt": 1, + "height": None, + "width": None, + "latents": None, + "num_inference_steps": 40, + "sigmas": None, + "use_kv_cache": True, + "attention_kwargs": None, + "output_type": "pil", + }, + "component_configs": {"image_processor": {"vae_scale_factor": 16}, "guider": {"guidance_scale": 1.0}}, + }, +} + + +def get_dummy_condition_image(seed=0): + array = np.random.RandomState(seed).randint(0, 255, (IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8) + return PIL.Image.fromarray(array).convert("RGBA") + + +class QwenImage21ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = QwenImage21ModularPipeline + pipeline_blocks_class = QwenImage21AutoBlocks + pretrained_model_name_or_path = "akshan-main/tiny-qwenimage21-modular-pipe" + params = frozenset( + ["prompt", "negative_prompt", "height", "width", "output_resolution", "attention_kwargs", "image"] + ) + batch_params = frozenset(["prompt", "negative_prompt"]) + expected_workflow_blocks = QWENIMAGE21_TEXT2IMAGE_WORKFLOWS + expected_workflow_defaults = {"text2image": QWENIMAGE21_WORKFLOW_DEFAULTS["text2image"]} + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "dance monkey", + "negative_prompt": "bad quality", + "generator": self.get_generator(seed), + "num_inference_steps": 2, + "height": IMAGE_SIZE, + "width": IMAGE_SIZE, + "output_type": "pt", + } + + +class TestQwenImage21ModularPipelineFast(QwenImage21ModularPipelineTesterConfig, ModularPipelineTesterMixin): + def test_kv_cache_matches_uncached_denoising(self): + # The cache only stores the step-independent text keys and values, so with and without it the loop has + # to land on the same image. + pipe = self.get_pipeline().to(torch_device) + cached = pipe(**self.get_dummy_inputs(), use_kv_cache=True, output="images") + uncached = pipe(**self.get_dummy_inputs(), use_kv_cache=False, output="images") + assert_tensors_close(cached, uncached, atol=1e-4, rtol=1e-4) + + def test_sample_sigmas_config_sets_the_schedule(self): + pipe = self.get_pipeline().to(torch_device) + sample_sigmas = [1.0, 0.6, 0.2] + pipe.update_components(sample_sigmas=sample_sigmas) + timesteps = pipe(**self.get_dummy_inputs(), output="timesteps") + assert len(timesteps) == len(sample_sigmas) + # Explicit sigmas take precedence over the configured grid. + timesteps = pipe(**self.get_dummy_inputs(), sigmas=[1.0, 0.5], output="timesteps") + assert len(timesteps) == 2 + + +class TestQwenImage21ModularPipelineLoading(QwenImage21ModularPipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestQwenImage21ModularPipelineWorkflow(QwenImage21ModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass + + +class TestQwenImage21ModularPipelineMemory(QwenImage21ModularPipelineTesterConfig, ModularMemoryTesterMixin): + pass + + +class TestQwenImage21ModularPipelineGuider(QwenImage21ModularPipelineTesterConfig, ModularGuiderTesterMixin): + def test_guider_cfg(self): + # The tiny model moves the output by about 1e-2 under guidance, right at the default threshold. + super().test_guider_cfg(1e-3) + + +class QwenImage21ImageConditionedModularPipelineTesterConfig(QwenImage21ModularPipelineTesterConfig): + expected_workflow_blocks = QWENIMAGE21_IMAGE_CONDITIONED_WORKFLOWS + expected_workflow_defaults = {"image_conditioned": QWENIMAGE21_WORKFLOW_DEFAULTS["image_conditioned"]} + + def get_dummy_inputs(self, seed=0): + inputs = super().get_dummy_inputs(seed) + # Keep the condition image at its own size: the resize step would otherwise blow a tiny image up to the + # default 1024x1024 area before the text encoder and the VAE. + inputs["output_resolution"] = IMAGE_SIZE + inputs["image"] = get_dummy_condition_image() + return inputs + + +class TestQwenImage21ImageConditionedModularPipelineFast( + QwenImage21ImageConditionedModularPipelineTesterConfig, ModularPipelineTesterMixin +): + def test_multiple_condition_images(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["image"] = [get_dummy_condition_image(0), get_dummy_condition_image(1)] + images = pipe(**inputs, output="images") + assert images.shape == (1, 4, IMAGE_SIZE, IMAGE_SIZE) + + def test_output_size_follows_the_last_condition_image(self): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs.pop("height") + inputs.pop("width") + inputs["output_resolution"] = 2 * IMAGE_SIZE + inputs["image"] = [ + get_dummy_condition_image(), + PIL.Image.new("RGB", (4 * IMAGE_SIZE, IMAGE_SIZE)), + ] + state = pipe(**inputs) + assert state.get("height") == IMAGE_SIZE + assert state.get("width") == 4 * IMAGE_SIZE + assert state.get("images").shape == (1, 4, IMAGE_SIZE, 4 * IMAGE_SIZE) + + def test_kv_cache_matches_uncached_denoising(self): + pipe = self.get_pipeline().to(torch_device) + cached = pipe(**self.get_dummy_inputs(), use_kv_cache=True, output="images") + uncached = pipe(**self.get_dummy_inputs(), use_kv_cache=False, output="images") + assert_tensors_close(cached, uncached, atol=1e-4, rtol=1e-4) + + +class TestQwenImage21ImageConditionedModularPipelineLoading( + QwenImage21ImageConditionedModularPipelineTesterConfig, ModularLoadingTesterMixin +): + pass + + +class TestQwenImage21ImageConditionedModularPipelineWorkflow( + QwenImage21ImageConditionedModularPipelineTesterConfig, ModularWorkflowTesterMixin +): + pass + + +class TestQwenImage21ImageConditionedModularPipelineMemory( + QwenImage21ImageConditionedModularPipelineTesterConfig, ModularMemoryTesterMixin +): + pass + + +class TestQwenImage21ImageConditionedModularPipelineGuider( + QwenImage21ImageConditionedModularPipelineTesterConfig, ModularGuiderTesterMixin +): + def test_guider_cfg(self): + super().test_guider_cfg(1e-3) diff --git a/tests/modular_pipelines/testing_utils/utils.py b/tests/modular_pipelines/testing_utils/utils.py index 8b2f1e71ac26..c29e2c5ed3b2 100644 --- a/tests/modular_pipelines/testing_utils/utils.py +++ b/tests/modular_pipelines/testing_utils/utils.py @@ -68,7 +68,7 @@ def get_specified_components(path_or_repo_id, cache_dir=None): components = set() for k, v in config.items(): - if isinstance(v, (str, int, float, bool)): + if v is None or isinstance(v, (str, int, float, bool)): continue for entry in v: if isinstance(entry, dict) and (entry.get("repo") or entry.get("pretrained_model_name_or_path")):