diffusers_local/pipeline_z_image_control_unified.py

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


import inspect
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as T
from diffusers import AutoencoderKL, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.loaders import FromSingleFileMixin, ZImageLoraLoaderMixin
from diffusers.pipelines.z_image.pipeline_output import ZImagePipelineOutput
from diffusers.utils import logging
from diffusers.utils.torch_utils import randn_tensor
from PIL import Image, ImageFilter
from transformers import AutoTokenizer, PreTrainedModel

from diffusers_local.z_image_control_transformer_2d import ZImageControlTransformer2DModel


logger = logging.get_logger(__name__)


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,
):
    """
    Calculates the shift value `mu` for the scheduler based on the image sequence length.

    This function implements a linear interpolation to determine the shift value based on the input
    image's sequence length, scaling between a base and a maximum shift value.

    Args:
        image_seq_len (`int`):
            The sequence length of the image latents (height * width).
        base_seq_len (`int`, *optional*, defaults to 256):
            The base sequence length for the shift calculation.
        max_seq_len (`int`, *optional*, defaults to 4096):
            The maximum sequence length for the shift calculation.
        base_shift (`float`, *optional*, defaults to 0.5):
            The shift value corresponding to `base_seq_len`.
        max_shift (`float`, *optional*, defaults to 1.15):
            The shift value corresponding to `max_seq_len`.

    Returns:
        `float`: The calculated shift value `mu`.
    """
    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


def retrieve_latents(encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"):
    """
    Retrieves latents from a VAE encoder output.

    Args:
        encoder_output (`torch.Tensor`):
            The output of a VAE encoder.
        generator (`torch.Generator`, *optional*):
            A random number generator for sampling from the latent distribution.
        sample_mode (`str`, *optional*, defaults to "sample"):
            The method to retrieve latents. Can be "sample" to sample from the distribution or
            "argmax" to take the mode.

    Returns:
        `torch.Tensor`: The retrieved latents.
    """
    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")


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,
):
    """
    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.

    Args:
        scheduler (`SchedulerMixin`):
            The scheduler to get timesteps from.
        num_inference_steps (`int`):
            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
            must be `None`.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        timesteps (`List[int]`, *optional*):
            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
            `num_inference_steps` and `sigmas` must be `None`.
        sigmas (`List[float]`, *optional*):
            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
            `num_inference_steps` and `timesteps` must be `None`.

    Returns:
        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
        second element is the number of inference steps.
    """
    if timesteps is not None and sigmas is not None:
        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
    if timesteps is not None:
        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accepts_timesteps:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" timestep schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    elif sigmas is not None:
        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accept_sigmas:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" sigmas schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    else:
        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
        timesteps = scheduler.timesteps
    return timesteps, num_inference_steps


class ZImageControlUnifiedPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
    model_cpu_offload_seq = "text_encoder->vae->transformer"
    _optional_components = []
    _callback_tensor_inputs = ["latents", "prompt_embeds"]

    def __init__(
        self,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: AutoencoderKL,
        text_encoder: PreTrainedModel,
        tokenizer: AutoTokenizer,
        transformer: ZImageControlTransformer2DModel,
    ):
        """
        Initializes the ZImageControlUnifiedPipeline.

        Args:
            scheduler (`FlowMatchEulerDiscreteScheduler`):
                A scheduler to be used in combination with `transformer` to denoise the latents.
            vae (`AutoencoderKL`):
                Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
            text_encoder (`PreTrainedModel`):
                A pretrained text encoder model.
            tokenizer (`AutoTokenizer`):
                A tokenizer to prepare text prompts for the `text_encoder`.
            transformer (`ZImageControlTransformer2DModel`):
                The main transformer model for the diffusion process.
        """
        super().__init__()
        self.register_modules(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            scheduler=scheduler,
            transformer=transformer,
        )
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
        self.mask_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True)

    def encode_prompt(
        self,
        prompt: Union[str, List[str]],
        device: Optional[torch.device] = None,
        num_images_per_prompt: int = 1,
        do_classifier_free_guidance: bool = True,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        prompt_embeds: Optional[List[torch.FloatTensor]] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        max_sequence_length: int = 512,
    ):
        """
        Encodes the prompt into text embeddings.

        Args:
            prompt (`Union[str, List[str]]`):
                The prompt or prompts to guide the image generation.
            device (`Optional[torch.device]`):
                The device to move the embeddings to.
            num_images_per_prompt (`int`):
                The number of images to generate per prompt.
            do_classifier_free_guidance (`bool`):
                Whether to generate embeddings for classifier-free guidance.
            negative_prompt (`Optional[Union[str, List[str]]]`):
                The negative prompt or prompts.
            prompt_embeds (`Optional[List[torch.FloatTensor]]`):
                Pre-generated positive prompt embeddings.
            negative_prompt_embeds (`Optional[torch.FloatTensor]`):
                Pre-generated negative prompt embeddings.
            max_sequence_length (`int`):
                The maximum sequence length for tokenization.

        Returns:
            `Tuple[List[torch.Tensor], List[torch.Tensor]]`: A tuple containing the positive and negative prompt embeddings.
        """
        device = device or self._execution_device
        prompt = [prompt] if isinstance(prompt, str) else prompt

        if prompt_embeds is not None:
            pass
        else:
            prompt_embeds = self._encode_prompt(
                prompt=prompt,
                device=device,
                max_sequence_length=max_sequence_length,
            )
        if num_images_per_prompt > 1:
            prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)]

        if do_classifier_free_guidance:
            if negative_prompt_embeds is not None:
                pass
            else:
                if negative_prompt is None:
                    negative_prompt = [""] * len(prompt)
                else:
                    negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
                assert len(prompt) == len(negative_prompt)
                negative_prompt_embeds = self._encode_prompt(
                    prompt=negative_prompt,
                    device=device,
                    max_sequence_length=max_sequence_length,
                )

            if num_images_per_prompt > 1:
                negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]

        return prompt_embeds, negative_prompt_embeds

    def _encode_prompt(self, prompt: Union[str, List[str]], device: torch.device, max_sequence_length: int) -> List[torch.Tensor]:
        """
        Internal helper to encode a list of prompts into embeddings, applying chat templates if available.

        Args:
            prompt (`Union[str, List[str]]`):
                A list of strings to be encoded.
            device (`torch.device`):
                The target device for the embeddings.
            max_sequence_length (`int`):
                The maximum length for tokenization.

        Returns:
            `List[torch.Tensor]`: A list of embedding tensors, one for each prompt.
        """
        formatted_prompts = []
        for p in prompt:
            messages = [{"role": "user", "content": p}]
            if hasattr(self.tokenizer, "apply_chat_template"):
                formatted_prompts.append(self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True))
            else:
                formatted_prompts.append(p)

        text_inputs = self.tokenizer(
            formatted_prompts,
            padding="max_length",
            max_length=max_sequence_length,
            truncation=True,
            return_tensors="pt",
        ).to(device)

        prompt_masks = text_inputs.attention_mask.bool()

        with torch.no_grad():
            prompt_embeds_batch = self.text_encoder(input_ids=text_inputs.input_ids, attention_mask=prompt_masks, output_hidden_states=True).hidden_states[-2]

        embeddings_list = []
        for i in range(prompt_embeds_batch.shape[0]):
            embeddings_list.append(prompt_embeds_batch[i][prompt_masks[i]])

        return embeddings_list

    def get_timesteps(self, num_inference_steps, strength, device):
        """
        Calculates the timesteps for the scheduler based on the number of inference steps and strength.
        This is primarily used for image-to-image pipelines.

        Args:
            num_inference_steps (`int`): The total number of diffusion steps.
            strength (`float`): The strength of the denoising process. A value of 1.0 means full denoising.
            device (`torch.device`): The device to place the timesteps on.

        Returns:
            `Tuple[torch.Tensor, int]`: A tuple containing the timesteps and the number of steps to run.
        """
        init_timestep = min(num_inference_steps * strength, num_inference_steps)

        t_start = int(max(num_inference_steps - init_timestep, 0))
        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
        if hasattr(self.scheduler, "set_begin_index"):
            self.scheduler.set_begin_index(t_start * self.scheduler.order)

        return timesteps, num_inference_steps - t_start

    def prepare_latents(
        self,
        batch_size: int,
        num_channels_latents: int,
        height: int,
        width: int,
        dtype: torch.dtype,
        device: torch.device,
        generator: torch.Generator,
        image: Optional[PipelineImageInput] = None,
        timestep: Optional[torch.Tensor] = None,
        latents: Optional[torch.Tensor] = None,
    ):
        """
        Prepares the initial latents for the diffusion process.

        This function handles three cases:
        1. `latents` are provided: They are returned directly.
        2. `image` is None (Text-to-Image): Random noise is generated.
        3. `image` is provided (Image-to-Image): The image is encoded, and noise is added according to the timestep.

        Args:
            batch_size (`int`): The number of latents to generate.
            num_channels_latents (`int`): The number of channels in the latents.
            height (`int`): The height of the output image in pixels.
            width (`int`): The width of the output image in pixels.
            dtype (`torch.dtype`): The data type for the latents.
            device (`torch.device`): The device to create the latents on.
            generator (`torch.Generator`): A random generator for creating the initial noise.
            image (`Optional[PipelineImageInput]`): An initial image for img2img mode.
            timestep (`Optional[torch.Tensor]`): The starting timestep for adding noise in img2img mode.
            latents (`Optional[torch.Tensor]`): Pre-generated latents.

        Returns:
            `torch.Tensor`: The prepared latents.
        """
        latent_height = 2 * (int(height) // (self.vae_scale_factor * 2))
        latent_width = 2 * (int(width) // (self.vae_scale_factor * 2))
        shape = (batch_size, num_channels_latents, latent_height, latent_width)

        if latents is not None:
            return latents.to(device=device, dtype=dtype)

        if image is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
            return latents

        image_tensor = self.image_processor.preprocess(image, height=height, width=width).to(device=device, dtype=self.vae.dtype)
        with torch.no_grad():
            if image_tensor.shape[1] != num_channels_latents:
                if isinstance(generator, list):
                    image_latents = [retrieve_latents(self.vae.encode(image_tensor[i : i + 1]), generator=generator[i]) for i in range(image_tensor.shape[0])]
                    image_latents = torch.cat(image_latents, dim=0)
                else:
                    image_latents = retrieve_latents(self.vae.encode(image_tensor), generator=generator)

        image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
        image_latents = image_latents.to(dtype)

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

        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        latents = self.scheduler.scale_noise(image_latents, timestep, noise)

        return latents

    def _prepare_image_latents(
        self,
        image: PipelineImageInput,
        mask_image: PipelineImageInput,
        width: int,
        height: int,
        batch_size: int,
        num_images_per_prompt: int,
        device: torch.device,
        dtype: torch.dtype,
        do_preprocess: bool = True,
    ) -> torch.Tensor:
        """
        Generic function to encode an image into 5D latents for inpainting context.

        If `do_preprocess` is True, it processes the image (PIL/np).
        If `do_preprocess` is False, it assumes 'image' is already a ready-to-use tensor.

        Args:
            image (`PipelineImageInput`): The input image. Can be None to return zeros.
            width (`int`): The target width.
            height (`int`): The target height.
            batch_size (`int`): The prompt batch size.
            num_images_per_prompt (`int`): The number of images per prompt.
            device (`torch.device`): The target device.
            dtype (`torch.dtype`): The target data type.
            do_preprocess (`bool`): Whether to preprocess the image.

        Returns:
            `torch.Tensor`: A 5D tensor of the encoded image latents.
        """
        if image is None:
            latent_h = height // self.vae_scale_factor
            latent_w = width // self.vae_scale_factor
            shape = (batch_size * num_images_per_prompt, self.transformer.in_channels, 1, latent_h, latent_w)
            return torch.zeros(shape, device=device, dtype=dtype)

        if do_preprocess:
            image_tensor = self.image_processor.preprocess(image, height=height, width=width).to(device=device, dtype=self.vae.dtype)
        else:
            image_tensor = image.to(device=device, dtype=self.vae.dtype)

        if mask_image is not None:
            mask_condition = self.mask_processor.preprocess(mask_image, height=height, width=width).to(device=device, dtype=self.vae.dtype)
            # Tile para 3 canais (RGB)
            mask_condition = torch.tile(mask_condition, [1, 3, 1, 1])
            # Aplica máscara: mantém apenas áreas escuras (< 0.5)
            image_tensor = image_tensor * (mask_condition < 0.5)

        with torch.no_grad():
            latents = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax")
            latents = (latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor

        effective_batch_size = batch_size * num_images_per_prompt
        if latents.shape[0] != effective_batch_size:
            repeat_by = effective_batch_size // latents.shape[0]
            latents = latents.repeat_interleave(repeat_by, dim=0)

        return latents.to(dtype=dtype).unsqueeze(2)

    def _prepare_mask_latents(
        self,
        mask_image: PipelineImageInput,
        width: int,
        height: int,
        batch_size: int,
        num_images_per_prompt: int,
        reference_latents_shape: Tuple,
        device: torch.device,
        dtype: torch.dtype,
        invert_mask: bool = False,
        do_unsqueeze: bool = True,
    ) -> torch.Tensor:
        """
        Processes a MASK using the mask_processor, inverts it, resizes it, and formats it for the control_context.

        Args:
            mask_image (`PipelineImageInput`): The mask image. Can be None to return zeros.
            width (`int`): The target width.
            height (`int`): The target height.
            batch_size (`int`): The prompt batch size.
            num_images_per_prompt (`int`): The number of images per prompt.
            reference_latents_shape (`Tuple`): The shape of the inpainting latents for resizing.
            device (`torch.device`): The target device.
            dtype (`torch.dtype`): The target data type.

        Returns:
            `torch.Tensor`: A 5D tensor of the processed mask latents.
        """
        if mask_image is None:
            placeholder_shape = (
                batch_size * num_images_per_prompt,
                1,
                1,
                reference_latents_shape[-2],
                reference_latents_shape[-1],
            )
            return torch.zeros(placeholder_shape, device=device, dtype=dtype)

        mask_tensor = self.mask_processor.preprocess(mask_image, height=height, width=width)
        mask_tensor = mask_tensor.to(device=device, dtype=dtype)
    
        if invert_mask:
            mask_tensor = 1.0 - mask_tensor

        mask_latents = F.interpolate(mask_tensor, size=reference_latents_shape[-2:], mode="nearest")
        
        if do_unsqueeze:
            mask_latents = mask_latents.unsqueeze(2)
    
        return mask_latents

    def prepare_control_latents(
        self, image: PipelineImageInput, width: int, height: int, batch_size: int, num_images_per_prompt: int, device: torch.device, dtype: torch.dtype
    ) -> torch.Tensor:
        """
        Preprocesses a control image, ENCODES it with the VAE to latent space,
        and returns a 5D tensor ready for the transformer model.

        Args:
            image (`PipelineImageInput`): The control image. Can be None to return zeros.
            width (`int`): The target width.
            height (`int`): The target height.
            batch_size (`int`): The prompt batch size.
            num_images_per_prompt (`int`): The number of images per prompt.
            device (`torch.device`): The target device.
            dtype (`torch.dtype`): The target data type.

        Returns:
            `torch.Tensor`: A 5D tensor of the control image latents.
        """
        if image is None:
            latent_h = 2 * (int(height) // (self.vae_scale_factor * 2))
            latent_w = 2 * (int(width) // (self.vae_scale_factor * 2))
            return torch.zeros(
                (batch_size * num_images_per_prompt, self.transformer.in_channels, 1, latent_h, latent_w),
                device=device,
                dtype=dtype,
            )

        image_tensor = self.image_processor.preprocess(image, height=height, width=width).to(device=device, dtype=self.vae.dtype)
        with torch.no_grad():
            latents = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax")
            latents = (latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor

        effective_batch_size = batch_size * num_images_per_prompt
        if latents.shape[0] < effective_batch_size:
            latents = latents.repeat_interleave(effective_batch_size // latents.shape[0], dim=0)

        return latents.to(dtype=dtype).unsqueeze(2)

    def _expand_and_feather_mask(self, mask_image, expand_pixels=10, feather_radius=8, is_inpaint_mode=True):
        """
        Expands the white area of a mask using PyTorch for performance and then smooths its edges with Pillow.

        Args:
            mask_image (PIL.Image.Image | np.ndarray | torch.Tensor): The input mask.
            expand_pixels (int): How many pixels to expand the white area.
            feather_radius (int): The radius of the Gaussian blur for the gradient.
            is_inpaint_mode (bool): Flag to enable/disable the operation.

        Returns:
            PIL.Image.Image | np.ndarray | torch.Tensor: The processed mask, in the same format as the input.
        """
        if not is_inpaint_mode or (expand_pixels <= 0 and feather_radius <= 0):
            return mask_image

        # --- 1. CONVERSÃO PARA TENSOR PYTORCH ---
        input_type = type(mask_image)
        
        if isinstance(mask_image, Image.Image):
            # Converte PIL Image para Tensor
            mask_tensor = T.ToTensor()(mask_image.convert("L"))
        elif isinstance(mask_image, np.ndarray):
            # Converte NumPy array para Tensor
            mask_tensor = torch.from_numpy(mask_image).permute(2, 0, 1) if mask_image.ndim == 3 else torch.from_numpy(mask_image).unsqueeze(0)
        elif isinstance(mask_image, torch.Tensor):
            mask_tensor = mask_image
        else:
            raise TypeError(f"Unsupported mask type: {input_type}")

        # Garante que o tensor está no device e formato corretos (Batch, Canais, H, W)
        mask_tensor = mask_tensor.to(device=self.device, dtype=torch.float32)
        if mask_tensor.ndim == 3:
            mask_tensor = mask_tensor.unsqueeze(0) # Adiciona a dimensão do batch se necessário

        # --- 2. EXPANSÃO (DILATION) NA GPU COM PYTORCH ---
        if expand_pixels > 0:
            kernel_size = expand_pixels * 2 + 1
            padding = expand_pixels

            # Max pooling com stride=1 é a implementação de dilatação para tensores
            mask_tensor = F.max_pool2d(
                mask_tensor, 
                kernel_size=kernel_size, 
                stride=1, 
                padding=padding
            )

        # --- 3. CONVERSÃO DE VOLTA PARA PIL IMAGE ---
        # `ToPILImage` espera um tensor [C, H, W], então removemos a dimensão do batch
        to_pil = T.ToPILImage()
        mask_pil = to_pil(mask_tensor.squeeze(0).cpu())

        # --- 4. DEGRADÊ (FEATHERING / BLUR) COM PILLOW ---
        if feather_radius > 0:
            mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(radius=feather_radius))

        # --- 5. CONVERSÃO FINAL PARA O TIPO ORIGINAL ---
        if input_type is torch.Tensor:
            # Reconverte para Tensor se o input era um Tensor
            return T.ToTensor()(mask_pil).to(device=self.device, dtype=mask_image.dtype)
        elif input_type is np.ndarray:
            # Reconverte para NumPy array se o input era um array
            return np.array(mask_pil)
        else: # input_type is Image.Image
            return mask_pil

    def _apply_mask_blur(self, mask_image, mask_blur_radius, is_inpaint_mode):
        """
        Apply Gaussian blur to a mask image for inpainting operations.
        Args:
            mask_image (Image.Image | np.ndarray | torch.Tensor): The mask image to be blurred.
                Can be provided as a PIL Image, NumPy array, or PyTorch tensor.
            mask_blur_radius (float): The radius of the Gaussian blur filter in pixels.
                Only applied if is_inpaint_mode is True and mask_blur_radius > 0.
            is_inpaint_mode (bool): Flag indicating whether the pipeline is in inpainting mode.
                Blur is only applied when this is True.
        Returns:
            Image.Image | np.ndarray | torch.Tensor: The mask image with Gaussian blur applied
                if is_inpaint_mode is True and mask_blur_radius > 0. Otherwise, returns the
                original mask_image unchanged. The return type matches the input type.
        """
        mask_to_use = mask_image
        if is_inpaint_mode and mask_blur_radius > 0:
            if isinstance(mask_image, Image.Image):
                mask_pil = mask_image
            elif isinstance(mask_image, np.ndarray):
                mask_pil = Image.fromarray(mask_image)
            elif isinstance(mask_image, torch.Tensor):
                mask_pil = Image.fromarray(mask_image.cpu().numpy().astype(np.uint8))
            else:
                mask_pil = mask_image

            mask_to_use = mask_pil.filter(ImageFilter.GaussianBlur(radius=mask_blur_radius))
        return mask_to_use

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

    @property
    def do_classifier_free_guidance(self):
        return self._guidance_scale > 1

    @property
    def joint_attention_kwargs(self):
        return self._joint_attention_kwargs

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

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

    def __call__(
        self,
        prompt: Union[str, List[str]],
        image: Optional[PipelineImageInput] = None,
        mask_image: Optional[PipelineImageInput] = None,
        inpaint_mode: Literal["default", "diff", "diff+inpaint"] = "default",
        mask_blur_radius: float=8.0,
        control_image: Optional[PipelineImageInput] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 20,
        sigmas: Optional[List[float]] = None,
        strength: float = 1.0,
        guidance_scale: float = 4.0,
        cfg_normalization: bool = False,
        cfg_truncation: float = 1.0,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: int = 1,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        prompt_embeds: Optional[List[torch.FloatTensor]] = None,
        negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
        controlnet_conditioning_scale: float = 1.0,
        controlnet_refiner_conditioning_scale: float = 1.0,
        output_type: str = "pil",
        return_dict: bool = True,
        joint_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,
    ):
        r"""
        The main entry point for the Z-Image unified pipeline for generation.

        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
            image (`PipelineImageInput`, *optional*):
                The initial image for image-to-image or inpainting modes.
            mask_image (`PipelineImageInput`, *optional*):
                The mask image for inpainting. White areas are preserved, black areas are inpainted.
            inpaint_mode (`str`, *optional*, defaults to `"default"`):
                The inpainting mode. Can be "default", "diff", or "diff+inpaint". Determines how the inpainting
            process is handled.
            mask_blur_radius (`float`, *optional*, defaults to 8.0):
                The radius for blurring the edges of the inpainting mask to create a smoother transition.
            control_image (`PipelineImageInput`, *optional*):
                The conditioning image for control modes (e.g., Canny, depth).
            height (`int`, *optional*, defaults to 1024):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to 1024):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 20):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
            expense of slower inference.
            sigmas (`List[float]`, *optional*):
                Custom sigmas to use for the denoising process. If not defined, the scheduler's default behavior
            will be used.
            strength (`float`, *optional*, defaults to 1.0):
                Denoising strength for image-to-image. A value of 1.0 means the initial image is fully replaced,
            while a lower value preserves more of the original image structure. Only used in img2img mode.
            guidance_scale (`float`, *optional*, defaults to 4.0):
                The scale for classifier-free guidance. A value > 1 enables it. Higher values encourage images
            closer to the prompt, potentially at the cost of quality.
            cfg_normalization (`bool`, *optional*, defaults to False):
                Whether to apply normalization to the guidance, which can prevent oversaturation.
            cfg_truncation (`float`, *optional*, defaults to 1.0):
                A value between 0.0 and 1.0 that disables CFG for the final portion of the denoising steps,
            specified as a fraction of total steps. For example, 0.8 disables CFG for the last 20% of steps.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation.
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                A torch generator to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents to be used as inputs for image generation.
            prompt_embeds (`List[torch.FloatTensor]`, *optional*):
                Pre-generated positive text embeddings.
            negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
                Pre-generated negative text embeddings.
            controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
                The scale of the control conditioning influence.
            controlnet_refiner_conditioning_scale (`float`, *optional*, defaults to 1.0):
                The scale of the control refiner conditioning influence.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generated image. Choose between "pil" (`PIL.Image.Image`), "np.array", or "latent".
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a `ZImagePipelineOutput` instead of a plain tuple.
            joint_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary for the `AttentionProcessor`.
            callback_on_step_end (`Callable`, *optional*):
                A function that is called at the end of each denoising step.
            callback_on_step_end_tensor_inputs (`List`, *optional*):
                The list of tensor inputs for the `callback_on_step_end` function.
            max_sequence_length (`int`, *optional*, defaults to 512):
                Maximum sequence length to use with the `prompt`.

        Examples:

        Returns:
            [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`:
            If `return_dict` is True, a `ZImagePipelineOutput` is returned, otherwise a `tuple` with the generated images.
        """
        self._guidance_scale = guidance_scale
        self._joint_attention_kwargs = joint_attention_kwargs
        self._interrupt = False
        self._cfg_normalization = cfg_normalization
        self._cfg_truncation = cfg_truncation
        is_two_stage_control_model = self.transformer.control_in_dim > self.transformer.in_channels if hasattr(self.transformer, "control_in_dim") else False
        device = self._execution_device
        dtype = self.transformer.dtype
        vae_scale = self.vae_scale_factor * 2   
        has_inpaint_inputs = image is not None and mask_image is not None
        is_inpaint_control_mode = has_inpaint_inputs and inpaint_mode in ["default", "diff+inpaint"]        
        is_diff_mode = has_inpaint_inputs and inpaint_mode in ["diff", "diff+inpaint"]        
        is_img2img_mode = image is not None and not has_inpaint_inputs
        
        ref_image = control_image or image
        image_height = None
        image_width = None
        if ref_image is not None:
            if isinstance(ref_image, Image.Image):
                image_height, image_width = ref_image.height, ref_image.width
            else:
                image_height, image_width = ref_image.shape[-2], ref_image.shape[-1]

        height = height or image_height or 1024
        width = width or image_width or 1024

        if height % vae_scale != 0 or width % vae_scale != 0:
            raise ValueError(f"Height/width must be divisible by {vae_scale}.")

        batch_size = len(prompt) if isinstance(prompt, list) else 1 if prompt else len(prompt_embeds)
        effective_batch_size = batch_size * num_images_per_prompt

        if prompt_embeds is not None and prompt is None:
            if self.do_classifier_free_guidance and negative_prompt_embeds is None:
                raise ValueError(
                    "When `prompt_embeds` is provided without `prompt`, `negative_prompt_embeds` must also be provided for classifier-free guidance."
                )
        else:
            (
                prompt_embeds,
                negative_prompt_embeds,
            ) = self.encode_prompt(
                prompt=prompt,
                num_images_per_prompt=num_images_per_prompt,
                negative_prompt=negative_prompt,
                do_classifier_free_guidance=self.do_classifier_free_guidance,
                prompt_embeds=prompt_embeds,
                negative_prompt_embeds=negative_prompt_embeds,
                device=device,
                max_sequence_length=max_sequence_length,
            )

        if self.do_classifier_free_guidance:
            prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds
        else:
            prompt_embeds_model_input = prompt_embeds
  
        if control_image is not None or is_inpaint_control_mode:
            control_latents = self.prepare_control_latents(control_image, width, height, batch_size, num_images_per_prompt, device, dtype)

            if is_two_stage_control_model:                 
                image_for_inpaint = None if is_diff_mode and not is_inpaint_control_mode else image
                mask_for_inpaint = None if is_diff_mode and not is_inpaint_control_mode else mask_image                
               
                if is_inpaint_control_mode:
                    mask_for_inpaint = self._apply_mask_blur(mask_for_inpaint, mask_blur_radius, True)

                inpaint_latents = self._prepare_image_latents(
                    image_for_inpaint, mask_for_inpaint, width, height, batch_size, num_images_per_prompt, device, dtype
                )
                
                mask_latents = self._prepare_mask_latents(
                    mask_for_inpaint,
                    width,
                    height,
                    batch_size,
                    num_images_per_prompt,
                    inpaint_latents.shape,
                    device,
                    dtype,
                    invert_mask=is_inpaint_control_mode,
                    do_unsqueeze=True,
                )
                control_context = torch.cat([control_latents, mask_latents, inpaint_latents], dim=1)
            else:
                control_context = control_latents
        else:
            control_context = None

        if self.do_classifier_free_guidance:
            control_context_model_input = control_context.repeat(2, 1, 1, 1, 1)
        else:
            control_context_model_input = control_context

        image_seq_len = (height // (self.vae_scale_factor * 2)) * (width // (self.vae_scale_factor * 2))
        mu = calculate_shift(image_seq_len)
        self.scheduler.sigma_min = 0.0
        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas, mu=mu)
        self._num_timesteps = len(timesteps)

        if is_img2img_mode:
            strength = min(strength, 1.0)
        else:
            strength = 1.0

        if strength < 1.0:
            init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
            t_start = max(num_inference_steps - init_timestep, 0)
            timesteps = timesteps[t_start * self.scheduler.order :]
            num_steps_to_run = len(timesteps) // self.scheduler.order
        else:
            num_steps_to_run = num_inference_steps

        latent_timestep = timesteps[:1].repeat(effective_batch_size) if strength < 1.0 else None

        use_image_for_latents = is_img2img_mode
        
        latents = self.prepare_latents(
            effective_batch_size,
            self.transformer.in_channels,
            height,
            width,
            torch.float32,
            device,
            generator,
            image=image if use_image_for_latents else None,
            timestep=latent_timestep if use_image_for_latents else None,
            latents=latents,
        )
        
        if is_diff_mode:            
            original_image_tensor = self.image_processor.preprocess(image, height=height, width=width).to(device=device, dtype=self.vae.dtype)
            with torch.no_grad():
                original_clean_latents = retrieve_latents(self.vae.encode(original_image_tensor), sample_mode="argmax")
            original_clean_latents = (original_clean_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
            original_clean_latents = original_clean_latents.to(dtype)
            
            noise = randn_tensor(original_clean_latents.shape, generator=generator, device=device, dtype=dtype)            
            latents_list = []
            step_indices = [(self.scheduler.timesteps == t).nonzero().item() for t in timesteps]
            for i in step_indices:
                sigma = self.scheduler.sigmas[i]
                noisy_latent = (1.0 - sigma) * original_clean_latents + sigma * noise
                latents_list.append(noisy_latent)
                
            original_latents_trajectory = torch.cat(latents_list, dim=0)            
            blurred_mask_image = self._apply_mask_blur(mask_image, mask_blur_radius, True)
            map_processed = self._prepare_mask_latents(
                blurred_mask_image,
                width,
                height,
                batch_size,
                num_images_per_prompt,
                latents.shape,
                device,
                dtype,
                invert_mask=True,
                do_unsqueeze=False,
            )
          
            thresholds = torch.arange(len(timesteps), device=device, dtype=dtype) / len(timesteps)          
            thresholds = thresholds.view(-1, 1, 1, 1)
            time_masks = map_processed > thresholds
        
        num_warmup_steps = len(timesteps) - num_steps_to_run * self.scheduler.order
        with torch.inference_mode():
            with self.progress_bar(total=num_steps_to_run) as progress_bar:
                for i, t in enumerate(timesteps):
                    if self.interrupt:
                        continue
                                   
                    if is_diff_mode:
                        if i == 0:
                            latents = original_latents_trajectory[:1]
                        else:
                            current_mask = time_masks[i].to(latents.dtype)
                            current_original_latent = original_latents_trajectory[i:i+1]
                            
                            if current_mask.ndim == 3:
                                current_mask = current_mask.unsqueeze(1)
                                
                            latents = current_original_latent * current_mask + latents * (1 - current_mask)
  
                    timestep = t.expand(latents.shape[0])
                    timestep = (1000 - timestep) / 1000
                    
                    t_norm = timestep[0].item()
                    current_guidance_scale = self.guidance_scale
                    if self.do_classifier_free_guidance and self._cfg_truncation is not None and float(self._cfg_truncation) <= 1:
                        if t_norm > self._cfg_truncation:
                            current_guidance_scale = 0.0
                    apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0

                    if apply_cfg:
                        latent_model_input = latents.repeat(2, 1, 1, 1)
                        timestep_model_input = timestep.repeat(2)
                    else:
                        latent_model_input = latents
                        timestep_model_input = timestep

                    latent_model_input = latent_model_input.to(self.transformer.dtype)
                    latent_model_input = latent_model_input.unsqueeze(2)
                    latent_model_input_list = list(latent_model_input.unbind(dim=0))

                    model_out_list = self.transformer(
                        x=latent_model_input_list,
                        t=timestep_model_input,
                        cap_feats=prompt_embeds_model_input,
                        control_context=control_context_model_input,
                        conditioning_scale=controlnet_conditioning_scale,
                        refiner_conditioning_scale=controlnet_refiner_conditioning_scale,
                    )[0]

                    if apply_cfg:
                        pos_out = model_out_list[:effective_batch_size]
                        neg_out = model_out_list[effective_batch_size:]

                        noise_pred = []
                        for j in range(effective_batch_size):
                            pos = pos_out[j].float()
                            neg = neg_out[j].float()

                            pred = pos + current_guidance_scale * (pos - neg)

                            if self._cfg_normalization and float(self._cfg_normalization) > 0.0:
                                ori_pos_norm = torch.linalg.vector_norm(pos)
                                new_pos_norm = torch.linalg.vector_norm(pred)
                                max_new_norm = ori_pos_norm * float(self._cfg_normalization)
                                if new_pos_norm > max_new_norm:
                                    pred = pred * (max_new_norm / new_pos_norm)

                            noise_pred.append(pred)

                        noise_pred = torch.stack(noise_pred, dim=0)
                    else:
                        noise_pred = torch.stack([t.float() for t in model_out_list], dim=0)

                    noise_pred = noise_pred.squeeze(2)
                    noise_pred = -noise_pred

                    latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents).prev_sample

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

                        if isinstance(callback_outputs, dict):
                            latents = callback_outputs.pop("latents", latents)
                            prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                            negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)

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

        if output_type != "latent":
            latents = latents.to(self.vae.dtype)
            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
            with torch.no_grad():
                image = self.vae.decode(latents, return_dict=False)[0]
            image = self.image_processor.postprocess(image, output_type=output_type)
        else:
            image = latents

        self.maybe_free_model_hooks()

        if not return_dict:
            return (image,)

        return ZImagePipelineOutput(images=image)