diffusers_src/src/diffusers/modular_pipelines/qwenimage/encoders.py

# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Text and VAE encoder blocks for QwenImage pipelines.
"""

import PIL
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor

from ...configuration_utils import FrozenDict
from ...guiders import ClassifierFreeGuidance
from ...image_processor import InpaintProcessor, VaeImageProcessor, is_valid_image, is_valid_image_imagelist
from ...models import AutoencoderKLQwenImage, QwenImageControlNetModel, QwenImageMultiControlNetModel
from ...pipelines.qwenimage.pipeline_qwenimage_edit import calculate_dimensions
from ...utils import logging
from ...utils.torch_utils import unwrap_module
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .modular_pipeline import QwenImageModularPipeline
from .prompt_templates import (
    QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE,
    QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE,
    QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX,
    QWENIMAGE_EDIT_PROMPT_TEMPLATE,
    QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX,
    QWENIMAGE_LAYERED_CAPTION_PROMPT_CN,
    QWENIMAGE_LAYERED_CAPTION_PROMPT_EN,
    QWENIMAGE_PROMPT_TEMPLATE,
    QWENIMAGE_PROMPT_TEMPLATE_START_IDX,
)


logger = logging.get_logger(__name__)


def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor):
    bool_mask = mask.bool()
    valid_lengths = bool_mask.sum(dim=1)
    selected = hidden_states[bool_mask]
    split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
    return split_result


def get_qwen_prompt_embeds(
    text_encoder,
    tokenizer,
    prompt: str | list[str] = None,
    prompt_template_encode: str = QWENIMAGE_PROMPT_TEMPLATE,
    prompt_template_encode_start_idx: int = QWENIMAGE_PROMPT_TEMPLATE_START_IDX,
    tokenizer_max_length: int = 1024,
    device: torch.device | None = None,
):
    prompt = [prompt] if isinstance(prompt, str) else prompt

    template = prompt_template_encode
    drop_idx = prompt_template_encode_start_idx
    txt = [template.format(e) for e in prompt]
    txt_tokens = tokenizer(
        txt, max_length=tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt"
    ).to(device)
    encoder_hidden_states = text_encoder(
        input_ids=txt_tokens.input_ids,
        attention_mask=txt_tokens.attention_mask,
        output_hidden_states=True,
    )
    hidden_states = encoder_hidden_states.hidden_states[-1]

    split_hidden_states = _extract_masked_hidden(hidden_states, txt_tokens.attention_mask)
    split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
    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]
    )

    prompt_embeds = prompt_embeds.to(device=device)

    return prompt_embeds, encoder_attention_mask


def get_qwen_prompt_embeds_edit(
    text_encoder,
    processor,
    prompt: str | list[str] = None,
    image: torch.Tensor | None = None,
    prompt_template_encode: str = QWENIMAGE_EDIT_PROMPT_TEMPLATE,
    prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX,
    device: torch.device | None = None,
):
    prompt = [prompt] if isinstance(prompt, str) else prompt

    template = prompt_template_encode
    drop_idx = prompt_template_encode_start_idx
    txt = [template.format(e) for e in prompt]

    model_inputs = processor(
        text=txt,
        images=image,
        padding=True,
        return_tensors="pt",
    ).to(device)

    outputs = text_encoder(
        input_ids=model_inputs["input_ids"],
        attention_mask=model_inputs["attention_mask"],
        pixel_values=model_inputs.get("pixel_values"),
        image_grid_thw=model_inputs.get("image_grid_thw"),
        output_hidden_states=True,
    )

    hidden_states = outputs.hidden_states[-1]
    split_hidden_states = _extract_masked_hidden(hidden_states, model_inputs["attention_mask"])
    split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
    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]
    )

    prompt_embeds = prompt_embeds.to(device=device)

    return prompt_embeds, encoder_attention_mask


def get_qwen_prompt_embeds_edit_plus(
    text_encoder,
    processor,
    prompt: str | list[str] = None,
    image: torch.Tensor | list[PIL.Image.Image, PIL.Image.Image] | None = None,
    prompt_template_encode: str = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE,
    img_template_encode: str = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE,
    prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX,
    device: torch.device | None = None,
):
    prompt = [prompt] if isinstance(prompt, str) else prompt
    if isinstance(image, list):
        base_img_prompt = ""
        for i, img in enumerate(image):
            base_img_prompt += img_template_encode.format(i + 1)
    elif image is not None:
        base_img_prompt = img_template_encode.format(1)
    else:
        base_img_prompt = ""

    template = prompt_template_encode

    drop_idx = prompt_template_encode_start_idx
    txt = [template.format(base_img_prompt + e) for e in prompt]

    model_inputs = processor(
        text=txt,
        images=image,
        padding=True,
        return_tensors="pt",
    ).to(device)
    outputs = text_encoder(
        input_ids=model_inputs["input_ids"],
        attention_mask=model_inputs["attention_mask"],
        pixel_values=model_inputs.get("pixel_values"),
        image_grid_thw=model_inputs.get("image_grid_thw"),
        output_hidden_states=True,
    )

    hidden_states = outputs.hidden_states[-1]
    split_hidden_states = _extract_masked_hidden(hidden_states, model_inputs["attention_mask"])
    split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
    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]
    )

    prompt_embeds = prompt_embeds.to(device=device)
    return prompt_embeds, encoder_attention_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")


# Modified from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._encode_vae_image
def encode_vae_image(
    image: torch.Tensor,
    vae: AutoencoderKLQwenImage,
    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


# ====================
# 1. RESIZE
# ====================
# In QwenImage pipelines, resize is a separate step because the resized image is used in VL encoding and vae encoder blocks:
#
#   image (PIL.Image.Image)
#       │
#       ▼
#   resized_image ([PIL.Image.Image])
#       │
#       ├──► text_encoder ──► prompt_embeds, prompt_embeds_mask
#       │    (VL encoding needs the resized image for vision-language fusion)
#       │
#       └──► image_processor ──► processed_image (torch.Tensor, pixel space)
#                │
#                ▼
#            vae_encoder ──► image_latents (torch.Tensor, latent space)
#
# In most of our other pipelines, resizing is done as part of the image preprocessing step.
# ====================


# auto_docstring
class QwenImageEditResizeStep(ModularPipelineBlocks):
    """
    Image Resize step that resize the image to target area while maintaining the aspect ratio.

      Components:
          image_resize_processor (`VaeImageProcessor`)

      Inputs:
          image (`Image | list`):
              Reference image(s) for denoising. Can be a single image or list of images.

      Outputs:
          resized_image (`list`):
              The resized images
    """

    model_name = "qwenimage-edit"

    @property
    def description(self) -> str:
        return "Image Resize step that resize the image to target area while maintaining the aspect ratio."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_resize_processor",
                VaeImageProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [InputParam.template("image")]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="resized_image",
                type_hint=list[PIL.Image.Image],
                description="The resized images",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: 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]

        image_width, image_height = images[0].size
        calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height)

        resized_images = [
            components.image_resize_processor.resize(image, height=calculated_height, width=calculated_width)
            for image in images
        ]

        block_state.resized_image = resized_images
        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageLayeredResizeStep(ModularPipelineBlocks):
    """
    Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while
    maintaining the aspect ratio.

      Components:
          image_resize_processor (`VaeImageProcessor`)

      Inputs:
          image (`Image | list`):
              Reference image(s) for denoising. Can be a single image or list of images.
          resolution (`int`, *optional*, defaults to 640):
              The target area to resize the image to, can be 1024 or 640

      Outputs:
          resized_image (`list`):
              The resized images
    """

    model_name = "qwenimage-layered"

    @property
    def description(self) -> str:
        return "Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while maintaining the aspect ratio."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_resize_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="resolution",
                default=640,
                type_hint=int,
                description="The target area to resize the image to, can be 1024 or 640",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="resized_image",
                type_hint=list[PIL.Image.Image],
                description="The resized images",
            )
        ]

    @staticmethod
    def check_inputs(resolution: int):
        if resolution not in [1024, 640]:
            raise ValueError(f"Resolution must be 1024 or 640 but is {resolution}")

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        self.check_inputs(resolution=block_state.resolution)

        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]

        image_width, image_height = images[0].size
        target_area = block_state.resolution * block_state.resolution
        calculated_width, calculated_height, _ = calculate_dimensions(target_area, image_width / image_height)

        resized_images = [
            components.image_resize_processor.resize(image, height=calculated_height, width=calculated_width)
            for image in images
        ]

        block_state.resized_image = resized_images
        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageEditPlusResizeStep(ModularPipelineBlocks):
    """
    Resize images for QwenImage Edit Plus pipeline.
      Produces two outputs: resized_image (1024x1024) for VAE encoding, resized_cond_image (384x384) for VL text
      encoding. Each image is resized independently based on its own aspect ratio.

      Components:
          image_resize_processor (`VaeImageProcessor`)

      Inputs:
          image (`Image | list`):
              Reference image(s) for denoising. Can be a single image or list of images.

      Outputs:
          resized_image (`list`):
              Images resized to 1024x1024 target area for VAE encoding
          resized_cond_image (`list`):
              Images resized to 384x384 target area for VL text encoding
    """

    model_name = "qwenimage-edit-plus"

    @property
    def description(self) -> str:
        return (
            "Resize images for QwenImage Edit Plus pipeline.\n"
            "Produces two outputs: resized_image (1024x1024) for VAE encoding, "
            "resized_cond_image (384x384) for VL text encoding.\n"
            "Each image is resized independently based on its own aspect ratio."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_resize_processor",
                VaeImageProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        # image
        return [InputParam.template("image")]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="resized_image",
                type_hint=list[PIL.Image.Image],
                description="Images resized to 1024x1024 target area for VAE encoding",
            ),
            OutputParam(
                name="resized_cond_image",
                type_hint=list[PIL.Image.Image],
                description="Images resized to 384x384 target area for VL text encoding",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: 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]

        # Resize each image independently based on its own aspect ratio
        resized_images = []
        resized_cond_images = []
        for image in images:
            image_width, image_height = image.size

            # For VAE encoder (1024x1024 target area)
            vae_width, vae_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height)
            resized_images.append(components.image_resize_processor.resize(image, height=vae_height, width=vae_width))

            # For VL text encoder (384x384 target area)
            vl_width, vl_height, _ = calculate_dimensions(384 * 384, image_width / image_height)
            resized_cond_images.append(
                components.image_resize_processor.resize(image, height=vl_height, width=vl_width)
            )

        block_state.resized_image = resized_images
        block_state.resized_cond_image = resized_cond_images
        self.set_block_state(state, block_state)
        return components, state


# ====================
# 2. GET IMAGE PROMPT
# ====================


# auto_docstring
class QwenImageLayeredGetImagePromptStep(ModularPipelineBlocks):
    """
    Auto-caption step that generates a text prompt from the input image if none is provided.
      Uses the VL model (text_encoder) to generate a description of the image. If prompt is already provided, this step
      passes through unchanged.

      Components:
          text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`)

      Inputs:
          prompt (`str`, *optional*):
              The prompt or prompts to guide image generation.
          resized_image (`Image`):
              The image to generate caption from, should be resized use the resize step
          use_en_prompt (`bool`, *optional*, defaults to False):
              Whether to use English prompt template

      Outputs:
          prompt (`str`):
              The prompt or prompts to guide image generation. If not provided, updated using image caption
    """

    model_name = "qwenimage-layered"

    def __init__(self):
        self.image_caption_prompt_en = QWENIMAGE_LAYERED_CAPTION_PROMPT_EN
        self.image_caption_prompt_cn = QWENIMAGE_LAYERED_CAPTION_PROMPT_CN
        super().__init__()

    @property
    def description(self) -> str:
        return (
            "Auto-caption step that generates a text prompt from the input image if none is provided.\n"
            "Uses the VL model (text_encoder) to generate a description of the image.\n"
            "If prompt is already provided, this step passes through unchanged."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
            ComponentSpec("processor", Qwen2VLProcessor),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template(
                "prompt", required=False
            ),  # it is not required for qwenimage-layered, unlike other pipelines
            InputParam(
                name="resized_image",
                required=True,
                type_hint=PIL.Image.Image,
                description="The image to generate caption from, should be resized use the resize step",
            ),
            InputParam(
                name="use_en_prompt",
                default=False,
                type_hint=bool,
                description="Whether to use English prompt template",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="prompt",
                type_hint=str,
                description="The prompt or prompts to guide image generation. If not provided, updated using image caption",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        device = components._execution_device

        # If prompt is empty or None, generate caption from image
        if block_state.prompt is None or block_state.prompt == "" or block_state.prompt == " ":
            if block_state.use_en_prompt:
                caption_prompt = self.image_caption_prompt_en
            else:
                caption_prompt = self.image_caption_prompt_cn

            model_inputs = components.processor(
                text=caption_prompt,
                images=block_state.resized_image,
                padding=True,
                return_tensors="pt",
            ).to(device)

            generated_ids = components.text_encoder.generate(**model_inputs, max_new_tokens=512)
            generated_ids_trimmed = [
                out_ids[len(in_ids) :] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
            ]
            output_text = components.processor.batch_decode(
                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
            )[0]

            block_state.prompt = output_text.strip()

        self.set_block_state(state, block_state)
        return components, state


# ====================
# 3. TEXT ENCODER
# ====================


# auto_docstring
class QwenImageTextEncoderStep(ModularPipelineBlocks):
    """
    Text Encoder step that generates text embeddings to guide the image generation.

      Components:
          text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`):
          The tokenizer to use 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.
          max_sequence_length (`int`, *optional*, defaults to 1024):
              Maximum sequence length for prompt encoding.

      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.
    """

    model_name = "qwenimage"

    def __init__(self):
        self.prompt_template_encode = QWENIMAGE_PROMPT_TEMPLATE
        self.prompt_template_encode_start_idx = QWENIMAGE_PROMPT_TEMPLATE_START_IDX
        self.tokenizer_max_length = 1024
        super().__init__()

    @property
    def description(self) -> str:
        return "Text Encoder step that generates text embeddings to guide the image generation."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration, description="The text encoder to use"),
            ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer to use"),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 4.0}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("prompt"),
            InputParam.template("negative_prompt"),
            InputParam.template("max_sequence_length", default=1024),
        ]

    @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"),
        ]

    @staticmethod
    def check_inputs(prompt, negative_prompt, max_sequence_length):
        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)}")

        if max_sequence_length is not None and max_sequence_length > 1024:
            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        device = components._execution_device
        self.check_inputs(block_state.prompt, block_state.negative_prompt, block_state.max_sequence_length)

        block_state.prompt_embeds, block_state.prompt_embeds_mask = get_qwen_prompt_embeds(
            components.text_encoder,
            components.tokenizer,
            prompt=block_state.prompt,
            prompt_template_encode=self.prompt_template_encode,
            prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
            tokenizer_max_length=self.tokenizer_max_length,
            device=device,
        )

        block_state.prompt_embeds = block_state.prompt_embeds[:, : block_state.max_sequence_length]
        block_state.prompt_embeds_mask = block_state.prompt_embeds_mask[:, : block_state.max_sequence_length]

        block_state.negative_prompt_embeds = None
        block_state.negative_prompt_embeds_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 = get_qwen_prompt_embeds(
                components.text_encoder,
                components.tokenizer,
                prompt=negative_prompt,
                prompt_template_encode=self.prompt_template_encode,
                prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
                tokenizer_max_length=self.tokenizer_max_length,
                device=device,
            )
            block_state.negative_prompt_embeds = block_state.negative_prompt_embeds[
                :, : block_state.max_sequence_length
            ]
            block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask[
                :, : block_state.max_sequence_length
            ]

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageEditTextEncoderStep(ModularPipelineBlocks):
    """
    Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image
    generation.

      Components:
          text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) 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 (`Image`):
              The image prompt to encode, should be resized using 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.
    """

    model_name = "qwenimage"

    def __init__(self):
        self.prompt_template_encode = QWENIMAGE_EDIT_PROMPT_TEMPLATE
        self.prompt_template_encode_start_idx = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX
        super().__init__()

    @property
    def description(self) -> str:
        return "Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image generation."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
            ComponentSpec("processor", Qwen2VLProcessor),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 4.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=PIL.Image.Image,
                description="The image prompt to encode, should be resized using 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"),
        ]

    @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: QwenImageModularPipeline, state: 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 = get_qwen_prompt_embeds_edit(
            components.text_encoder,
            components.processor,
            prompt=block_state.prompt,
            image=block_state.resized_image,
            prompt_template_encode=self.prompt_template_encode,
            prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
            device=device,
        )

        block_state.negative_prompt_embeds = None
        block_state.negative_prompt_embeds_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 = get_qwen_prompt_embeds_edit(
                components.text_encoder,
                components.processor,
                prompt=negative_prompt,
                image=block_state.resized_image,
                prompt_template_encode=self.prompt_template_encode,
                prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
                device=device,
            )

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageEditPlusTextEncoderStep(ModularPipelineBlocks):
    """
    Text Encoder step for QwenImage Edit Plus that processes prompt and multiple images together to generate text
    embeddings for guiding image generation.

      Components:
          text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) 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_cond_image (`Tensor`):
              The image(s) to encode, can be a single image or list of images, should be resized to 384x384 using
              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.
    """

    model_name = "qwenimage-edit-plus"

    def __init__(self):
        self.prompt_template_encode = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE
        self.img_template_encode = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE
        self.prompt_template_encode_start_idx = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX
        super().__init__()

    @property
    def description(self) -> str:
        return (
            "Text Encoder step for QwenImage Edit Plus that processes prompt and multiple images together "
            "to generate text embeddings for guiding image generation."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
            ComponentSpec("processor", Qwen2VLProcessor),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 4.0}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("prompt"),
            InputParam.template("negative_prompt"),
            InputParam(
                name="resized_cond_image",
                required=True,
                type_hint=torch.Tensor,
                description="The image(s) to encode, can be a single image or list of images, should be resized to 384x384 using 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"),
        ]

    @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: QwenImageModularPipeline, state: 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 = get_qwen_prompt_embeds_edit_plus(
            components.text_encoder,
            components.processor,
            prompt=block_state.prompt,
            image=block_state.resized_cond_image,
            prompt_template_encode=self.prompt_template_encode,
            img_template_encode=self.img_template_encode,
            prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
            device=device,
        )

        block_state.negative_prompt_embeds = None
        block_state.negative_prompt_embeds_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 = (
                get_qwen_prompt_embeds_edit_plus(
                    components.text_encoder,
                    components.processor,
                    prompt=negative_prompt,
                    image=block_state.resized_cond_image,
                    prompt_template_encode=self.prompt_template_encode,
                    img_template_encode=self.img_template_encode,
                    prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
                    device=device,
                )
            )

        self.set_block_state(state, block_state)
        return components, state


# ====================
# 4. IMAGE PREPROCESS
# ====================


# auto_docstring
class QwenImageInpaintProcessImagesInputStep(ModularPipelineBlocks):
    """
    Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be
    resized to the given height and width.

      Components:
          image_mask_processor (`InpaintProcessor`)

      Inputs:
          mask_image (`Image`):
              Mask image for inpainting.
          image (`Image | list`):
              Reference image(s) for denoising. Can be a single image or list of images.
          height (`int`, *optional*):
              The height in pixels of the generated image.
          width (`int`, *optional*):
              The width in pixels of the generated image.
          padding_mask_crop (`int`, *optional*):
              Padding for mask cropping in inpainting.

      Outputs:
          processed_image (`Tensor`):
              The processed image
          processed_mask_image (`Tensor`):
              The processed mask image
          mask_overlay_kwargs (`dict`):
              The kwargs for the postprocess step to apply the mask overlay
    """

    model_name = "qwenimage"

    @property
    def description(self) -> str:
        return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be resized to the given height and width."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_mask_processor",
                InpaintProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("mask_image"),
            InputParam.template("image"),
            InputParam.template("height"),
            InputParam.template("width"),
            InputParam.template("padding_mask_crop"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="processed_image",
                type_hint=torch.Tensor,
                description="The processed image",
            ),
            OutputParam(
                name="processed_mask_image",
                type_hint=torch.Tensor,
                description="The processed mask image",
            ),
            OutputParam(
                name="mask_overlay_kwargs",
                type_hint=dict,
                description="The kwargs for the postprocess step to apply the mask overlay",
            ),
        ]

    @staticmethod
    def check_inputs(height, width, vae_scale_factor):
        if height is not None and height % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")

        if width is not None and width % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        self.check_inputs(
            height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
        )
        height = block_state.height or components.default_height
        width = block_state.width or components.default_width

        block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = (
            components.image_mask_processor.preprocess(
                image=block_state.image,
                mask=block_state.mask_image,
                height=height,
                width=width,
                padding_mask_crop=block_state.padding_mask_crop,
            )
        )

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageEditInpaintProcessImagesInputStep(ModularPipelineBlocks):
    """
    Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be
    resized first.

      Components:
          image_mask_processor (`InpaintProcessor`)

      Inputs:
          mask_image (`Image`):
              Mask image for inpainting.
          resized_image (`Image`):
              The resized image. should be generated using a resize step
          padding_mask_crop (`int`, *optional*):
              Padding for mask cropping in inpainting.

      Outputs:
          processed_image (`Tensor`):
              The processed image
          processed_mask_image (`Tensor`):
              The processed mask image
          mask_overlay_kwargs (`dict`):
              The kwargs for the postprocess step to apply the mask overlay
    """

    model_name = "qwenimage-edit"

    @property
    def description(self) -> str:
        return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be resized first."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_mask_processor",
                InpaintProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("mask_image"),
            InputParam(
                name="resized_image",
                required=True,
                type_hint=PIL.Image.Image,
                description="The resized image. should be generated using a resize step",
            ),
            InputParam.template("padding_mask_crop"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(name="processed_image", type_hint=torch.Tensor, description="The processed image"),
            OutputParam(
                name="processed_mask_image",
                type_hint=torch.Tensor,
                description="The processed mask image",
            ),
            OutputParam(
                name="mask_overlay_kwargs",
                type_hint=dict,
                description="The kwargs for the postprocess step to apply the mask overlay",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        width, height = block_state.resized_image[0].size

        block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = (
            components.image_mask_processor.preprocess(
                image=block_state.resized_image,
                mask=block_state.mask_image,
                height=height,
                width=width,
                padding_mask_crop=block_state.padding_mask_crop,
            )
        )

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageProcessImagesInputStep(ModularPipelineBlocks):
    """
    Image Preprocess step. will resize the image to the given height and width.

      Components:
          image_processor (`VaeImageProcessor`)

      Inputs:
          image (`Image | list`):
              Reference image(s) for denoising. Can be a single image or list of images.
          height (`int`, *optional*):
              The height in pixels of the generated image.
          width (`int`, *optional*):
              The width in pixels of the generated image.

      Outputs:
          processed_image (`Tensor`):
              The processed image
    """

    model_name = "qwenimage"

    @property
    def description(self) -> str:
        return "Image Preprocess step. will resize the image to the given height and width."

    @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.template("height"),
            InputParam.template("width"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="processed_image",
                type_hint=torch.Tensor,
                description="The processed image",
            )
        ]

    @staticmethod
    def check_inputs(height, width, vae_scale_factor):
        if height is not None and height % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")

        if width is not None and width % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        self.check_inputs(
            height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
        )
        height = block_state.height or components.default_height
        width = block_state.width or components.default_width

        block_state.processed_image = components.image_processor.preprocess(
            image=block_state.image,
            height=height,
            width=width,
        )

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageEditProcessImagesInputStep(ModularPipelineBlocks):
    """
    Image Preprocess step. Images needs to be resized first.

      Components:
          image_processor (`VaeImageProcessor`)

      Inputs:
          resized_image (`list`):
              The resized image. should be generated using a resize step

      Outputs:
          processed_image (`Tensor`):
              The processed image
    """

    model_name = "qwenimage-edit"

    @property
    def description(self) -> str:
        return "Image Preprocess step. Images needs to be resized first."

    @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="The resized image. should be generated using a resize step",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="processed_image",
                type_hint=torch.Tensor,
                description="The processed image",
            )
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        width, height = block_state.resized_image[0].size

        block_state.processed_image = components.image_processor.preprocess(
            image=block_state.resized_image,
            height=height,
            width=width,
        )

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class QwenImageEditPlusProcessImagesInputStep(ModularPipelineBlocks):
    """
    Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of
    processed images.

      Components:
          image_processor (`VaeImageProcessor`)

      Inputs:
          resized_image (`list`):
              The resized image. should be generated using a resize step

      Outputs:
          processed_image (`Tensor`):
              The processed image
    """

    model_name = "qwenimage-edit-plus"

    @property
    def description(self) -> str:
        return "Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of processed images."

    @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="The resized image. should be generated using a resize step",
            )
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="processed_image",
                type_hint=torch.Tensor,
                description="The processed image",
            )
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        image = block_state.resized_image

        is_image_list = isinstance(image, list)
        if not is_image_list:
            image = [image]

        processed_images = []
        for img in image:
            img_width, img_height = img.size
            processed_images.append(
                components.image_processor.preprocess(image=img, height=img_height, width=img_width)
            )

        if is_image_list:
            block_state.processed_image = processed_images
        else:
            block_state.processed_image = processed_images[0]

        self.set_block_state(state, block_state)
        return components, state


# ====================
# 5. VAE ENCODER
# ====================


# auto_docstring
class QwenImageVaeEncoderStep(ModularPipelineBlocks):
    """
    VAE Encoder step that converts processed_image into latent representations image_latents.
      Handles both single images and lists of images with varied resolutions.

      Components:
          vae (`AutoencoderKLQwenImage`)

      Inputs:
          processed_image (`Tensor`):
              The image tensor to encode
          generator (`Generator`, *optional*):
              Torch generator for deterministic generation.

      Outputs:
          image_latents (`Tensor`):
              The latent representation of the input image.
    """

    model_name = "qwenimage"

    def __init__(self, input: InputParam | None = None, output: OutputParam | None = None):
        """Initialize a VAE encoder step for converting images to latent representations.

        Handles both single images and lists of images. When input is a list, outputs a list of latents. When input is
        a single tensor, outputs a single latent tensor.

        Args:
            input (InputParam, optional): Input parameter for the processed image. Defaults to "processed_image".
            output (OutputParam, optional): Output parameter for the image latents. Defaults to "image_latents".
        """
        if input is None:
            input = InputParam(
                name="processed_image", required=True, type_hint=torch.Tensor, description="The image tensor to encode"
            )

        if output is None:
            output = OutputParam.template("image_latents")

        if not isinstance(input, InputParam):
            raise ValueError(f"input must be InputParam but is {type(input)}")
        if not isinstance(output, OutputParam):
            raise ValueError(f"output must be OutputParam but is {type(output)}")

        self._input = input
        self._output = output
        self._image_input_name = input.name
        self._image_latents_output_name = output.name
        super().__init__()

    @property
    def description(self) -> str:
        return (
            f"VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n"
            "Handles both single images and lists of images with varied resolutions."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("vae", AutoencoderKLQwenImage)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            self._input,  # default is "processed_image"
            InputParam.template("generator"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [self._output]  # default is "image_latents"

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        device = components._execution_device
        dtype = components.vae.dtype

        image = getattr(block_state, self._image_input_name)
        is_image_list = isinstance(image, list)
        if not is_image_list:
            image = [image]

        # Handle both single image and list of images
        image_latents = []
        for img in image:
            image_latents.append(
                encode_vae_image(
                    image=img,
                    vae=components.vae,
                    generator=block_state.generator,
                    device=device,
                    dtype=dtype,
                    latent_channels=components.num_channels_latents,
                )
            )
        if not is_image_list:
            image_latents = image_latents[0]

        setattr(block_state, self._image_latents_output_name, image_latents)

        self.set_block_state(state, block_state)

        return components, state


# auto_docstring
class QwenImageControlNetVaeEncoderStep(ModularPipelineBlocks):
    """
    VAE Encoder step that converts `control_image` into latent representations control_image_latents.

      Components:
          vae (`AutoencoderKLQwenImage`) controlnet (`QwenImageControlNetModel`) control_image_processor
          (`VaeImageProcessor`)

      Inputs:
          control_image (`Image`):
              Control image for ControlNet conditioning.
          height (`int`, *optional*):
              The height in pixels of the generated image.
          width (`int`, *optional*):
              The width in pixels of the generated image.
          generator (`Generator`, *optional*):
              Torch generator for deterministic generation.

      Outputs:
          control_image_latents (`Tensor`):
              The latents representing the control image
    """

    model_name = "qwenimage"

    @property
    def description(self) -> str:
        return "VAE Encoder step that converts `control_image` into latent representations control_image_latents.\n"

    @property
    def expected_components(self) -> list[ComponentSpec]:
        components = [
            ComponentSpec("vae", AutoencoderKLQwenImage),
            ComponentSpec("controlnet", QwenImageControlNetModel),
            ComponentSpec(
                "control_image_processor",
                VaeImageProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]
        return components

    @property
    def inputs(self) -> list[InputParam]:
        inputs = [
            InputParam.template("control_image"),
            InputParam.template("height"),
            InputParam.template("width"),
            InputParam.template("generator"),
        ]
        return inputs

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "control_image_latents",
                type_hint=torch.Tensor,
                description="The latents representing the control image",
            )
        ]

    @staticmethod
    def check_inputs(height, width, vae_scale_factor):
        if height is not None and height % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")

        if width is not None and width % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        self.check_inputs(block_state.height, block_state.width, components.vae_scale_factor)

        device = components._execution_device
        dtype = components.vae.dtype

        height = block_state.height or components.default_height
        width = block_state.width or components.default_width

        controlnet = unwrap_module(components.controlnet)
        if isinstance(controlnet, QwenImageMultiControlNetModel) and not isinstance(block_state.control_image, list):
            block_state.control_image = [block_state.control_image]

        if isinstance(controlnet, QwenImageMultiControlNetModel):
            block_state.control_image_latents = []
            for control_image_ in block_state.control_image:
                control_image_ = components.control_image_processor.preprocess(
                    image=control_image_,
                    height=height,
                    width=width,
                )

                control_image_latents_ = encode_vae_image(
                    image=control_image_,
                    vae=components.vae,
                    generator=block_state.generator,
                    device=device,
                    dtype=dtype,
                    latent_channels=components.num_channels_latents,
                    sample_mode="sample",
                )
                block_state.control_image_latents.append(control_image_latents_)

        elif isinstance(controlnet, QwenImageControlNetModel):
            control_image = components.control_image_processor.preprocess(
                image=block_state.control_image,
                height=height,
                width=width,
            )
            block_state.control_image_latents = encode_vae_image(
                image=control_image,
                vae=components.vae,
                generator=block_state.generator,
                device=device,
                dtype=dtype,
                latent_channels=components.num_channels_latents,
                sample_mode="sample",
            )

        else:
            raise ValueError(
                f"Expected controlnet to be a QwenImageControlNetModel or QwenImageMultiControlNetModel, got {type(controlnet)}"
            )

        self.set_block_state(state, block_state)

        return components, state


# ====================
# 6. PERMUTE LATENTS
# ====================


# auto_docstring
class QwenImageLayeredPermuteLatentsStep(ModularPipelineBlocks):
    """
    Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing.

      Inputs:
          image_latents (`Tensor`):
              image latents used to guide the image generation. Can be generated from vae_encoder step.

      Outputs:
          image_latents (`Tensor`):
              The latent representation of the input image. (permuted from [B, C, 1, H, W] to [B, 1, C, H, W])
    """

    model_name = "qwenimage-layered"

    @property
    def description(self) -> str:
        return "Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing."

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("image_latents"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam.template("image_latents", note="permuted from [B, C, 1, H, W] to [B, 1, C, H, W]"),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        # Permute: (B, C, 1, H, W) -> (B, 1, C, H, W)
        latents = block_state.image_latents
        block_state.image_latents = latents.permute(0, 2, 1, 3, 4)

        self.set_block_state(state, block_state)
        return components, state