qwenimage/pipeline_qwenimage_edit_plus.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.

import inspect
import math
from typing import Any, Callable, Dict, List, Optional, Union

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageOps
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor

from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.loaders import QwenImageLoraLoaderMixin
from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput

if is_torch_xla_available():
    import torch_xla.core.xla_model as xm

    XLA_AVAILABLE = True
else:
    XLA_AVAILABLE = False

logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
>>> from diffusers import QwenImageEditPlusPipeline
>>> from diffusers.utils import load_image

>>> pipe = QwenImageEditPlusPipeline.from_pretrained(
...     "Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16
... ).to("cuda")

>>> image = load_image(
... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
... ).convert("RGB")

>>> prompt = "Make Pikachu hold a sign that says 'Qwen Edit is awesome', yarn art style, detailed, vibrant colors"

>>> out = pipe(image=image, prompt=prompt, num_inference_steps=50).images[0]
>>> out.save("qwenimage_edit_plus.png")
```
"""

CONDITION_IMAGE_SIZE = 384 * 384
VAE_IMAGE_SIZE = 1024 * 1024


def pad_to_aspect(img: Image.Image, target_w: int, target_h: int) -> Image.Image:
    """Pad (letterbox) to target aspect ratio without warping."""
    return ImageOps.pad(
        img.convert("RGB"),
        (int(target_w), int(target_h)),
        method=Image.Resampling.LANCZOS,
        color=(0, 0, 0),
        centering=(0.5, 0.5),
    )


def choose_condition_area(canvas_area: int, base_area: int = CONDITION_IMAGE_SIZE) -> int:
    """Choose a conditioning target area derived from canvas area with sensible bounds."""
    scaled = int(canvas_area * (base_area / (1024 * 1024)))
    return int(min(base_area, max(256 * 256, scaled)))


# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
def calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
):
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    mu = image_seq_len * m + b
    return mu


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
    scheduler,
    num_inference_steps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
    timesteps: Optional[List[int]] = None,
    sigmas: Optional[List[float]] = None,
    **kwargs,
):
    if timesteps is not None and sigmas is not None:
        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one.")

    if timesteps is not None:
        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accepts_timesteps:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom timesteps."
            )
        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)

    elif sigmas is not None:
        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accept_sigmas:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom sigmas."
            )
        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)

    else:
        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
        timesteps = scheduler.timesteps

    return timesteps, num_inference_steps


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"):
    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
        return encoder_output.latent_dist.sample(generator)
    if hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
        return encoder_output.latent_dist.mode()
    if hasattr(encoder_output, "latents"):
        return encoder_output.latents
    raise AttributeError("Could not access latents of provided encoder_output")


def calculate_dimensions(target_area: int, ratio: float, multiple: int = 32):
    """
    Area-based sizing while snapping to a chosen lattice multiple.
    Used for canvas sizing AND conditioning sizing (anti-drift).
    """
    m = int(multiple) if multiple else 32
    m = max(1, m)

    width = math.sqrt(float(target_area) * float(ratio))
    height = width / float(ratio)

    width = round(width / m) * m
    height = round(height / m) * m
    return int(width), int(height)


# Optional: decoder VAE (Wan2x)
_ALT_VAE_WAN2X = None

# Track desired tiling state for the optional decoder VAE, so it stays consistent across lazy loads.
_ALT_VAE_WAN2X_TILING_ENABLED = False


def _set_vae_tiling(model: Any, enabled: bool) -> bool:
    """
    Best-effort tiling toggle for a VAE-like module.
    Returns True if a tiling method existed and was called, False otherwise.
    """
    if model is None:
        return False
    try:
        if enabled:
            if hasattr(model, "enable_tiling"):
                model.enable_tiling()
                return True
            if hasattr(model, "enable_vae_tiling"):
                model.enable_vae_tiling()
                return True
        else:
            if hasattr(model, "disable_tiling"):
                model.disable_tiling()
                return True
            if hasattr(model, "disable_vae_tiling"):
                model.disable_vae_tiling()
                return True
    except Exception as e:
        # Don't hard-fail inference if tiling toggle fails for an alt decoder.
        logger.warning(f"VAE tiling toggle failed on {type(model)}: {e}")
        return False
    return False



def _get_wan2x_vae(device: torch.device, dtype: torch.dtype):
    """
    Decoder-only finetune that outputs 2x resolution via pixel-shuffle.
    Lazy-loaded so it doesn't impact startup unless used.
    """
    global _ALT_VAE_WAN2X, _ALT_VAE_WAN2X_TILING_ENABLED
    if _ALT_VAE_WAN2X is None:
        from diffusers import AutoencoderKLWan

        _ALT_VAE_WAN2X = AutoencoderKLWan.from_pretrained(
            "spacepxl/Wan2.1-VAE-upscale2x",
            subfolder="diffusers/Wan2.1_VAE_upscale2x_imageonly_real_v1",
            torch_dtype=dtype,
        )
        _ALT_VAE_WAN2X.eval()

        # Apply last requested tiling immediately on first load (if supported).
        _set_vae_tiling(_ALT_VAE_WAN2X, _ALT_VAE_WAN2X_TILING_ENABLED)

    _ALT_VAE_WAN2X = _ALT_VAE_WAN2X.to(device=device, dtype=dtype)

    # Re-apply after moving to device, just in case.
    _set_vae_tiling(_ALT_VAE_WAN2X, _ALT_VAE_WAN2X_TILING_ENABLED)

    return _ALT_VAE_WAN2X


class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
    r"""
    The Qwen-Image-Edit pipeline for image editing.
    """

    model_cpu_offload_seq = "text_encoder->transformer->vae"
    _callback_tensor_inputs = ["latents", "prompt_embeds"]

    def __init__(
        self,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: AutoencoderKLQwenImage,
        text_encoder: Qwen2_5_VLForConditionalGeneration,
        tokenizer: Qwen2Tokenizer,
        processor: Qwen2VLProcessor,
        transformer: QwenImageTransformer2DModel,
    ):
        super().__init__()
        self.register_modules(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            processor=processor,
            transformer=transformer,
            scheduler=scheduler,
        )

        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
        self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16

        # QwenImage latents are turned into 2x2 patches and packed; multiply scale-factor by patch size
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
        self.tokenizer_max_length = 1024


        # Track tiling state (applies to both primary VAE and optional decoder VAE)
        self._vae_tiling_enabled = False
        self.prompt_template_encode = (
            "<|im_start|>system\n"
            "Describe the key features of the input image (color, shape, size, texture, objects, background), "
            "then explain how the user's text instruction should alter or modify the image.\n"
            "Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate."
            "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
        )
        self.prompt_template_encode_start_idx = 64
        self.default_sample_size = 128


        # ------------------------------------------------------------
        # VAE tiling control (applies to both primary VAE and optional decoder VAE)
        # ------------------------------------------------------------
        # Expose a stable API so app.py can call pipe.enable_vae_tiling()/disable_vae_tiling()
        # regardless of which decoder VAE is selected at runtime.
        
    def set_vae_tiling(self, enabled: bool) -> None:
        global _ALT_VAE_WAN2X_TILING_ENABLED, _ALT_VAE_WAN2X

        enabled = bool(enabled)
        self._vae_tiling_enabled = enabled

        # 1) Primary VAE (Qwen)
        _set_vae_tiling(getattr(self, "vae", None), enabled)

        # 2) Optional decoder VAE (Wan2x): store desired global state; apply now if already loaded.
        _ALT_VAE_WAN2X_TILING_ENABLED = enabled
        if _ALT_VAE_WAN2X is not None:
            _set_vae_tiling(_ALT_VAE_WAN2X, enabled)

    def enable_vae_tiling(self) -> None:
        self.set_vae_tiling(True)

    def disable_vae_tiling(self) -> None:
        self.set_vae_tiling(False)
    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
        bool_mask = mask.bool()
        valid_lengths = bool_mask.sum(dim=1)
        selected = hidden_states[bool_mask]
        split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
        return split_result

    def _get_qwen_prompt_embeds(
        self,
        prompt: Union[str, List[str]] = None,
        image: Optional[torch.Tensor] = None,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        device = device or self._execution_device
        dtype = dtype or self.text_encoder.dtype

        prompt = [prompt] if isinstance(prompt, str) else prompt
        img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"

        if isinstance(image, list):
            base_img_prompt = ""
            for i, _ in enumerate(image):
                base_img_prompt += img_prompt_template.format(i + 1)
        elif image is not None:
            base_img_prompt = img_prompt_template.format(1)
        else:
            base_img_prompt = ""

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

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

        outputs = self.text_encoder(
            input_ids=model_inputs.input_ids,
            attention_mask=model_inputs.attention_mask,
            pixel_values=model_inputs.pixel_values,
            image_grid_thw=model_inputs.image_grid_thw,
            output_hidden_states=True,
        )

        hidden_states = outputs.hidden_states[-1]
        split_hidden_states = self._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(dtype=dtype, device=device)
        return prompt_embeds, encoder_attention_mask

    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
    def encode_prompt(
        self,
        prompt: Union[str, List[str]],
        image: Optional[torch.Tensor] = None,
        device: Optional[torch.device] = None,
        num_images_per_prompt: int = 1,
        prompt_embeds: Optional[torch.Tensor] = None,
        prompt_embeds_mask: Optional[torch.Tensor] = None,
        max_sequence_length: int = 1024,
    ):
        device = device or self._execution_device
        prompt = [prompt] if isinstance(prompt, str) else prompt
        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]

        if prompt_embeds is None:
            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)

        _, seq_len, _ = prompt_embeds.shape

        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)

        return prompt_embeds, prompt_embeds_mask

    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs
    def check_inputs(
        self,
        prompt,
        height,
        width,
        negative_prompt=None,
        prompt_embeds=None,
        negative_prompt_embeds=None,
        prompt_embeds_mask=None,
        negative_prompt_embeds_mask=None,
        callback_on_step_end_tensor_inputs=None,
        max_sequence_length=None,
    ):
        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
            logger.warning(
                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. "
                "Dimensions will be resized accordingly."
            )

        if callback_on_step_end_tensor_inputs is not None and not all(
            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
        ):
            raise ValueError(
                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found "
                f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
            )

        if prompt is not None and prompt_embeds is not None:
            raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.")
        if prompt is None and prompt_embeds is None:
            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
        if prompt is not None and (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 negative_prompt_embeds is not None:
            raise ValueError("Cannot forward both `negative_prompt` and `negative_prompt_embeds`.")

        if prompt_embeds is not None and prompt_embeds_mask is None:
            raise ValueError("If `prompt_embeds` are provided, `prompt_embeds_mask` must also be passed.")

        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
            raise ValueError("If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` must also be passed.")

        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}")

    @staticmethod
    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
        latents = latents.permute(0, 2, 4, 1, 3, 5)
        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
        return latents

    @staticmethod
    def _unpack_latents(latents, height, width, vae_scale_factor):
        batch_size, _, channels = latents.shape
        height = 2 * (int(height) // (vae_scale_factor * 2))
        width = 2 * (int(width) // (vae_scale_factor * 2))
        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
        latents = latents.permute(0, 3, 1, 4, 2, 5)
        latents = latents.reshape(batch_size, channels // 4, 1, height, width)
        return latents

    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
        if isinstance(generator, list):
            image_latents = [
                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
                for i in range(image.shape[0])
            ]
            image_latents = torch.cat(image_latents, dim=0)
        else:
            image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")

        latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.latent_channels, 1, 1, 1).to(
            image_latents.device, image_latents.dtype
        )
        latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.latent_channels, 1, 1, 1).to(
            image_latents.device, image_latents.dtype
        )
        image_latents = (image_latents - latents_mean) / latents_std
        return image_latents

    def prepare_latents(
        self,
        images,
        batch_size,
        num_channels_latents,
        height,
        width,
        dtype,
        device,
        generator,
        latents=None,
    ):
        height = 2 * (int(height) // (self.vae_scale_factor * 2))
        width = 2 * (int(width) // (self.vae_scale_factor * 2))
        shape = (batch_size, 1, num_channels_latents, height, width)

        image_latents = None
        if images is not None:
            if not isinstance(images, list):
                images = [images]
            all_image_latents = []

            for image in images:
                image = image.to(device=device, dtype=dtype)
                if image.shape[1] != self.latent_channels:
                    image_latents = self._encode_vae_image(image=image, generator=generator)
                else:
                    image_latents = image

                if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
                    additional_image_per_prompt = batch_size // image_latents.shape[0]
                    image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
                elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
                    raise ValueError(
                        f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
                    )

                image_latent_height, image_latent_width = image_latents.shape[3:]
                image_latents = self._pack_latents(
                    image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
                )
                all_image_latents.append(image_latents)

            image_latents = torch.cat(all_image_latents, dim=1)

        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You passed a list of generators of length {len(generator)}, but requested an effective batch size of {batch_size}."
            )

        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
        else:
            latents = latents.to(device=device, dtype=dtype)

        return latents, image_latents

    @property
    def guidance_scale(self):
        return self._guidance_scale

    @property
    def attention_kwargs(self):
        return self._attention_kwargs

    @property
    def num_timesteps(self):
        return self._num_timesteps

    @property
    def current_timestep(self):
        return self._current_timestep

    @property
    def interrupt(self):
        return self._interrupt

    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
        self,
        image: Optional[PipelineImageInput] = None,
        prompt: Union[str, List[str]] = None,
        negative_prompt: Union[str, List[str]] = None,
        true_cfg_scale: float = 4.0,
        height: Optional[int] = None,
        width: Optional[int] = None,
        condition_area: Optional[int] = None,
        vae_image_indices: Optional[List[int]] = None,
        pad_to_canvas: bool = True,
        # NEW: lattice + VAE ref override
        resolution_multiple: Optional[int] = None,
        vae_ref_area: Optional[int] = None,
        vae_ref_start_index: int = 2,
        # Optional: decoder swap
        decoder_vae: str = "qwen",  # "qwen" | "wan2x"
        keep_decoder_2x: bool = False,
        # standard args
        num_inference_steps: int = 50,
        sigmas: Optional[List[float]] = None,
        guidance_scale: Optional[float] = None,
        num_images_per_prompt: int = 1,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        prompt_embeds_mask: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
        max_sequence_length: int = 512,
    ):
        """Run Qwen-Image-Edit inference.

        Examples:
        """
        # ---- determine input size ----
        if isinstance(image, list):
            image_size = image[0].size
        else:
            image_size = image.size

        # Lattice multiple used throughout (canvas sizing + condition sizing)
        multiple_of = int(resolution_multiple) if resolution_multiple is not None else (self.vae_scale_factor * 2)
        multiple_of = max(1, multiple_of)

        calculated_width, calculated_height = calculate_dimensions(
            1024 * 1024, float(image_size[0]) / float(image_size[1]), multiple=multiple_of
        )
        height = height or calculated_height
        width = width or calculated_width

        width = (int(width) // multiple_of) * multiple_of
        height = (int(height) // multiple_of) * multiple_of

        # ---- validate ----
        self.check_inputs(
            prompt,
            height,
            width,
            negative_prompt=negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            prompt_embeds_mask=prompt_embeds_mask,
            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
            max_sequence_length=max_sequence_length,
        )

        self._guidance_scale = guidance_scale
        self._attention_kwargs = attention_kwargs
        self._current_timestep = None
        self._interrupt = False

        # ---- call params ----
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device

        # ---- preprocess ----
        condition_images = None
        vae_images = None
        vae_image_sizes: List[tuple[int, int]] = []

        # support pre-latent tensors (rare, but keep compatibility)
        if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
            if not isinstance(image, list):
                image = [image]

            canvas_area = int(width) * int(height)
            cond_area = int(condition_area) if condition_area is not None else choose_condition_area(canvas_area)

            cond_w, cond_h = calculate_dimensions(cond_area, float(width) / float(height), multiple=multiple_of)

            # Optional VAE ref override sizing (applied only to indices >= vae_ref_start_index)
            ref_w = ref_h = None
            if vae_ref_area is not None:
                try:
                    ref_w, ref_h = calculate_dimensions(
                        int(vae_ref_area),
                        float(width) / float(height),
                        multiple=multiple_of,
                    )
                except Exception:
                    ref_w = ref_h = None

            condition_images = []
            vae_images = []

            if vae_image_indices is None:
                vae_image_indices = list(range(len(image)))
            vae_set = set(int(i) for i in vae_image_indices)

            for idx, img in enumerate(image):
                pil = img.convert("RGB") if isinstance(img, Image.Image) else img

                if pad_to_canvas and isinstance(pil, Image.Image):
                    pil = pad_to_aspect(pil, int(width), int(height))

                # conditioning stream (always)
                condition_images.append(self.image_processor.resize(pil, cond_h, cond_w))

                # VAE stream (selective)
                if idx in vae_set:
                    if (ref_w is not None) and (ref_h is not None) and (int(idx) >= int(vae_ref_start_index)):
                        vw, vh = int(ref_w), int(ref_h)
                    else:
                        vw, vh = int(width), int(height)

                    vae_image_sizes.append((vw, vh))
                    vae_images.append(self.image_processor.preprocess(pil, int(vh), int(vw)).unsqueeze(2))

            has_neg_prompt = negative_prompt is not None or (
                negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
            )
            if true_cfg_scale > 1 and not has_neg_prompt:
                logger.warning(
                    f"true_cfg_scale={true_cfg_scale} but CFG disabled because no negative prompt was provided."
                )
            if true_cfg_scale <= 1 and has_neg_prompt:
                logger.warning("negative_prompt provided but CFG disabled because true_cfg_scale <= 1")

            do_true_cfg = (true_cfg_scale > 1) and has_neg_prompt

            prompt_embeds, prompt_embeds_mask = self.encode_prompt(
                image=condition_images,
                prompt=prompt,
                prompt_embeds=prompt_embeds,
                prompt_embeds_mask=prompt_embeds_mask,
                device=device,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
            )

            if do_true_cfg:
                negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
                    image=condition_images,
                    prompt=negative_prompt,
                    prompt_embeds=negative_prompt_embeds,
                    prompt_embeds_mask=negative_prompt_embeds_mask,
                    device=device,
                    num_images_per_prompt=num_images_per_prompt,
                    max_sequence_length=max_sequence_length,
                )

            # ---- prepare latents ----
            num_channels_latents = self.transformer.config.in_channels // 4
            latents, image_latents = self.prepare_latents(
                vae_images,
                batch_size * num_images_per_prompt,
                num_channels_latents,
                height,
                width,
                prompt_embeds.dtype,
                device,
                generator,
                latents,
            )

            img_shapes = [
                [
                    (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2),
                    *[
                        (1, vae_h // self.vae_scale_factor // 2, vae_w // self.vae_scale_factor // 2)
                        for (vae_w, vae_h) in vae_image_sizes
                    ],
                ]
            ] * batch_size

        else:
            raise ValueError(
                "This Space pipeline expects `image` as PIL/np inputs (not pre-latents) in this setup."
            )

        # ---- timesteps ----
        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas

        image_seq_len = latents.shape[1]
        mu = calculate_shift(
            image_seq_len,
            self.scheduler.config.get("base_image_seq_len", 256),
            self.scheduler.config.get("max_image_seq_len", 4096),
            self.scheduler.config.get("base_shift", 0.5),
            self.scheduler.config.get("max_shift", 1.15),
        )
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu
        )

        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
        self._num_timesteps = len(timesteps)

        # guidance-distilled models need explicit guidance input
        if self.transformer.config.guidance_embeds and guidance_scale is None:
            raise ValueError("guidance_scale is required for guidance-distilled model.")
        if self.transformer.config.guidance_embeds:
            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0])
        else:
            if guidance_scale is not None:
                logger.warning("guidance_scale passed but ignored since model is not guidance-distilled.")
            guidance = None

        if self.attention_kwargs is None:
            self._attention_kwargs = {}

        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
        image_rotary_emb = self.transformer.pos_embed(img_shapes, txt_seq_lens, device=latents.device)

        do_true_cfg = (
            (true_cfg_scale > 1)
            and (negative_prompt_embeds is not None)
            and (negative_prompt_embeds_mask is not None)
        )
        if do_true_cfg:
            negative_txt_seq_lens = negative_prompt_embeds_mask.sum(dim=1).tolist()
            uncond_image_rotary_emb = self.transformer.pos_embed(img_shapes, negative_txt_seq_lens, device=latents.device)
        else:
            uncond_image_rotary_emb = None

        # ---- denoise ----
        self.scheduler.set_begin_index(0)
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
                    continue
                self._current_timestep = t

                latent_model_input = latents
                if image_latents is not None:
                    latent_model_input = torch.cat([latents, image_latents], dim=1)

                timestep = t.expand(latents.shape[0]).to(latents.dtype)

                with self.transformer.cache_context("cond"):
                    noise_pred = self.transformer(
                        hidden_states=latent_model_input,
                        timestep=timestep / 1000,
                        guidance=guidance,
                        encoder_hidden_states_mask=prompt_embeds_mask,
                        encoder_hidden_states=prompt_embeds,
                        image_rotary_emb=image_rotary_emb,
                        attention_kwargs=self.attention_kwargs,
                        return_dict=False,
                    )[0]
                noise_pred = noise_pred[:, : latents.size(1)]

                if do_true_cfg:
                    with self.transformer.cache_context("uncond"):
                        neg_noise_pred = self.transformer(
                            hidden_states=latent_model_input,
                            timestep=timestep / 1000,
                            guidance=guidance,
                            encoder_hidden_states_mask=negative_prompt_embeds_mask,
                            encoder_hidden_states=negative_prompt_embeds,
                            image_rotary_emb=uncond_image_rotary_emb,
                            attention_kwargs=self.attention_kwargs,
                            return_dict=False,
                        )[0]
                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]

                    comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
                    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
                    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
                    noise_pred = comb_pred * (cond_norm / (noise_norm + 1e-8))

                latents_dtype = latents.dtype
                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
                if latents.dtype != latents_dtype and torch.backends.mps.is_available():
                    latents = latents.to(latents_dtype)

                if callback_on_step_end is not None:
                    callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
                    latents = callback_outputs.pop("latents", latents)
                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)

                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()

                if XLA_AVAILABLE:
                    xm.mark_step()

        self._current_timestep = None

        # ---- decode ----
        if output_type == "latent":
            image_out = latents
        else:
            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
            latents = latents.to(self.vae.dtype)

            latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(
                latents.device, latents.dtype
            )
            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
                latents.device, latents.dtype
            )
            latents = latents / latents_std + latents_mean

            if decoder_vae == "wan2x":
                alt_vae = _get_wan2x_vae(latents.device, self.vae.dtype)
                decoder_out = alt_vae.decode(latents, return_dict=False)[0]  # [B, 12, F, H, W]
                img_2x = F.pixel_shuffle(decoder_out[:, :, 0], upscale_factor=2)  # [B, 3, 2H, 2W]
                if keep_decoder_2x:
                    decoded = img_2x
                else:
                    decoded = F.interpolate(img_2x, size=(int(height), int(width)), mode="area")
            else:
                decoded = self.vae.decode(latents, return_dict=False)[0][:, :, 0]

            image_out = self.image_processor.postprocess(decoded, output_type=output_type)

        self.maybe_free_model_hooks()

        if not return_dict:
            return (image_out,)
        return QwenImagePipelineOutput(images=image_out)