pipeline_bria.py

from typing import Any, Callable

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPVisionModelWithProjection,
    T5EncoderModel,
    T5TokenizerFast,
)

from ...image_processor import VaeImageProcessor
from ...loaders import FluxLoraLoaderMixin
from ...models import AutoencoderKL
from ...models.transformers.transformer_bria import BriaTransformer2DModel
from ...pipelines import DiffusionPipeline
from ...pipelines.bria.pipeline_output import BriaPipelineOutput
from ...pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
from ...schedulers import (
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    FlowMatchEulerDiscreteScheduler,
    KarrasDiffusionSchedulers,
)
from ...utils import (
    USE_PEFT_BACKEND,
    is_torch_xla_available,
    logging,
    replace_example_docstring,
    scale_lora_layers,
    unscale_lora_layers,
)
from ...utils.torch_utils import randn_tensor


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 BriaPipeline

        >>> pipe = BriaPipeline.from_pretrained("briaai/BRIA-3.2", torch_dtype=torch.bfloat16)
        >>> pipe.to("cuda")
        # BRIA's T5 text encoder is sensitive to precision. We need to cast it to bfloat16 and keep the final layer in float32.

        >>> pipe.text_encoder = pipe.text_encoder.to(dtype=torch.bfloat16)
        >>> for block in pipe.text_encoder.encoder.block:
        ...     block.layer[-1].DenseReluDense.wo.to(dtype=torch.float32)
        # BRIA's VAE is not supported in mixed precision, so we use float32.

        >>> if pipe.vae.config.shift_factor == 0:
        ...     pipe.vae.to(dtype=torch.float32)

        >>> prompt = "Photorealistic food photography of a stack of fluffy pancakes on a white plate, with maple syrup being poured over them. On top of the pancakes are the words 'BRIA 3.2' in bold, yellow, 3D letters. The background is dark and out of focus."
        >>> image = pipe(prompt).images[0]
        >>> image.save("bria.png")
        ```
"""


def is_ng_none(negative_prompt):
    return (
        negative_prompt is None
        or negative_prompt == ""
        or (isinstance(negative_prompt, list) and negative_prompt[0] is None)
        or (type(negative_prompt) == list and negative_prompt[0] == "")
    )


def get_original_sigmas(num_train_timesteps=1000, num_inference_steps=1000):
    timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
    sigmas = timesteps / num_train_timesteps

    inds = [int(ind) for ind in np.linspace(0, num_train_timesteps - 1, num_inference_steps)]
    new_sigmas = sigmas[inds]
    return new_sigmas


class BriaPipeline(DiffusionPipeline):
    r"""
    Based on FluxPipeline with several changes:
    - no pooled embeddings
    - We use zero padding for prompts
    - No guidance embedding since this is not a distilled version

    Args:
        transformer ([`BriaTransformer2DModel`]):
            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
        vae ([`AutoencoderKL`]):
            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
        text_encoder ([`T5EncoderModel`]):
            Frozen text-encoder. Bria uses
            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
            [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
        tokenizer (`T5TokenizerFast`):
            Tokenizer of class
            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
    """

    model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae"
    _optional_components = ["image_encoder", "feature_extractor"]
    _callback_tensor_inputs = ["latents", "prompt_embeds"]

    def __init__(
        self,
        transformer: BriaTransformer2DModel,
        scheduler: FlowMatchEulerDiscreteScheduler | KarrasDiffusionSchedulers,
        vae: AutoencoderKL,
        text_encoder: T5EncoderModel,
        tokenizer: T5TokenizerFast,
        image_encoder: CLIPVisionModelWithProjection = None,
        feature_extractor: CLIPImageProcessor = None,
    ):
        self.register_modules(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            transformer=transformer,
            scheduler=scheduler,
            image_encoder=image_encoder,
            feature_extractor=feature_extractor,
        )

        self.vae_scale_factor = (
            2 ** (len(self.vae.config.block_out_channels)) if hasattr(self, "vae") and self.vae is not None else 16
        )
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
        self.default_sample_size = 64  # due to patchify=> 128,128 => res of 1k,1k

        if self.vae.config.shift_factor is None:
            self.vae.config.shift_factor = 0
            self.vae.to(dtype=torch.float32)

    def encode_prompt(
        self,
        prompt: str | list[str],
        device: torch.device | None = None,
        num_images_per_prompt: int = 1,
        do_classifier_free_guidance: bool = True,
        negative_prompt: str | list[str] | None = None,
        prompt_embeds: torch.FloatTensor | None = None,
        negative_prompt_embeds: torch.FloatTensor | None = None,
        max_sequence_length: int = 128,
        lora_scale: float | None = None,
    ):
        r"""

        Args:
            prompt (`str` or `list[str]`, *optional*):
                prompt to be encoded
            device: (`torch.device`):
                torch device
            num_images_per_prompt (`int`):
                number of images that should be generated per prompt
            do_classifier_free_guidance (`bool`):
                whether to use classifier free guidance or not
            negative_prompt (`str` or `list[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
        """
        device = device or self._execution_device

        # set lora scale so that monkey patched LoRA
        # function of text encoder can correctly access it
        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
            self._lora_scale = lora_scale

            # dynamically adjust the LoRA scale
            if self.text_encoder is not None and USE_PEFT_BACKEND:
                scale_lora_layers(self.text_encoder, lora_scale)

        prompt = [prompt] if isinstance(prompt, str) else prompt
        if prompt is not None:
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        if prompt_embeds is None:
            prompt_embeds = self._get_t5_prompt_embeds(
                prompt=prompt,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
                device=device,
            )

        if do_classifier_free_guidance and negative_prompt_embeds is None:
            if not is_ng_none(negative_prompt):
                negative_prompt = (
                    batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
                )

                if prompt is not None and type(prompt) is not type(negative_prompt):
                    raise TypeError(
                        f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
                        f" {type(prompt)}."
                    )
                elif batch_size != len(negative_prompt):
                    raise ValueError(
                        f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
                        f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                        " the batch size of `prompt`."
                    )

                negative_prompt_embeds = self._get_t5_prompt_embeds(
                    prompt=negative_prompt,
                    num_images_per_prompt=num_images_per_prompt,
                    max_sequence_length=max_sequence_length,
                    device=device,
                )
            else:
                negative_prompt_embeds = torch.zeros_like(prompt_embeds)

        if self.text_encoder is not None:
            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
                # Retrieve the original scale by scaling back the LoRA layers
                unscale_lora_layers(self.text_encoder, lora_scale)

        text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(device=device)
        text_ids = text_ids.repeat(num_images_per_prompt, 1, 1)

        return prompt_embeds, negative_prompt_embeds, text_ids

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

    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
    # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`
    # corresponds to doing no classifier free guidance.
    @property
    def do_classifier_free_guidance(self):
        return self._guidance_scale > 1

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

    @attention_kwargs.setter
    def attention_kwargs(self, value):
        self._attention_kwargs = value

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

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

    def check_inputs(
        self,
        prompt,
        height,
        width,
        negative_prompt=None,
        prompt_embeds=None,
        negative_prompt_embeds=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 {[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(
                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
                " only forward one of the two."
            )
        elif prompt is None and prompt_embeds is None:
            raise ValueError(
                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
            )
        elif 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(
                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
            )

        if prompt_embeds is not None and negative_prompt_embeds is not None:
            if prompt_embeds.shape != negative_prompt_embeds.shape:
                raise ValueError(
                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
                    f" {negative_prompt_embeds.shape}."
                )

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

    def _get_t5_prompt_embeds(
        self,
        prompt: str | list[str] = None,
        num_images_per_prompt: int = 1,
        max_sequence_length: int = 128,
        device: torch.device | None = None,
    ):
        tokenizer = self.tokenizer
        text_encoder = self.text_encoder
        device = device or text_encoder.device

        prompt = [prompt] if isinstance(prompt, str) else prompt
        batch_size = len(prompt)
        prompt_embeds_list = []
        for p in prompt:
            text_inputs = tokenizer(
                p,
                # padding="max_length",
                max_length=max_sequence_length,
                truncation=True,
                add_special_tokens=True,
                return_tensors="pt",
            )
            text_input_ids = text_inputs.input_ids
            untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids

            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                text_input_ids, untruncated_ids
            ):
                removed_text = tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
                logger.warning(
                    "The following part of your input was truncated because `max_sequence_length` is set to "
                    f" {max_sequence_length} tokens: {removed_text}"
                )

            prompt_embeds = text_encoder(text_input_ids.to(device))[0]

            # Concat zeros to max_sequence
            b, seq_len, dim = prompt_embeds.shape
            if seq_len < max_sequence_length:
                padding = torch.zeros(
                    (b, max_sequence_length - seq_len, dim), dtype=prompt_embeds.dtype, device=prompt_embeds.device
                )
                prompt_embeds = torch.concat([prompt_embeds, padding], dim=1)
            prompt_embeds_list.append(prompt_embeds)

        prompt_embeds = torch.concat(prompt_embeds_list, dim=0)
        prompt_embeds = prompt_embeds.to(device=device)

        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, max_sequence_length, -1)
        prompt_embeds = prompt_embeds.to(dtype=self.transformer.dtype)
        return prompt_embeds

    def prepare_latents(
        self,
        batch_size,
        num_channels_latents,
        height,
        width,
        dtype,
        device,
        generator,
        latents=None,
    ):
        # VAE applies 8x compression on images but we must also account for packing which requires
        # latent height and width to be divisible by 2.
        height = 2 * (int(height) // self.vae_scale_factor)
        width = 2 * (int(width) // self.vae_scale_factor)

        shape = (batch_size, num_channels_latents, height, width)

        if latents is not None:
            latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
            return latents.to(device=device, dtype=dtype), latent_image_ids

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

        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)

        latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)

        return latents, latent_image_ids

    @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, num_patches, channels = latents.shape

        height = height // vae_scale_factor
        width = width // vae_scale_factor

        latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
        latents = latents.permute(0, 3, 1, 4, 2, 5)

        latents = latents.reshape(batch_size, channels // (2 * 2), height * 2, width * 2)

        return latents

    @staticmethod
    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
        latent_image_ids = torch.zeros(height, width, 3)
        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]

        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape

        latent_image_ids = latent_image_ids.repeat(batch_size, 1, 1, 1)
        latent_image_ids = latent_image_ids.reshape(
            batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
        )

        return latent_image_ids.to(device=device, dtype=dtype)

    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
        self,
        prompt: str | list[str] = None,
        height: int | None = None,
        width: int | None = None,
        num_inference_steps: int = 30,
        timesteps: list[int] = None,
        guidance_scale: float = 5,
        negative_prompt: str | list[str] | None = None,
        num_images_per_prompt: int | None = 1,
        generator: torch.Generator | list[torch.Generator] | None = None,
        latents: torch.FloatTensor | None = None,
        prompt_embeds: torch.FloatTensor | None = None,
        negative_prompt_embeds: torch.FloatTensor | None = None,
        output_type: str | None = "pil",
        return_dict: bool = True,
        attention_kwargs: dict[str, Any] | None = None,
        callback_on_step_end: Callable[[int, int], None] | None = None,
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
        max_sequence_length: int = 128,
        clip_value: None | float = None,
        normalize: bool = False,
    ):
        r"""
                Function invoked when calling the 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`. instead.
                    height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                        The height in pixels of the generated image. This is set to 1024 by default for the best
                        results.
                    width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                        The width in pixels of the generated image. This is set to 1024 by default for the best
                        results.
                    num_inference_steps (`int`, *optional*, defaults to 50):
                        The number of denoising steps. More denoising steps usually lead to a higher quality image at
                        the expense of slower inference.
                    timesteps (`list[int]`, *optional*):
                        Custom timesteps to use for the denoising process with schedulers which support a `timesteps`
                        argument in their `set_timesteps` method. If not defined, the default behavior when
                        `num_inference_steps` is passed will be used. Must be in descending order.
                    guidance_scale (`float`, *optional*, defaults to 5.0):
        <<<<<<< HEAD
                        Guidance scale as defined in [Classifier-Free Diffusion
                        Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                        of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting
                        `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely
                        linked to the text `prompt`, usually at the expense of lower image quality.
                    negative_prompt (`str` or `list[str]`, *optional*):
        =======
                        Guidance scale as defined in [Classifier-Free Diffusion
                        Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of
                        equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is
                        enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images
                        that are closely linked to the text `prompt`, usually at the expense of lower image quality.
                    negative_prompt (`str` or `list[str]`, *optional*):
        >>>>>>> main
                        The prompt or prompts not to guide the image generation. If not defined, one has to pass
                        `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if
                        `guidance_scale` is less than `1`).
                    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*):
                        One or a list of [torch
                        generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
                        generation deterministic.
                    latents (`torch.FloatTensor`, *optional*):
                        Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for
                        image generation. Can be used to tweak the same generation with different prompts. If not
                        provided, a latents tensor will be generated by sampling using the supplied random `generator`.
                    prompt_embeds (`torch.FloatTensor`, *optional*):
                        Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                        weighting. If not provided, text embeddings will be generated from `prompt` input argument.
                    negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                        Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                        weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt`
                        input argument.
                    output_type (`str`, *optional*, defaults to `"pil"`):
                        The output format of the generate image. Choose between
                        [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
                    return_dict (`bool`, *optional*, defaults to `True`):
                        Whether or not to return a [`~pipelines.bria.BriaPipelineOutput`] instead of a plain tuple.
                    attention_kwargs (`dict`, *optional*):
                        A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined
                        under `self.processor` in
                        [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
                    callback_on_step_end (`Callable`, *optional*):
                        A function that calls at the end of each denoising steps during the inference. The function is
                        called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int,
                        timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as
                        specified by `callback_on_step_end_tensor_inputs`.
                    callback_on_step_end_tensor_inputs (`list`, *optional*):
                        The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the
                        list will be passed as `callback_kwargs` argument. You will only be able to include variables
                        listed in the `._callback_tensor_inputs` attribute of your pipeline class.
                    max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`.

                Examples:

                  Returns:
                    [`~pipelines.bria.BriaPipelineOutput`] or `tuple`: [`~pipelines.bria.BriaPipelineOutput`] if
                    `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list
                    with the generated images.
        """

        height = height or self.default_sample_size * self.vae_scale_factor
        width = width or self.default_sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt=prompt,
            height=height,
            width=width,
            prompt_embeds=prompt_embeds,
            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._interrupt = False

        # 2. Define call parameters
        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

        lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None

        (prompt_embeds, negative_prompt_embeds, text_ids) = self.encode_prompt(
            prompt=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,
            num_images_per_prompt=num_images_per_prompt,
            max_sequence_length=max_sequence_length,
            lora_scale=lora_scale,
        )

        if self.do_classifier_free_guidance:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)

        # 5. Prepare latent variables
        num_channels_latents = self.transformer.config.in_channels // 4  # due to patch=2, we devide by 4
        latents, latent_image_ids = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        if (
            isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler)
            and self.scheduler.config["use_dynamic_shifting"]
        ):
            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
            image_seq_len = latents.shape[1]

            mu = calculate_shift(
                image_seq_len,
                self.scheduler.config.base_image_seq_len,
                self.scheduler.config.max_image_seq_len,
                self.scheduler.config.base_shift,
                self.scheduler.config.max_shift,
            )
            timesteps, num_inference_steps = retrieve_timesteps(
                self.scheduler,
                num_inference_steps,
                device,
                timesteps,
                sigmas,
                mu=mu,
            )
        else:
            # 4. Prepare timesteps
            # Sample from training sigmas
            if isinstance(self.scheduler, DDIMScheduler) or isinstance(
                self.scheduler, EulerAncestralDiscreteScheduler
            ):
                timesteps, num_inference_steps = retrieve_timesteps(
                    self.scheduler, num_inference_steps, device, None, None
                )
            else:
                sigmas = get_original_sigmas(
                    num_train_timesteps=self.scheduler.config.num_train_timesteps,
                    num_inference_steps=num_inference_steps,
                )
                timesteps, num_inference_steps = retrieve_timesteps(
                    self.scheduler, num_inference_steps, device, timesteps, sigmas=sigmas
                )

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

        if len(latent_image_ids.shape) == 3:
            latent_image_ids = latent_image_ids[0]
        if len(text_ids.shape) == 3:
            text_ids = text_ids[0]

        # 6. Denoising loop
        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

                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                if type(self.scheduler) != FlowMatchEulerDiscreteScheduler:
                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                timestep = t.expand(latent_model_input.shape[0])

                # This is predicts "v" from flow-matching or eps from diffusion
                noise_pred = self.transformer(
                    hidden_states=latent_model_input,
                    timestep=timestep,
                    encoder_hidden_states=prompt_embeds,
                    attention_kwargs=self.attention_kwargs,
                    return_dict=False,
                    txt_ids=text_ids,
                    img_ids=latent_image_ids,
                )[0]

                # perform guidance
                if self.do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    cfg_noise_pred_text = noise_pred_text.std()
                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)

                if normalize:
                    noise_pred = noise_pred * (0.7 * (cfg_noise_pred_text / noise_pred.std())) + 0.3 * noise_pred

                if clip_value:
                    assert clip_value > 0
                    noise_pred = noise_pred.clip(-clip_value, clip_value)

                # compute the previous noisy sample x_t -> x_t-1
                latents_dtype = latents.dtype
                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]

                if latents.dtype != latents_dtype:
                    if torch.backends.mps.is_available():
                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
                        latents = latents.to(latents_dtype)

                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)

                    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)

                # call the callback, if provided
                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()

        if output_type == "latent":
            image = latents

        else:
            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
            latents = (latents.to(dtype=torch.float32) / self.vae.config.scaling_factor) + self.vae.config.shift_factor
            image = self.vae.decode(latents.to(dtype=self.vae.dtype), return_dict=False)[0]
            image = self.image_processor.postprocess(image, output_type=output_type)

        # Offload all models
        self.maybe_free_model_hooks()

        if not return_dict:
            return (image,)

        return BriaPipelineOutput(images=image)