diff --git a/docs/source/en/api/pipelines/qwenimage21.md b/docs/source/en/api/pipelines/qwenimage21.md index 64461ed8c0d0..e48f12ef2ec9 100644 --- a/docs/source/en/api/pipelines/qwenimage21.md +++ b/docs/source/en/api/pipelines/qwenimage21.md @@ -62,6 +62,83 @@ pipe.transformer.set_attn_processor(QwenImage21FlexAttnProcessor()) pipe.transformer.compile() ``` +## Modular generation and inpainting + +[`QwenImage21ModularPipeline`] selects text-to-image, image-conditioned generation, or inpainting from the inputs. +It reuses the QwenImage21 transformer, RGBA VAE, and Qwen3-VL processor from the standard checkpoint. + +```python +import torch +from diffusers import ModularPipeline +from diffusers.utils import load_image + +pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image-2.1") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +images = pipe(prompt="A capybara wearing a wizard hat", output="images") + +source = load_image("source.png") +mask = load_image("mask.png").convert("L") +images = pipe( + prompt="Give the capybara a red scarf", + image=source, + mask_image=mask, + strength=1.0, + generator=torch.Generator("cpu").manual_seed(0), + output="images", +) +images[0].save("inpaint.png") +``` + +White mask pixels are repainted; black pixels restore the source latents after each denoising step. Preservation +is in latent space, so decoding can change unmasked pixels through VAE reconstruction. `strength` must be in +`(0, 1]` and leave at least one denoising step. Lower values start from a less noisy source image. + +For multiple-reference inpainting, keep one source in `image` and pass the extra images as `reference_images`: + +```python +images = pipe( + prompt="Use the scarf from image 2 and the fabric pattern from image 3 in the masked area of image 1", + image=source, + mask_image=mask, + reference_images=[load_image("scarf.png"), load_image("fabric.png")], + output="images", +) +``` + +The source is the first condition image, followed by references in the supplied order. All condition images feed +both Qwen3-VL and the VAE. References can have different aspect ratios; only the source has a repaint mask. The +source, mask, and references are shared across a batch of prompts. Nested per-prompt reference lists are unsupported. +Without a mask, `image` can be a single image or a flat list for image-conditioned generation. + +`height` and `width` must be multiples of 32. `output_resolution` defaults to 1024 and controls the area used to +resize condition images, independently of explicit output dimensions. Inpainting derives omitted dimensions from +the source; ordinary image-conditioned generation derives them from the last condition image. + +The default guider disables classifier-free guidance. Configure guidance through the component rather than a +pipeline argument: + +```python +from diffusers import ClassifierFreeGuidance + +pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=3.0)) +images = pipe(prompt="A capybara", negative_prompt="blurry", output="images") +``` + +`use_kv_cache=True` caches step-independent condition tokens when the transformer has `causal_condition` enabled. +Each generation and guidance branch has its own cache. As with the standard pipeline, reduced-precision cached +and uncached outputs can differ because the attention layouts differ. + +For standalone encoding or denoising, obtain a workflow with +`pipe.blocks.get_workflow("text2image")`, `"image_conditioned"`, or `"inpainting"`, and initialize a pipeline from +its individual blocks using `init_pipeline()`. Encoders return unexpanded embeddings; the denoise input block +applies `num_images_per_prompt`, so encoded features can be reused for different output counts. + +## QwenImage21ModularPipeline + +[[autodoc]] QwenImage21ModularPipeline + ## QwenImage21Pipeline [[autodoc]] QwenImage21Pipeline diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2825e9888c98..e86adad3529e 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -557,6 +557,8 @@ "MiniMaxH3ModularPipeline", "MiniMaxMusic3Blocks", "MiniMaxMusic3ModularPipeline", + "QwenImage21AutoBlocks", + "QwenImage21ModularPipeline", "QwenImageAutoBlocks", "QwenImageEditAutoBlocks", "QwenImageEditModularPipeline", @@ -1418,6 +1420,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 8c3f9ccc62c8..d188eada29ce 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -91,6 +91,7 @@ "Krea2TurboAutoBlocks", "Krea2TurboModularPipeline", ] + _import_structure["qwenimage21"] = ["QwenImage21AutoBlocks", "QwenImage21ModularPipeline"] _import_structure["qwenimage"] = [ "QwenImageAutoBlocks", "QwenImageModularPipeline", @@ -221,6 +222,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 e9e5463c1e72..575389b405f7 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -143,6 +143,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("ideogram4", _create_default_map_fn("Ideogram4ModularPipeline")), ("krea2", _krea2_map_fn), ("qwenimage", _create_default_map_fn("QwenImageModularPipeline")), + ("qwenimage21", _create_default_map_fn("QwenImage21ModularPipeline")), ("qwenimage-edit", _create_default_map_fn("QwenImageEditModularPipeline")), ("qwenimage-edit-plus", _create_default_map_fn("QwenImageEditPlusModularPipeline")), ("qwenimage-layered", _create_default_map_fn("QwenImageLayeredModularPipeline")), 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..557d86f0466d --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/before_denoise.py @@ -0,0 +1,298 @@ +# 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 numpy as np +import torch +import torch.nn.functional as F + +from ...models import QwenImage21Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class QwenImage21TextInputsStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Expand prompt embeddings and vision masks to the requested output batch." + + @property + def inputs(self): + return [ + InputParam.template("prompt_embeds", required=True), + InputParam.template("negative_prompt_embeds", required=False), + InputParam.template("prompt_embeds_mask", required=False), + InputParam.template("negative_prompt_embeds_mask", required=False), + InputParam( + "image_pad_mask", + required=True, + type_hint=torch.Tensor, + description="Positive prompt vision positions.", + ), + InputParam( + "negative_image_pad_mask", type_hint=torch.Tensor, description="Negative prompt vision positions." + ), + InputParam.template("num_images_per_prompt", default=1), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds_mask"), + OutputParam("image_pad_mask", type_hint=torch.Tensor, description="Expanded positive vision mask."), + OutputParam( + "negative_image_pad_mask", type_hint=torch.Tensor, description="Expanded negative vision mask." + ), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + if not isinstance(block_state.num_images_per_prompt, int) or block_state.num_images_per_prompt < 1: + raise ValueError("`num_images_per_prompt` must be a positive integer.") + for name in ( + "prompt_embeds", + "negative_prompt_embeds", + "prompt_embeds_mask", + "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).to( + components._execution_device + ) + if name.endswith("embeds_mask") and value.all(): + value = None + setattr(block_state, name, value) + self.set_block_state(state, block_state) + return components, state + + +class QwenImage21PrepareLatentsStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Prepare target noise and spatial metadata for unpatched QwenImage21 latents." + + @property + def expected_components(self): + return [ComponentSpec("transformer", QwenImage21Transformer2DModel)] + + @property + def inputs(self): + return [ + InputParam.template("prompt_embeds", required=True), + InputParam.template("height"), + InputParam.template("width"), + InputParam("output_resolution", default=1024, type_hint=int, description="Default output side length."), + InputParam.template("generator"), + InputParam.template("latents"), + InputParam("condition_latents", type_hint=torch.Tensor, description="Packed condition tokens."), + InputParam("condition_shapes", type_hint=list, description="Spatial shape of each condition image."), + InputParam( + "image_pad_mask", + required=True, + type_hint=torch.Tensor, + description="Positive prompt vision positions.", + ), + InputParam( + "negative_image_pad_mask", type_hint=torch.Tensor, description="Negative prompt vision positions." + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Initial target noise."), + OutputParam("height", type_hint=int, description="Output height in pixels."), + OutputParam("width", type_hint=int, description="Output width in pixels."), + OutputParam( + "condition_latents", + type_hint=torch.Tensor, + description="Condition tokens expanded to the output batch.", + ), + OutputParam( + "img_shapes", type_hint=list, description="Condition and target grid shapes for rotary embeddings." + ), + OutputParam( + "img_mask", type_hint=torch.Tensor, description="Joint positive prompt and target vision positions." + ), + OutputParam( + "negative_img_mask", + type_hint=torch.Tensor, + description="Joint negative prompt and target vision positions.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + height = block_state.height or block_state.output_resolution + width = block_state.width or block_state.output_resolution + if min(height, width) < 32 or height % 32 or width % 32: + raise ValueError("`height` and `width` must be positive multiples of 32.") + block_state.height, block_state.width = height, width + batch = block_state.prompt_embeds.shape[0] + channels = components.transformer.config.in_channels + dtype, device = block_state.prompt_embeds.dtype, components._execution_device + shape = (batch, 1, channels, height // 16, width // 16) + if block_state.latents is None: + noise = randn_tensor(shape, generator=block_state.generator, dtype=dtype, device=device) + block_state.latents = noise.view(batch, channels, -1).transpose(1, 2) + else: + if block_state.latents.shape != (batch, height * width // 256, channels): + raise ValueError("`latents` must have shape (effective batch, height * width / 256, latent channels).") + block_state.latents = block_state.latents.to(device=device, dtype=dtype) + if block_state.condition_latents is not None: + block_state.condition_latents = block_state.condition_latents.to(device=device, dtype=dtype).expand( + batch, -1, -1 + ) + block_state.img_shapes = [(block_state.condition_shapes or []) + [(1, height // 16, width // 16)]] * batch + target_slots = block_state.latents.shape[1] // 4 + block_state.img_mask = torch.cat( + [block_state.image_pad_mask, block_state.image_pad_mask.new_ones(batch, target_slots)], dim=1 + ) + block_state.negative_img_mask = None + if block_state.negative_image_pad_mask is not None: + block_state.negative_img_mask = torch.cat( + [ + block_state.negative_image_pad_mask, + block_state.negative_image_pad_mask.new_ones(batch, target_slots), + ], + dim=1, + ) + self.set_block_state(state, block_state) + return components, state + + +class QwenImage21SetTimestepsStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Set the flow-matching schedule using the target image sequence length." + + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [ + InputParam.template("num_inference_steps", default=40), + InputParam.template("sigmas"), + InputParam.template("latents", required=True), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="Denoising timesteps."), + OutputParam("num_inference_steps", type_hint=int, description="Number of denoising timesteps."), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + if block_state.num_inference_steps < 1: + raise ValueError("`num_inference_steps` must be positive.") + config = components.scheduler.config + base_len, max_len = config.get("base_image_seq_len", 256), config.get("max_image_seq_len", 4096) + base_shift, max_shift = config.get("base_shift", 0.5), config.get("max_shift", 1.15) + m = (max_shift - base_shift) / (max_len - base_len) + mu = block_state.latents.shape[1] * m + base_shift - m * base_len + sigmas = block_state.sigmas + if sigmas is None: + sigmas = np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps) + components.scheduler.set_timesteps(sigmas=sigmas, device=components._execution_device, mu=mu) + components.scheduler.set_begin_index(0) + block_state.timesteps = components.scheduler.timesteps + block_state.num_inference_steps = len(block_state.timesteps) + self.set_block_state(state, block_state) + return components, state + + +class QwenImage21PrepareInpaintStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Apply strength, retain source noise, and prepare the latent repaint mask." + + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam("source_latents", required=True, type_hint=torch.Tensor, description="Encoded source tokens."), + InputParam("processed_mask", required=True, type_hint=torch.Tensor, description="Binary repaint mask."), + InputParam("timesteps", required=True, type_hint=torch.Tensor, description="Full denoising schedule."), + InputParam.template("num_inference_steps", default=40), + InputParam.template("strength", default=1.0), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Source latents with initial noise."), + OutputParam( + "initial_noise", type_hint=torch.Tensor, description="Noise reused when preserving the source." + ), + OutputParam( + "source_latents", type_hint=torch.Tensor, description="Source tokens expanded to the output batch." + ), + OutputParam("mask", type_hint=torch.Tensor, description="Repaint weights for target latent tokens."), + OutputParam("timesteps", type_hint=torch.Tensor, description="Strength-adjusted denoising schedule."), + OutputParam("num_inference_steps", type_hint=int, description="Number of retained steps."), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + if not 0 < block_state.strength <= 1: + raise ValueError("`strength` must be in (0, 1].") + count = int(block_state.num_inference_steps * block_state.strength) + if count < 1: + raise ValueError("`strength` leaves no denoising steps; increase strength or the step count.") + start = block_state.num_inference_steps - count + block_state.timesteps = block_state.timesteps[start:] + block_state.num_inference_steps = count + components.scheduler.set_begin_index(start) + block_state.initial_noise = block_state.latents + block_state.source_latents = block_state.source_latents.to(block_state.latents).expand( + block_state.latents.shape[0], -1, -1 + ) + block_state.latents = components.scheduler.scale_noise( + block_state.source_latents, block_state.timesteps[:1], block_state.initial_noise + ) + mask = F.interpolate( + block_state.processed_mask, size=(block_state.height // 16, block_state.width // 16), mode="nearest" + ) + block_state.mask = mask.flatten(2).transpose(1, 2).to(block_state.latents) + 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..fc2aca0c2035 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/decoders.py @@ -0,0 +1,73 @@ +# 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 ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class QwenImage21DecodeStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Unpack and decode QwenImage21 target tokens to RGBA images." + + @property + def expected_components(self): + return [ + ComponentSpec("vae", AutoencoderKLQwenImage21), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("output_type", default="pil"), + ] + + @property + def intermediate_outputs(self): + return [OutputParam.template("images")] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + latents = block_state.latents + if block_state.output_type == "latent": + block_state.images = latents + else: + latents = ( + latents.transpose(1, 2) + .reshape(latents.shape[0], latents.shape[-1], 1, block_state.height // 16, block_state.width // 16) + .to(components.vae.dtype) + ) + mean = latents.new_tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1) + std = latents.new_tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1) + images = components.vae.decode(latents * std + mean, 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..506624d28d81 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/denoise.py @@ -0,0 +1,316 @@ +# 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 ...guiders import ClassifierFreeGuidance +from ...models import QwenImage21Transformer2DModel +from ...models.transformers.transformer_qwenimage21 import QwenImage21KVCache +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ..modular_pipeline import LoopSequentialPipelineBlocks, ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class QwenImage21LoopDenoiser(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Predict target flow with separate condition and guidance KV histories." + + @property + def expected_components(self): + return [ + ComponentSpec("transformer", QwenImage21Transformer2DModel), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam.template("prompt_embeds", required=True), + InputParam.template("negative_prompt_embeds", required=False), + InputParam.template("prompt_embeds_mask", required=False), + InputParam.template("negative_prompt_embeds_mask", required=False), + InputParam.template("attention_kwargs"), + InputParam( + "condition_latents", type_hint=torch.Tensor, description="Condition image tokens preceding the target." + ), + InputParam("img_shapes", required=True, type_hint=list, description="Condition and target grid shapes."), + InputParam("img_mask", required=True, type_hint=torch.Tensor, description="Joint positive vision mask."), + InputParam("negative_img_mask", type_hint=torch.Tensor, description="Joint negative vision mask."), + ] + + @property + def intermediate_outputs(self): + return [OutputParam("noise_pred", type_hint=torch.Tensor, description="Guided target flow prediction.")] + + @torch.no_grad() + def __call__(self, components, block_state, i, t): + guider_inputs = { + "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds), + "encoder_hidden_states_mask": (block_state.prompt_embeds_mask, block_state.negative_prompt_embeds_mask), + "img_mask": (block_state.img_mask, block_state.negative_img_mask), + } + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + batches = components.guider.prepare_inputs(guider_inputs) + latent_input = block_state.latents + if block_state.condition_latents is not None: + latent_input = torch.cat([block_state.condition_latents, latent_input], dim=1) + timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype) / 1000 + for batch in batches: + components.guider.prepare_models(components.transformer) + context = getattr(batch, components.guider._identifier_key) + cache, mode = None, None + if block_state.cache_enabled: + if context not in block_state.kv_caches: + block_state.kv_caches[context] = QwenImage21KVCache(len(components.transformer.transformer_blocks)) + mode = "extract" + else: + mode = "cached" + cache = block_state.kv_caches[context] + try: + with components.transformer.cache_context(context): + batch.noise_pred = components.transformer( + hidden_states=latent_input, + timestep=timestep, + img_shapes=block_state.img_shapes, + attention_kwargs=block_state.attention_kwargs, + kv_cache=cache, + kv_cache_mode=mode, + return_dict=False, + **{name: getattr(batch, name) for name in guider_inputs}, + )[0][:, -block_state.latents.shape[1] :] + finally: + components.guider.cleanup_models(components.transformer) + block_state.noise_pred = components.guider(batches)[0] + return components, block_state + + +class QwenImage21LoopStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Advance the flow-matching scheduler." + + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam("noise_pred", required=True, type_hint=torch.Tensor, description="Predicted target flow."), + ] + + @property + def intermediate_outputs(self): + return [OutputParam.template("latents")] + + @torch.no_grad() + def __call__(self, components, block_state, i, t): + dtype = block_state.latents.dtype + block_state.latents = components.scheduler.step( + block_state.noise_pred, t, block_state.latents, return_dict=False + )[0] + if torch.backends.mps.is_available(): + block_state.latents = block_state.latents.to(dtype) + return components, block_state + + +class QwenImage21LoopInpaintStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Restore unmasked source tokens at the next step's noise level." + + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam("source_latents", required=True, type_hint=torch.Tensor, description="Clean source tokens."), + InputParam("initial_noise", required=True, type_hint=torch.Tensor, description="Initial target noise."), + InputParam("mask", required=True, type_hint=torch.Tensor, description="Latent repaint mask."), + ] + + @property + def intermediate_outputs(self): + return [OutputParam.template("latents")] + + @torch.no_grad() + def __call__(self, components, block_state, i, t): + source = block_state.source_latents + if i + 1 < len(block_state.timesteps): + source = components.scheduler.scale_noise( + source, block_state.timesteps[i + 1 : i + 2], block_state.initial_noise + ) + block_state.latents = (1 - block_state.mask) * source + block_state.mask * block_state.latents + return components, block_state + + +# auto_docstring +class QwenImage21DenoiseStep(LoopSequentialPipelineBlocks): + """ + Denoise target tokens with per-call causal-condition KV caches. + + Components: + transformer (`QwenImage21Transformer2DModel`) guider (`ClassifierFreeGuidance`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + use_kv_cache (`bool`, *optional*, defaults to True): + Cache step-independent text and condition-image keys and values. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. 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. + prompt_embeds_mask (`Tensor`, *optional*): + mask for the text embeddings. 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. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + condition_latents (`Tensor`, *optional*): + Condition image tokens preceding the target. + img_shapes (`list`): + Condition and target grid shapes. + img_mask (`Tensor`): + Joint positive vision mask. + negative_img_mask (`Tensor`, *optional*): + Joint negative vision mask. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = [QwenImage21LoopDenoiser, QwenImage21LoopStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "Denoise target tokens with per-call causal-condition KV caches." + + @property + def loop_inputs(self): + return [ + InputParam("timesteps", required=True, type_hint=torch.Tensor, description="Denoising timesteps."), + InputParam.template("num_inference_steps", default=40), + InputParam( + "use_kv_cache", + default=True, + type_hint=bool, + description="Cache step-independent text and condition-image keys and values.", + ), + ] + + @property + def loop_intermediate_outputs(self): + return [ + OutputParam("kv_caches", type_hint=dict, description="Temporary KV caches, cleared after generation."), + OutputParam("cache_enabled", type_hint=bool, description="Whether causal-condition caching is active."), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.kv_caches = {} + block_state.cache_enabled = block_state.use_kv_cache and components.transformer.config.causal_condition + try: + with self.progress_bar(total=len(block_state.timesteps)) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + progress_bar.update() + finally: + block_state.kv_caches.clear() + components.guider.set_state(step=0, num_inference_steps=None, timestep=None) + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImage21InpaintDenoiseStep(QwenImage21DenoiseStep): + """ + Denoise the masked target while preserving source tokens outside the mask. + + Components: + transformer (`QwenImage21Transformer2DModel`) guider (`ClassifierFreeGuidance`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + use_kv_cache (`bool`, *optional*, defaults to True): + Cache step-independent text and condition-image keys and values. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. 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. + prompt_embeds_mask (`Tensor`, *optional*): + mask for the text embeddings. 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. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + condition_latents (`Tensor`, *optional*): + Condition image tokens preceding the target. + img_shapes (`list`): + Condition and target grid shapes. + img_mask (`Tensor`): + Joint positive vision mask. + negative_img_mask (`Tensor`, *optional*): + Joint negative vision mask. + source_latents (`Tensor`): + Clean source tokens. + initial_noise (`Tensor`): + Initial target noise. + mask (`Tensor`): + Latent repaint mask. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = [QwenImage21LoopDenoiser, QwenImage21LoopStep, QwenImage21LoopInpaintStep] + block_names = ["denoiser", "scheduler", "blend"] + + @property + def description(self): + return "Denoise the masked target while preserving source tokens outside the mask." diff --git a/src/diffusers/modular_pipelines/qwenimage21/encoders.py b/src/diffusers/modular_pipelines/qwenimage21/encoders.py new file mode 100644 index 000000000000..fa15801f1a7b --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/encoders.py @@ -0,0 +1,303 @@ +# 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 PIL import Image as PILImage +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLQwenImage21 +from ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +def get_qwenimage21_prompt_embeds(text_encoder, processor, prompt, image, device): + sys_prompt = "Comprehend and analyze the provided prompt." + prefix = f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n" + suffix = "{}<|im_end|>\n<|im_start|>assistant\n" + prompt_template_t2i = prefix + suffix + prompt_template_ti2i = prefix + "<|vision_start|><|image_pad|><|vision_end|>" + suffix + sys_message = [{"role": "system", "content": [{"type": "text", "text": sys_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] + prompt = [prompt] if isinstance(prompt, str) else prompt + # Qwen has no bos token, so an empty string leaves the encoder with nothing to read. + prompt = [" " if not p else p for p in prompt] + is_t2i = image is None + + if is_t2i: + prompts = [prompt_template_t2i.format(t) for t in prompt] + else: + prompts = [] + condition_pil_list = [] + for t in prompt: + n_imgs = len(image) + replace = "<|vision_start|><|image_pad|><|vision_end|>" + for i in range(2, n_imgs + 1): + replace += f" <|vision_start|><|image_pad|><|vision_end|>" + template = prompt_template_ti2i.replace("<|vision_start|><|image_pad|><|vision_end|>", replace) + prompts.append(template.format(t)) + # Each prompt's template repeats the `<|image_pad|>` placeholders, so hand the processor one set of + # images per prompt, in the order the placeholders appear. + for _ in prompt: + for img in image: + if not isinstance(img, PILImage.Image): + img = PILImage.fromarray(img) + if img.mode == "RGBA": + # The checkpoint was trained with the alpha composited over white for the vision encoder. + # Only this copy is flattened; the VAE still reads all four channels. + white = PILImage.new("RGB", img.size, (255, 255, 255)) + white.paste(img, mask=img.getchannel("A")) + img = white + condition_pil_list.append(img) + + # Left padding, as the checkpoint was trained with. `_extract_masked_hidden` drops the padding either way, + # but the side decides the positions the encoder sees for a batch of prompts of different lengths. + processor_kwargs = { + "text": prompts, + "padding": True, + "padding_side": "left", + "return_tensors": "pt", + } + if not is_t2i: + processor_kwargs["images"] = condition_pil_list + + model_inputs = processor(**processor_kwargs).to(device) + + forward_kwargs = { + "input_ids": model_inputs.input_ids, + "attention_mask": model_inputs.attention_mask, + "output_hidden_states": True, + } + if not is_t2i and hasattr(model_inputs, "pixel_values"): + forward_kwargs.update(pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw) + if hasattr(model_inputs, "mm_token_type_ids"): + forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids + + # `hidden_states[-1]` has to be the last decoder layer's output, before the text encoder's final RMSNorm: + # that is what the transformer was trained on. It is what transformers 4.x returns there, but from + # transformers 5.0 the output capturing ties that entry to `last_hidden_state`, so it comes back normalized + # instead — a third of the signal the transformer reads, which shows up first in rendered text. A forward hook + # returning the module's input replaces its output, which neutralizes the norm for this call on either version. + # TODO: replace this with `tie_last_hidden_states=False` in the text encoder's config, which + # huggingface/transformers#48087 adds, once that ships in a stable transformers release (5.18). + text_model = getattr(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] + + split_hidden_states = list( + torch.split( + hidden_states[model_inputs.attention_mask.bool()], model_inputs.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) + for sample_ids, sample_mask in zip(model_inputs.input_ids, model_inputs.attention_mask) + ] + image_pad_mask = [e[drop_idx:] for e in image_pad_mask] + + 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 + + +class QwenImage21TextEncoderStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Encode prompts and optional condition images together with Qwen3-VL." + + @property + def expected_components(self): + 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): + return [ + InputParam.template("prompt", required=True), + InputParam.template("negative_prompt"), + InputParam("condition_images", type_hint=list, description="Resized images to encode with each prompt."), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds_mask"), + OutputParam( + "image_pad_mask", type_hint=torch.Tensor, description="Vision token positions in each positive prompt." + ), + OutputParam( + "negative_image_pad_mask", + type_hint=torch.Tensor, + description="Vision token positions in each negative prompt.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + prompt = block_state.prompt + if isinstance(prompt, str): + prompt = [prompt] + if not prompt or not all(isinstance(p, str) for p in prompt): + raise ValueError("`prompt` must be a string or a nonempty list of strings.") + 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.condition_images, + components._execution_device, + ) + ) + block_state.negative_prompt_embeds = None + block_state.negative_prompt_embeds_mask = None + block_state.negative_image_pad_mask = None + components.guider.set_state(step=0, num_inference_steps=None, timestep=None) + if components.guider.num_conditions > 1: + negative_prompt = block_state.negative_prompt + if negative_prompt is None: + negative_prompt = "" + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * len(prompt) + if len(negative_prompt) != len(prompt): + raise ValueError("`negative_prompt` must have the same batch size as `prompt`.") + ( + 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, + negative_prompt, + block_state.condition_images, + components._execution_device, + ) + self.set_block_state(state, block_state) + return components, state + + +def encode_image(vae, image): + latents = vae.encode(image.to(device=vae.device, dtype=vae.dtype)).latent_dist.mode() + mean = latents.new_tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1) + std = latents.new_tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1) + return ((latents - mean) / std).flatten(2).transpose(1, 2) + + +class QwenImage21VaeEncoderStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Encode condition images into unpatched, normalized latent tokens." + + @property + def expected_components(self): + return [ComponentSpec("vae", AutoencoderKLQwenImage21)] + + @property + def inputs(self): + return [ + InputParam("vae_images", required=True, type_hint=list, description="Normalized RGBA condition tensors.") + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "condition_latents", type_hint=torch.Tensor, description="Concatenated condition image tokens." + ), + OutputParam( + "condition_shapes", + type_hint=list, + description="Frame, height and width for each condition latent grid.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.condition_latents = torch.cat( + [encode_image(components.vae, image) for image in block_state.vae_images], dim=1 + ) + block_state.condition_shapes = [ + (1, image.shape[-2] // 16, image.shape[-1] // 16) for image in block_state.vae_images + ] + self.set_block_state(state, block_state) + return components, state + + +class QwenImage21InpaintVaeEncoderStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Encode the source image at the target resolution for inpainting preservation." + + @property + def expected_components(self): + return [ComponentSpec("vae", AutoencoderKLQwenImage21)] + + @property + def inputs(self): + return [ + InputParam( + "source_image", required=True, type_hint=torch.Tensor, description="Normalized source RGBA image." + ) + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "source_latents", + type_hint=torch.Tensor, + description="Source tokens used to preserve the unmasked area.", + ) + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.source_latents = encode_image(components.vae, block_state.source_image.unsqueeze(2)) + 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..cc0a3885e739 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/inputs.py @@ -0,0 +1,179 @@ +# 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 torch +from PIL import Image + +from ...configuration_utils import FrozenDict +from ...image_processor import VaeImageProcessor +from ..modular_pipeline import ModularPipelineBlocks +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +def prepare_condition_images(image_processor, images, output_resolution): + images = images if isinstance(images, list) else [images] + if not images: + raise ValueError("Provide at least one condition image.") + resized, tensors = [], [] + for image in images: + if isinstance(image, np.ndarray): + image = Image.fromarray(image) + if not isinstance(image, Image.Image): + raise ValueError("Condition images must be PIL images or numpy arrays in a flat list.") + image = image.convert("RGBA") + ratio = image.width / image.height + width = round(math.sqrt(output_resolution**2 * ratio) / 32) * 32 + height = round(math.sqrt(output_resolution**2 / ratio) / 32) * 32 + if min(width, height) < 32: + raise ValueError("The condition image aspect ratio produces a dimension below 32 pixels.") + resized.append(image_processor.resize(image, width=width, height=height)) + tensors.append(image_processor.preprocess(image, width=width, height=height).unsqueeze(2)) + return resized, tensors + + +class QwenImage21ProcessImagesStep(ModularPipelineBlocks): + model_name = "qwenimage21" + + @property + def description(self): + return "Resize condition images for the vision encoder and RGBA VAE." + + @property + def expected_components(self): + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ) + ] + + @property + def inputs(self): + return [ + InputParam.template("image", required=True), + InputParam( + "output_resolution", + default=1024, + type_hint=int, + description="Target side length used to resize condition images.", + ), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "condition_images", + type_hint=list, + description="Resized RGBA images for joint vision and text encoding.", + ), + OutputParam("vae_images", type_hint=list, description="Normalized RGBA condition tensors."), + OutputParam("height", type_hint=int, description="Output height in pixels."), + OutputParam("width", type_hint=int, description="Output width in pixels."), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.condition_images, block_state.vae_images = prepare_condition_images( + components.image_processor, block_state.image, block_state.output_resolution + ) + block_state.height = block_state.height or block_state.condition_images[-1].height + block_state.width = block_state.width or block_state.condition_images[-1].width + self.set_block_state(state, block_state) + return components, state + + +class QwenImage21ProcessInpaintStep(QwenImage21ProcessImagesStep): + @property + def description(self): + return "Prepare one source image, its repaint mask, and optional additional reference images." + + @property + def expected_components(self): + return super().expected_components + [ + ComponentSpec( + "mask_processor", + VaeImageProcessor, + config=FrozenDict( + {"vae_scale_factor": 16, "do_normalize": False, "do_binarize": True, "do_convert_grayscale": True} + ), + default_creation_method="from_config", + ) + ] + + @property + def inputs(self): + return super().inputs + [ + InputParam.template("mask_image", required=True), + InputParam( + "reference_images", + type_hint=list, + description="Additional PIL or numpy reference images, shared by every prompt. The source is always the first condition image.", + ), + ] + + @property + def intermediate_outputs(self): + return super().intermediate_outputs + [ + OutputParam( + "source_image", + type_hint=torch.Tensor, + description="Normalized source RGBA image at the output resolution.", + ), + OutputParam( + "processed_mask", + type_hint=torch.Tensor, + description="Binary repaint mask at the output resolution; white is repainted.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + source = block_state.image + if isinstance(source, np.ndarray): + source = Image.fromarray(source) + if not isinstance(source, Image.Image): + raise ValueError( + "Inpainting requires one source PIL or numpy image; pass extra images as `reference_images`." + ) + references = block_state.reference_images + if references is not None and not isinstance(references, list): + raise ValueError("`reference_images` must be a flat list of images.") + block_state.condition_images, block_state.vae_images = prepare_condition_images( + components.image_processor, [source] + (references or []), block_state.output_resolution + ) + height = block_state.height or block_state.condition_images[0].height + width = block_state.width or block_state.condition_images[0].width + if min(height, width) < 32 or height % 32 or width % 32: + raise ValueError("`height` and `width` must be positive multiples of 32.") + block_state.height, block_state.width = height, width + block_state.source_image = components.image_processor.preprocess( + source.convert("RGBA"), height=height, width=width + ) + block_state.processed_mask = components.mask_processor.preprocess( + block_state.mask_image, height=height, width=width + ) + if block_state.processed_mask.shape[0] != 1: + raise ValueError("Inpainting requires one mask shared by the prompt batch.") + 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..9b8ca8704216 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py @@ -0,0 +1,612 @@ +# 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 ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + QwenImage21PrepareInpaintStep, + QwenImage21PrepareLatentsStep, + QwenImage21SetTimestepsStep, + QwenImage21TextInputsStep, +) +from .decoders import QwenImage21DecodeStep +from .denoise import QwenImage21DenoiseStep, QwenImage21InpaintDenoiseStep +from .encoders import QwenImage21InpaintVaeEncoderStep, QwenImage21TextEncoderStep, QwenImage21VaeEncoderStep +from .inputs import QwenImage21ProcessImagesStep, QwenImage21ProcessInpaintStep + + +QwenImage21CoreDenoiseBlocks = InsertableDict( + [ + ("input", QwenImage21TextInputsStep()), + ("prepare_latents", QwenImage21PrepareLatentsStep()), + ("set_timesteps", QwenImage21SetTimestepsStep()), + ("denoise", QwenImage21DenoiseStep()), + ] +) + + +# auto_docstring +class QwenImage21CoreDenoiseStep(SequentialPipelineBlocks): + """ + Prepare and denoise target latents from reusable text and image embeddings. + + Components: + transformer (`QwenImage21Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. 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. + prompt_embeds_mask (`Tensor`, *optional*): + mask for the text embeddings. 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`): + Positive prompt vision positions. + negative_image_pad_mask (`Tensor`, *optional*): + Negative prompt vision positions. + 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. + output_resolution (`int`, *optional*, defaults to 1024): + Default output side length. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + condition_latents (`Tensor`, *optional*): + Packed condition tokens. + condition_shapes (`list`, *optional*): + Spatial shape of each condition image. + 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 step-independent text and condition-image keys and values. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage21" + block_classes = QwenImage21CoreDenoiseBlocks.values() + block_names = QwenImage21CoreDenoiseBlocks.keys() + + @property + def description(self): + return "Prepare and denoise target latents from reusable text and image embeddings." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +QwenImage21InpaintCoreDenoiseBlocks = InsertableDict( + [ + ("input", QwenImage21TextInputsStep()), + ("prepare_latents", QwenImage21PrepareLatentsStep()), + ("set_timesteps", QwenImage21SetTimestepsStep()), + ("prepare_inpaint", QwenImage21PrepareInpaintStep()), + ("denoise", QwenImage21InpaintDenoiseStep()), + ] +) + + +# auto_docstring +class QwenImage21InpaintCoreDenoiseStep(QwenImage21CoreDenoiseStep): + """ + Prepare and denoise a masked source with optional reference-image conditioning. + + Components: + transformer (`QwenImage21Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. 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. + prompt_embeds_mask (`Tensor`, *optional*): + mask for the text embeddings. 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`): + Positive prompt vision positions. + negative_image_pad_mask (`Tensor`, *optional*): + Negative prompt vision positions. + 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. + output_resolution (`int`, *optional*, defaults to 1024): + Default output side length. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + condition_latents (`Tensor`, *optional*): + Packed condition tokens. + condition_shapes (`list`, *optional*): + Spatial shape of each condition image. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + source_latents (`Tensor`): + Encoded source tokens. + processed_mask (`Tensor`): + Binary repaint mask. + strength (`float`, *optional*, defaults to 1.0): + Strength for img2img/inpainting. + use_kv_cache (`bool`, *optional*, defaults to True): + Cache step-independent text and condition-image keys and values. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = QwenImage21InpaintCoreDenoiseBlocks.values() + block_names = QwenImage21InpaintCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Prepare and denoise a masked source with optional reference-image conditioning." + + +QwenImage21Text2ImageBlocks = InsertableDict( + [ + ("text_encoder", QwenImage21TextEncoderStep()), + ("denoise", QwenImage21CoreDenoiseStep()), + ("decode", QwenImage21DecodeStep()), + ] +) + + +# auto_docstring +class QwenImage21Text2ImageStep(SequentialPipelineBlocks): + """ + Generate RGBA images from text with Qwen-Image 2.1. + + Components: + text_encoder (`Qwen3VLForConditionalGeneration`) processor (`Qwen3VLProcessor`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImage21Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLQwenImage21`) image_processor (`VaeImageProcessor`) + + 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. + condition_images (`list`, *optional*): + Resized images to encode with each prompt. + 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. + output_resolution (`int`, *optional*, defaults to 1024): + Default output side length. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + condition_latents (`Tensor`, *optional*): + Packed condition tokens. + condition_shapes (`list`, *optional*): + Spatial shape of each condition image. + 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 step-independent text and condition-image keys and values. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Vision token positions in each positive prompt. + negative_image_pad_mask (`Tensor`): + Vision token positions in each negative prompt. + latents (`Tensor`): + Initial target noise. + height (`int`): + Output height in pixels. + width (`int`): + Output width in pixels. + condition_latents (`Tensor`): + Condition tokens expanded to the output batch. + img_shapes (`list`): + Condition and target grid shapes for rotary embeddings. + img_mask (`Tensor`): + Joint positive prompt and target vision positions. + negative_img_mask (`Tensor`): + Joint negative prompt and target vision positions. + timesteps (`Tensor`): + Denoising timesteps. + num_inference_steps (`int`): + Number of denoising timesteps. + noise_pred (`Tensor`): + Guided target flow prediction. + kv_caches (`dict`): + Temporary KV caches, cleared after generation. + cache_enabled (`bool`): + Whether causal-condition caching is active. + images (`list`): + Generated images. + """ + + model_name = "qwenimage21" + block_classes = QwenImage21Text2ImageBlocks.values() + block_names = QwenImage21Text2ImageBlocks.keys() + + @property + def description(self): + return "Generate RGBA images from text with Qwen-Image 2.1." + + +QwenImage21ImageConditionedBlocks = InsertableDict( + [ + ("preprocess", QwenImage21ProcessImagesStep()), + ("text_encoder", QwenImage21TextEncoderStep()), + ("vae_encoder", QwenImage21VaeEncoderStep()), + ("denoise", QwenImage21CoreDenoiseStep()), + ("decode", QwenImage21DecodeStep()), + ] +) + + +# auto_docstring +class QwenImage21ImageConditionedStep(SequentialPipelineBlocks): + """ + Generate images conditioned on one or more reference images. + + Components: + image_processor (`VaeImageProcessor`) text_encoder (`Qwen3VLForConditionalGeneration`) processor + (`Qwen3VLProcessor`) guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLQwenImage21`) transformer + (`QwenImage21Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + 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 resize condition images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + 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. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image 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 step-independent text and condition-image keys and values. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + condition_images (`list`): + Resized RGBA images for joint vision and text encoding. + vae_images (`list`): + Normalized RGBA condition tensors. + height (`int`): + Output height in pixels. + width (`int`): + Output width in pixels. + prompt_embeds (`Tensor`): + The prompt embeddings. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Vision token positions in each positive prompt. + negative_image_pad_mask (`Tensor`): + Vision token positions in each negative prompt. + condition_latents (`Tensor`): + Concatenated condition image tokens. + condition_shapes (`list`): + Frame, height and width for each condition latent grid. + latents (`Tensor`): + Initial target noise. + img_shapes (`list`): + Condition and target grid shapes for rotary embeddings. + img_mask (`Tensor`): + Joint positive prompt and target vision positions. + negative_img_mask (`Tensor`): + Joint negative prompt and target vision positions. + timesteps (`Tensor`): + Denoising timesteps. + num_inference_steps (`int`): + Number of denoising timesteps. + noise_pred (`Tensor`): + Guided target flow prediction. + kv_caches (`dict`): + Temporary KV caches, cleared after generation. + cache_enabled (`bool`): + Whether causal-condition caching is active. + images (`list`): + Generated images. + """ + + model_name = "qwenimage21" + block_classes = QwenImage21ImageConditionedBlocks.values() + block_names = QwenImage21ImageConditionedBlocks.keys() + + @property + def description(self): + return "Generate images conditioned on one or more reference images." + + +QwenImage21InpaintBlocks = InsertableDict( + [ + ("preprocess", QwenImage21ProcessInpaintStep()), + ("text_encoder", QwenImage21TextEncoderStep()), + ("vae_encoder", QwenImage21VaeEncoderStep()), + ("source_encoder", QwenImage21InpaintVaeEncoderStep()), + ("denoise", QwenImage21InpaintCoreDenoiseStep()), + ("decode", QwenImage21DecodeStep()), + ] +) + + +# auto_docstring +class QwenImage21InpaintStep(SequentialPipelineBlocks): + """ + Repaint one source image using a mask and optional additional references. + + Components: + image_processor (`VaeImageProcessor`) mask_processor (`VaeImageProcessor`) text_encoder + (`Qwen3VLForConditionalGeneration`) processor (`Qwen3VLProcessor`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLQwenImage21`) transformer (`QwenImage21Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + 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 resize condition images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + mask_image (`Image`): + Mask image for inpainting. + reference_images (`list`, *optional*): + Additional PIL or numpy reference images, shared by every prompt. The source is always the first + condition image. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + 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. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 1.0): + Strength for img2img/inpainting. + use_kv_cache (`bool`, *optional*, defaults to True): + Cache step-independent text and condition-image keys and values. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + condition_images (`list`): + Resized RGBA images for joint vision and text encoding. + vae_images (`list`): + Normalized RGBA condition tensors. + height (`int`): + Output height in pixels. + width (`int`): + Output width in pixels. + source_image (`Tensor`): + Normalized source RGBA image at the output resolution. + processed_mask (`Tensor`): + Binary repaint mask at the output resolution; white is repainted. + prompt_embeds (`Tensor`): + The prompt embeddings. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + image_pad_mask (`Tensor`): + Vision token positions in each positive prompt. + negative_image_pad_mask (`Tensor`): + Vision token positions in each negative prompt. + condition_latents (`Tensor`): + Concatenated condition image tokens. + condition_shapes (`list`): + Frame, height and width for each condition latent grid. + source_latents (`Tensor`): + Source tokens used to preserve the unmasked area. + latents (`Tensor`): + Initial target noise. + img_shapes (`list`): + Condition and target grid shapes for rotary embeddings. + img_mask (`Tensor`): + Joint positive prompt and target vision positions. + negative_img_mask (`Tensor`): + Joint negative prompt and target vision positions. + timesteps (`Tensor`): + Denoising timesteps. + num_inference_steps (`int`): + Number of denoising timesteps. + initial_noise (`Tensor`): + Noise reused when preserving the source. + mask (`Tensor`): + Repaint weights for target latent tokens. + noise_pred (`Tensor`): + Guided target flow prediction. + kv_caches (`dict`): + Temporary KV caches, cleared after generation. + cache_enabled (`bool`): + Whether causal-condition caching is active. + images (`list`): + Generated images. + """ + + model_name = "qwenimage21" + block_classes = QwenImage21InpaintBlocks.values() + block_names = QwenImage21InpaintBlocks.keys() + + @property + def description(self): + return "Repaint one source image using a mask and optional additional references." + + +# auto_docstring +class QwenImage21AutoBlocks(ConditionalPipelineBlocks): + """ + Qwen-Image 2.1 text-to-image, image-conditioned generation, and inpainting with optional references. + + Components: + image_processor (`VaeImageProcessor`) mask_processor (`VaeImageProcessor`) text_encoder + (`Qwen3VLForConditionalGeneration`) processor (`Qwen3VLProcessor`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLQwenImage21`) transformer (`QwenImage21Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + 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 resize condition images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + mask_image (`Image`, *optional*): + Mask image for inpainting. + reference_images (`list`, *optional*): + Additional PIL or numpy reference images, shared by every prompt. The source is always the first + condition image. + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + 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. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 1.0): + Strength for img2img/inpainting. + use_kv_cache (`bool`, *optional*, defaults to True): + Cache step-independent text and condition-image keys and values. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + condition_images (`list`, *optional*): + Resized images to encode with each prompt. + condition_latents (`Tensor`, *optional*): + Packed condition tokens. + condition_shapes (`list`, *optional*): + Spatial shape of each condition image. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage21" + block_classes = [QwenImage21InpaintStep, QwenImage21ImageConditionedStep, QwenImage21Text2ImageStep] + block_names = ["inpainting", "image_conditioned", "text2image"] + block_trigger_inputs = ["mask_image", "image", "reference_images"] + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"prompt": True, "image": True}, + "inpainting": {"prompt": True, "image": True, "mask_image": True}, + } + + def select_block(self, mask_image=None, image=None, reference_images=None): + if mask_image is not None: + if image is None: + raise ValueError("`mask_image` requires a source `image`.") + return "inpainting" + if reference_images is not None: + raise ValueError( + "`reference_images` requires inpainting with `image` and `mask_image`. For unmasked generation, pass the condition images as `image`." + ) + return "image_conditioned" if image is not None else "text2image" + + @property + def available_workflows(self): + return list(self._workflow_map) + + def get_workflow(self, workflow_name: str): + if workflow_name not in self._workflow_map: + raise ValueError(f"Unknown workflow {workflow_name!r}. Available workflows: {self.available_workflows}") + return self.get_execution_blocks(**self._workflow_map[workflow_name]) + + @property + def description(self): + return "Qwen-Image 2.1 text-to-image, image-conditioned generation, and inpainting with optional references." + + @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..553793d801cf --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage21/modular_pipeline.py @@ -0,0 +1,22 @@ +# 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 ..modular_pipeline import ModularPipeline + + +class QwenImage21ModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): + """Modular Qwen-Image 2.1 generation and inpainting with optional reference images.""" + + default_blocks_name = "QwenImage21AutoBlocks" 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 ed724e7de751..d15ca9a0c97b 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -632,6 +632,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..8b137891791f --- /dev/null +++ b/tests/modular_pipelines/qwenimage21/__init__.py @@ -0,0 +1 @@ + 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..fbaaee2913e0 --- /dev/null +++ b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py @@ -0,0 +1,296 @@ +# 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 numpy as np +import pytest +import torch +from PIL import Image + +from diffusers import ( + ClassifierFreeGuidance, + FlowMatchEulerDiscreteScheduler, + ModularPipeline, + QwenImage21AutoBlocks, + QwenImage21ModularPipeline, + QwenImage21Pipeline, +) +from diffusers.modular_pipelines.qwenimage21.modular_blocks_qwenimage21 import QwenImage21CoreDenoiseStep + +from ...pipelines.qwenimage21.test_qwenimage21 import QwenImage21PipelineTesterConfig +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularGuiderTesterMixin, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) + + +@pytest.fixture(scope="module") +def tiny_checkpoint(tmp_path_factory): + root = tmp_path_factory.mktemp("qwenimage21") + standard = root / "standard" + components = QwenImage21PipelineTesterConfig().get_dummy_components() + components["scheduler"] = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=True, max_image_seq_len=8192, base_shift=0.5, max_shift=0.9, shift_terminal=0.02 + ) + QwenImage21Pipeline(**components).save_pretrained(standard) + pipe = ModularPipeline.from_pretrained(str(standard)) + pipe.load_components(dtype=torch.float32) + pipe.save_pretrained(str(root / "modular")) + return root + + +def condition_image(seed=0, size=(32, 32)): + return Image.fromarray(np.random.RandomState(seed).randint(0, 256, (*size, 4), dtype=np.uint8)) + + +class QwenImage21ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = QwenImage21ModularPipeline + pipeline_blocks_class = QwenImage21AutoBlocks + pretrained_model_name_or_path = None + params = frozenset(["prompt", "negative_prompt", "height", "width", "image", "mask_image", "reference_images"]) + batch_params = frozenset(["prompt", "negative_prompt"]) + expected_workflow_blocks = { + "text2image": [ + ("text_encoder", "QwenImage21TextEncoderStep"), + ("denoise.input", "QwenImage21TextInputsStep"), + ("denoise.prepare_latents", "QwenImage21PrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImage21SetTimestepsStep"), + ("denoise.denoise", "QwenImage21DenoiseStep"), + ("decode", "QwenImage21DecodeStep"), + ], + "image_conditioned": [ + ("preprocess", "QwenImage21ProcessImagesStep"), + ("text_encoder", "QwenImage21TextEncoderStep"), + ("vae_encoder", "QwenImage21VaeEncoderStep"), + ("denoise.input", "QwenImage21TextInputsStep"), + ("denoise.prepare_latents", "QwenImage21PrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImage21SetTimestepsStep"), + ("denoise.denoise", "QwenImage21DenoiseStep"), + ("decode", "QwenImage21DecodeStep"), + ], + "inpainting": [ + ("preprocess", "QwenImage21ProcessInpaintStep"), + ("text_encoder", "QwenImage21TextEncoderStep"), + ("vae_encoder", "QwenImage21VaeEncoderStep"), + ("source_encoder", "QwenImage21InpaintVaeEncoderStep"), + ("denoise.input", "QwenImage21TextInputsStep"), + ("denoise.prepare_latents", "QwenImage21PrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImage21SetTimestepsStep"), + ("denoise.prepare_inpaint", "QwenImage21PrepareInpaintStep"), + ("denoise.denoise", "QwenImage21InpaintDenoiseStep"), + ("decode", "QwenImage21DecodeStep"), + ], + } + + @pytest.fixture(scope="class", autouse=True) + def checkpoint(self, request, tiny_checkpoint): + request.cls.pretrained_model_name_or_path = str(tiny_checkpoint / "modular") + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "a cat", + "height": 32, + "width": 32, + "output_resolution": 32, + "num_inference_steps": 2, + "generator": self.get_generator(seed), + "output_type": "pt", + } + + +class TestQwenImage21ModularPipeline(QwenImage21ModularPipelineTesterConfig, ModularPipelineTesterMixin): + def test_inference_batch_single_identical(self): + # The tiny RGBA VAE amplifies the observed 1e-7 CUDA latent difference to 1.6e-4 in decoded pixels. + super().test_inference_batch_single_identical(expected_max_diff=3e-4) + + @pytest.mark.parametrize("references", [0, 1, 2]) + @pytest.mark.parametrize("use_kv_cache", [False, True]) + @pytest.mark.parametrize("guidance", [1.0, 3.0]) + def test_standard_parity(self, tiny_checkpoint, references, use_kv_cache, guidance): + reference = QwenImage21Pipeline.from_pretrained(tiny_checkpoint / "standard", dtype=torch.float32) + pipe = self.get_pipeline() + pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=guidance)) + inputs = self.get_dummy_inputs() + inputs.update(use_kv_cache=use_kv_cache, negative_prompt="blurry") + if references: + inputs["image"] = [condition_image(i) for i in range(references)] + expected = reference(**inputs, true_cfg_scale=guidance).images + inputs["generator"] = self.get_generator() + actual = pipe(**inputs, output="images") + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + @pytest.mark.parametrize("references", [0, 2]) + @pytest.mark.parametrize("strength", [0.5, 1.0]) + def test_inpaint_black_mask_preserves_source_latents(self, references, strength): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.update( + image=condition_image(), + mask_image=Image.new("L", (32, 32), 0), + reference_images=[condition_image(i + 1) for i in range(references)], + strength=strength, + ) + state = pipe(**inputs) + torch.testing.assert_close(state.get("latents"), state.get("source_latents"), atol=0, rtol=0) + assert torch.isfinite(state.get("images")).all() + + def test_inpaint_white_mask_matches_conditioned_generation(self): + pipe = self.get_pipeline() + source, reference = condition_image(), condition_image(1) + expected = pipe(**self.get_dummy_inputs(), image=[source, reference], output="images") + actual = pipe( + **self.get_dummy_inputs(), + image=source, + reference_images=[reference], + mask_image=Image.new("L", (32, 32), 255), + output="images", + ) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_multireference_inpaint_batch_and_partial_mask(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.update( + prompt=["a cat", "a dog"], + height=64, + width=64, + output_resolution=64, + num_images_per_prompt=2, + image=condition_image(), + reference_images=[condition_image(1, (32, 64)), condition_image(2)], + mask_image=Image.fromarray(np.pad(np.full((32, 32), 255, dtype=np.uint8), 16)), + ) + state = pipe(**inputs) + assert state.get("images").shape == (4, 4, 64, 64) + latents, source, mask = state.get("latents"), state.get("source_latents"), state.get("mask") + torch.testing.assert_close(latents * (1 - mask), source * (1 - mask), atol=0, rtol=0) + assert (latents * mask - source * mask).abs().max() > 0 + assert state.get("condition_latents").shape[1] > 4 + + def test_reference_images_affect_masked_output(self): + pipe = self.get_pipeline() + inputs = {"image": condition_image(), "mask_image": Image.new("L", (32, 32), 255), "output": "images"} + first = pipe(**self.get_dummy_inputs(), **inputs) + second = pipe(**self.get_dummy_inputs(), reference_images=[condition_image(1)], **inputs) + assert (first - second).abs().max() > 1e-6 + + def test_repeated_calls_reset_kv_cache(self): + pipe = self.get_pipeline() + first = pipe(**self.get_dummy_inputs(), output="images") + pipe(**self.get_dummy_inputs(1), image=[condition_image(), condition_image(1)], output="images") + second = pipe(**self.get_dummy_inputs(), output="images") + torch.testing.assert_close(first, second, atol=0, rtol=0) + + def test_guidance_start_with_kv_cache(self): + pipe = self.get_pipeline() + pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=3.0, start=0.5)) + inputs = self.get_dummy_inputs() + inputs.update(num_inference_steps=4, negative_prompt="blurry", image=[condition_image(), condition_image(1)]) + expected = pipe(**inputs, use_kv_cache=False, output="images") + inputs["generator"] = self.get_generator() + actual = pipe(**inputs, use_kv_cache=True, output="images") + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) + + def test_inpaint_dimensions_follow_source(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.pop("height") + inputs.pop("width") + inputs.update( + output_resolution=64, + image=condition_image(size=(32, 64)), + mask_image=Image.new("L", (64, 32)), + reference_images=[condition_image(size=(64, 32))], + ) + state = pipe(**inputs) + assert state.get("height") == 32 + assert state.get("width") == 96 + assert state.get("images").shape == (1, 4, 32, 96) + assert state.get("condition_shapes") == [(1, 2, 6), (1, 6, 2)] + + def test_reusable_text_encoder(self): + pipe = self.get_pipeline() + blocks = pipe.blocks.get_workflow("text2image") + encoder = blocks.sub_blocks["text_encoder"].init_pipeline() + encoder.update_components(text_encoder=pipe.text_encoder, processor=pipe.processor, guider=pipe.guider) + encoded = encoder(prompt="a cat") + core = QwenImage21CoreDenoiseStep().init_pipeline() + core.update_components(transformer=pipe.transformer, scheduler=pipe.scheduler, guider=pipe.guider) + inputs = self.get_dummy_inputs() + inputs.pop("prompt") + inputs.pop("output_type") + state = core( + **inputs, + prompt_embeds=encoded.get("prompt_embeds"), + prompt_embeds_mask=encoded.get("prompt_embeds_mask"), + image_pad_mask=encoded.get("image_pad_mask"), + num_images_per_prompt=2, + ) + assert state.get("latents").shape == (2, 4, 8) + assert encoded.get("prompt_embeds").shape[0] == 1 + + @pytest.mark.parametrize("strength", [0, -0.1, 1.1, 0.1]) + def test_invalid_strength(self, strength): + pipe = self.get_pipeline() + with pytest.raises(ValueError, match="strength"): + pipe( + **self.get_dummy_inputs(), + image=condition_image(), + mask_image=Image.new("L", (32, 32)), + strength=strength, + ) + + def test_inpaint_rejects_multiple_sources(self): + pipe = self.get_pipeline() + with pytest.raises(ValueError, match="one source"): + pipe( + **self.get_dummy_inputs(), + image=[condition_image(), condition_image(1)], + mask_image=Image.new("L", (32, 32)), + ) + + def test_rejects_references_without_mask(self): + pipe = self.get_pipeline() + with pytest.raises(ValueError, match="reference_images"): + pipe(**self.get_dummy_inputs(), image=condition_image(), reference_images=[condition_image(1)]) + + def test_rejects_mask_without_source(self): + pipe = self.get_pipeline() + with pytest.raises(ValueError, match="requires a source"): + pipe(**self.get_dummy_inputs(), mask_image=Image.new("L", (32, 32))) + + def test_load_from_standard_index(self, tiny_checkpoint): + pipe = ModularPipeline.from_pretrained(str(tiny_checkpoint / "standard")) + assert isinstance(pipe, QwenImage21ModularPipeline) + pipe.load_components(dtype=torch.float32) + assert pipe(**self.get_dummy_inputs(), output="images").shape == (1, 4, 32, 32) + + +class TestQwenImage21ModularLoading(QwenImage21ModularPipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestQwenImage21ModularWorkflows(QwenImage21ModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass + + +class TestQwenImage21ModularGuiders(QwenImage21ModularPipelineTesterConfig, ModularGuiderTesterMixin): + pass + + +class TestQwenImage21ModularMemory(QwenImage21ModularPipelineTesterConfig, ModularMemoryTesterMixin): + pass