pipeline.py

import inspect
from typing import Any, Callable, Dict, List, Optional, Union
import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionModelWithProjection,
    T5EncoderModel,
    T5TokenizerFast,
)
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
from diffusers.models import AutoencoderKL, FluxTransformer2DModel
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import (
    USE_PEFT_BACKEND,
    deprecate,
    is_torch_xla_available,
    logging,
    replace_example_docstring,
    scale_lora_layers,
    unscale_lora_layers,
)
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
from diffusers import FluxKontextPipeline

PREFERRED_KONTEXT_RESOLUTIONS = [
    (672, 1568),
    (688, 1504),
    (720, 1456),
    (752, 1392),
    (800, 1328),
    (832, 1248),
    (880, 1184),
    (944, 1104),
    (1024, 1024),
    (1104, 944),
    (1184, 880),
    (1248, 832),
    (1328, 800),
    (1392, 752),
    (1456, 720),
    (1504, 688),
    (1568, 672),
]


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,
):
    r"""
    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


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


from diffusers import FluxKontextPipeline
from typing import Union, List, Optional
import torch


class CustomFluxKontextPipeline(FluxKontextPipeline):
	r"""
	Custom Flux Kontext pipeline with a wrapper text embedder.
	"""

	model_cpu_offload_seq = "text_embedder->image_encoder->transformer->vae"

	def __init__(
			self,
			scheduler,
			vae,
			text_embedder,  # Your custom text embedder wrapper
			transformer,
			aligner,
			image_encoder=None,
			feature_extractor=None,
	):
		# Don't call super().__init__() since parent expects text_encoder parameters
		# Instead, manually register modules
		from diffusers import DiffusionPipeline
		DiffusionPipeline.__init__(self)

		self.register_modules(
			vae=vae,
			text_embedder=text_embedder,
			transformer=transformer,
			scheduler=scheduler,
			aligner=aligner,
			image_encoder=image_encoder,
			feature_extractor=feature_extractor,
		)

		# Initialize the necessary attributes from parent
		self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
		self.latent_channels = self.vae.config.latent_channels
		from diffusers.image_processor import VaeImageProcessor
		self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
		self.default_sample_size = 128

	def encode_prompt(
			self,
			prompt: Union[str, List[str]],
			prompt_2: Optional[Union[str, List[str]]] = None,
			device: Optional[torch.device] = None,
			num_images_per_prompt: int = 1,
			prompt_embeds: Optional[torch.FloatTensor] = None,
			pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
			max_sequence_length: int = 512,
			lora_scale: Optional[float] = None,
			temperature=None,
			threshold=None,
	):
		device = device or self._execution_device

		if prompt_embeds is None:
			# Use your custom text embedder
			qwen_embeds, clip_image_embeds, perturbed_qwen_embeds, replace_ids, t5_tokenizer, batch_encoding = self.text_embedder(prompt)
			prompt_embeds, prompt_attention_mask, pooled_prompt_embeds, text_ids, _ = self.aligner(qwen_embeds,
			                                                                                       )
			prompt_embeds = prompt_embeds.to(device=device)
			pooled_prompt_embeds = pooled_prompt_embeds.to(device=device)
			text_ids = text_ids.to(device=device)
		else:
			# When embeddings are provided, create text_ids
			dtype = self.transformer.dtype
			text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)

		# Duplicate for num_images_per_prompt
		if num_images_per_prompt > 1:
			prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
			pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
			text_ids = text_ids.repeat(num_images_per_prompt, 1)

		return prompt_embeds, prompt_attention_mask, pooled_prompt_embeds, text_ids

	@torch.no_grad()
	def __call__(
			self,
			image: Optional[PipelineImageInput] = None,
			prompt: Union[str, List[str]] = None,
			prompt_2: Optional[Union[str, List[str]]] = None,
			negative_prompt: Union[str, List[str]] = "",
			negative_prompt_2: Optional[Union[str, List[str]]] = None,
			true_cfg_scale: float = 1.0,
			height: Optional[int] = None,
			width: Optional[int] = None,
			num_inference_steps: int = 28,
			sigmas: Optional[List[float]] = None,
			guidance_scale: float = 3.5,
			num_images_per_prompt: Optional[int] = 1,
			generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
			latents: Optional[torch.FloatTensor] = None,
			prompt_embeds: Optional[torch.FloatTensor] = None,
			pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
			ip_adapter_image: Optional[PipelineImageInput] = None,
			ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
			negative_ip_adapter_image: Optional[PipelineImageInput] = None,
			negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
			negative_prompt_embeds: Optional[torch.FloatTensor] = None,
			negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
			output_type: Optional[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,
			max_area: int = 1024 ** 2,
			_auto_resize: bool = True,
			temperature=None,
			threshold=None,
	):
		r"""
		Function invoked when calling the pipeline for generation.

		Args:
			image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
				`Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
				numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
				or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
				list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
				latents as `image`, but if passing latents directly it is not encoded again.
			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.
			prompt_2 (`str` or `List[str]`, *optional*):
				The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
				will be used instead.
			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 `true_cfg_scale` is
				not greater than `1`).
			negative_prompt_2 (`str` or `List[str]`, *optional*):
				The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
				`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
			true_cfg_scale (`float`, *optional*, defaults to 1.0):
				When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
			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.
			sigmas (`List[float]`, *optional*):
				Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
				their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
				will be used.
			guidance_scale (`float`, *optional*, defaults to 3.5):
				Embedded guidance scale is enabled by setting `guidance_scale` > 1. Higher `guidance_scale` encourages
				a model to generate images more aligned with prompt at the expense of lower image quality.

				Guidance-distilled models approximates true classifier-free guidance for `guidance_scale` > 1. Refer to
				the [paper](https://huggingface.co/papers/2210.03142) to learn more.
			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.
			pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
				Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
				If not provided, pooled text embeddings will be generated from `prompt` input argument.
			ip_adapter_image: (`PipelineImageInput`, *optional*):
				Optional image input to work with IP Adapters.
			ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
				Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
				IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
				provided, embeddings are computed from the `ip_adapter_image` input argument.
			negative_ip_adapter_image:
				(`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
			negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
				Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
				IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
				provided, embeddings are computed from the `ip_adapter_image` 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.
			negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
				Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
				weighting. If not provided, pooled 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.flux.FluxPipelineOutput`] instead of a plain tuple.
			joint_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 512):
				Maximum sequence length to use with the `prompt`.
			max_area (`int`, defaults to `1024 ** 2`):
				The maximum area of the generated image in pixels. The height and width will be adjusted to fit this
				area while maintaining the aspect ratio.

		Examples:

		Returns:
			[`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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

		original_height, original_width = height, width
		aspect_ratio = width / height

		"""
		width = round((max_area * aspect_ratio) ** 0.5)
		height = round((max_area / aspect_ratio) ** 0.5)
		multiple_of = self.vae_scale_factor * 2
		width = width // multiple_of * multiple_of
		height = height // multiple_of * multiple_of

		if height != original_height or width != original_width:
			print(
				f"Generation `height` and `width` have been adjusted to {height} and {width} to fit the model requirements."
			)
		"""

		# 1. Check inputs. Raise error if not correct
		self.check_inputs(
			prompt,
			prompt_2,
			height,
			width,
			negative_prompt=negative_prompt,
			negative_prompt_2=negative_prompt_2,
			prompt_embeds=prompt_embeds,
			negative_prompt_embeds=negative_prompt_embeds,
			pooled_prompt_embeds=pooled_prompt_embeds,
			negative_pooled_prompt_embeds=negative_pooled_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._joint_attention_kwargs = joint_attention_kwargs
		self._current_timestep = None
		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 = torch.device("cuda" if torch.cuda.is_available() else "cpu")

		lora_scale = (
			self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
		)
		has_neg_prompt = negative_prompt is not None or (
				negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
		)
		do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
		(
			prompt_embeds,
			prompt_attention_mask,
			pooled_prompt_embeds,
			text_ids,
		) = self.encode_prompt(
			prompt=prompt,
			prompt_2=prompt_2,
			prompt_embeds=prompt_embeds,
			pooled_prompt_embeds=pooled_prompt_embeds,
			device=device,
			num_images_per_prompt=num_images_per_prompt,
			max_sequence_length=max_sequence_length,
			lora_scale=lora_scale,
			temperature=temperature,
			threshold=threshold,
		)
		(
				negative_prompt_embeds,
				negative_prompt_attention_mask,
				negative_pooled_prompt_embeds,
				negative_text_ids,
			) = self.encode_prompt(
				prompt=negative_prompt,
				prompt_2=negative_prompt_2,
				prompt_embeds=negative_prompt_embeds,
				pooled_prompt_embeds=negative_pooled_prompt_embeds,
				device=device,
				num_images_per_prompt=num_images_per_prompt,
				max_sequence_length=max_sequence_length,
				lora_scale=lora_scale,
				temperature=temperature,
				threshold=threshold,
			)

		pooled_prompt_embeds = negative_pooled_prompt_embeds

		# 3. Preprocess image
		if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
			img = image[0] if isinstance(image, list) else image
			"""
			image_height, image_width = self.image_processor.get_default_height_width(img)
			aspect_ratio = image_width / image_height
			if _auto_resize:
				# Kontext is trained on specific resolutions, using one of them is recommended
				_, image_width, image_height = min(
					(abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS
				)
			image_width = image_width // multiple_of * multiple_of
			image_height = image_height // multiple_of * multiple_of
			"""
			image_height, image_width = original_height, original_width
			image = self.image_processor.resize(image, image_height, image_width)
			image = self.image_processor.preprocess(image, image_height, image_width)

		# 4. Prepare latent variables
		num_channels_latents = self.transformer.config.in_channels // 4
		latents, image_latents, latent_ids, image_ids = self.prepare_latents(
			image,
			batch_size * num_images_per_prompt,
			num_channels_latents,
			height,
			width,
			prompt_embeds.dtype,
			device,
			generator,
			latents,
		)
		if image_ids is not None:
			latent_ids = torch.cat([latent_ids, image_ids], dim=0)  # dim 0 is sequence dimension

		# 5. Prepare 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)

		# handle guidance
		if self.transformer.config.guidance_embeds:
			guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
			guidance = guidance.expand(latents.shape[0])
		else:
			guidance = None

		if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
				negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
		):
			negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
			negative_ip_adapter_image = [negative_ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters

		elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
				negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
		):
			ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
			ip_adapter_image = [ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters

		if self.joint_attention_kwargs is None:
			self._joint_attention_kwargs = {}

		image_embeds = None
		negative_image_embeds = None
		if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
			image_embeds = self.prepare_ip_adapter_image_embeds(
				ip_adapter_image,
				ip_adapter_image_embeds,
				device,
				batch_size * num_images_per_prompt,
			)
		if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
			negative_image_embeds = self.prepare_ip_adapter_image_embeds(
				negative_ip_adapter_image,
				negative_ip_adapter_image_embeds,
				device,
				batch_size * num_images_per_prompt,
			)

		# 6. Denoising loop
		# We set the index here to remove DtoH sync, helpful especially during compilation.
		# Check out more details here: https://github.com/huggingface/diffusers/pull/11696
		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
				if image_embeds is not None:
					self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds

				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)

				noise_pred = self.transformer(
					hidden_states=latent_model_input,
					timestep=timestep / 1000,
					guidance=guidance,
					pooled_projections=pooled_prompt_embeds,
					encoder_hidden_states=prompt_embeds,
					txt_ids=text_ids,
					img_ids=latent_ids,
					joint_attention_kwargs={'attention_mask': prompt_attention_mask},
					return_dict=False,
				)[0]
				noise_pred = noise_pred[:, : latents.size(1)]

				if do_true_cfg:
					if negative_image_embeds is not None:
						self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
					neg_noise_pred = self.transformer(
						hidden_states=latent_model_input,
						timestep=timestep / 1000,
						guidance=guidance,
						pooled_projections=negative_pooled_prompt_embeds,
						encoder_hidden_states=negative_prompt_embeds,
						txt_ids=negative_text_ids,
						img_ids=latent_ids,
						joint_attention_kwargs={'attention_mask': negative_prompt_attention_mask},
						return_dict=False,
					)[0]
					neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
					noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)

				# 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)

				# 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()

		self._current_timestep = None

		if output_type == "latent":
			image = latents
		else:
			latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
			latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor

			dtype = torch.bfloat16
			image = self.vae.decode(latents.to(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 FluxPipelineOutput(images=image)