#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
#
from __future__ import annotations

import argparse
import inspect
import os
import time
import warnings
from typing import Any, List, Optional

import numpy as np
import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
from diffusers.image_processor import VaeImageProcessor
from flux.content_filters import PixtralContentFilter
from huggingface_hub import snapshot_download

from demo_diffusion import path as path_module
from demo_diffusion.model import (
    CLIPModel,
    FLUXLoraLoader,
    FluxTransformerModel,
    T5Model,
    VAEEncoderModel,
    VAEModel,
    get_clip_embedding_dim,
    load,
    make_tokenizer,
)
from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline
from demo_diffusion.pipeline.type import PIPELINE_TYPE

TRT_LOGGER = trt.Logger(trt.Logger.ERROR)

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


class FluxKontextUtil:
    """
    Utility class for Flux Kontext-related operations including context dimension calculations
    and resolution handling.
    """

    @staticmethod
    def _get_context_dim(
        image_height: int,
        image_width: int,
        compression_factor: int,
    ) -> int:
        """
        Calculate the context dimension based on image dimensions and compression factor.

        Args:
            image_height: Height of the image
            image_width: Width of the image
            compression_factor: Compression factor applied to the image

        Returns:
            The calculated sequence length for the context
        """
        seq_len = (image_height // (2 * compression_factor)) * (image_width // (2 * compression_factor))
        return seq_len

    @staticmethod
    def get_context_latent_dim(
        version: str,
        kontext_resolution: tuple = None,
        compression_factor: int = 8,
        static_shape: bool = False,
    ):
        """
        Get the context latent dimensions for Flux Kontext models.

        Args:
            version: Model version string
            kontext_resolution: Tuple of (width, height) for kontext resolution
            compression_factor: Compression factor for the model
            static_shape: Whether to use static shapes

        Returns:
            Tuple of (min_context_latent_dim, context_latent_dim, max_context_latent_dim)
        """
        min_context_latent_dim, context_latent_dim, max_context_latent_dim = 0, 0, 0

        if version == "flux.1-kontext-dev":
            assert kontext_resolution is not None, "kontext_resolution must be provided for flux.1-kontext-dev"

            # get opt context size
            context_latent_dim = FluxKontextUtil._get_context_dim(
                image_height=kontext_resolution[1],
                image_width=kontext_resolution[0],
                compression_factor=compression_factor,
            )

            # get min context size
            _, min_context_width, min_context_height = min((w * h, w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS)
            min_context_latent_dim = (
                context_latent_dim
                if static_shape
                else FluxKontextUtil._get_context_dim(
                    image_height=min_context_height,
                    image_width=min_context_width,
                    compression_factor=compression_factor,
                )
            )

            # get max context size
            _, max_context_width, max_context_height = max((w * h, w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS)
            max_context_latent_dim = (
                context_latent_dim
                if static_shape
                else FluxKontextUtil._get_context_dim(
                    image_height=max_context_height,
                    image_width=max_context_width,
                    compression_factor=compression_factor,
                )
            )

        return min_context_latent_dim, context_latent_dim, max_context_latent_dim

    @staticmethod
    def get_preferred_resolutions():
        """
        Get the list of preferred Kontext resolutions.

        Returns:
            List of (width, height) tuples representing preferred resolutions
        """
        return PREFERRED_KONTEXT_RESOLUTIONS.copy()

    @staticmethod
    def get_min_max_kontext_dimensions():
        """
        Get the resolution tuples with minimum and maximum aspect ratios from preferred Kontext resolutions.

        Returns:
            Tuple of ((min_aspect_width, min_aspect_height), (max_aspect_width, max_aspect_height))
        """
        widths = [w for w, h in PREFERRED_KONTEXT_RESOLUTIONS]
        heights = [h for w, h in PREFERRED_KONTEXT_RESOLUTIONS]

        min_width = min(widths)
        max_width = max(widths)
        min_height = min(heights)
        max_height = max(heights)

        return (min_width, min_height), (max_width, max_height)

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.16,
):
    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


class FluxPipeline(DiffusionPipeline):
    """
    Application showcasing the acceleration of Flux pipelines using Nvidia TensorRT.
    """

    def __init__(
        self,
        version="flux.1-dev",
        pipeline_type=PIPELINE_TYPE.TXT2IMG,
        guidance_scale=3.5,
        max_sequence_length=512,
        calibration_dataset=None,
        t5_weight_streaming_budget_percentage=None,
        transformer_weight_streaming_budget_percentage=None,
        lora_scale: float = 1.0,
        lora_weight: Optional[List[float]] = None,
        lora_path: Optional[List[str]] = None,
        **kwargs,
    ):
        """
        Initializes the Flux pipeline.

        Args:
            version (`str`, defaults to `flux.1-dev`)
                Version of the underlying Flux model.
            guidance_scale (`float`, defaults to 3.5):
                Guidance scale is enabled by setting as > 1.
                Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality.
            max_sequence_length (`int`, defaults to 512):
                Maximum sequence length to use with the `prompt`.
            t5_weight_streaming_budget_percentage (`int`, defaults to None):
                Weight streaming budget as a percentage of the size of total streamable weights for the T5 model.
            transformer_weight_streaming_budget_percentage (`int`, defaults to None):
                Weight streaming budget as a percentage of the size of total streamable weights for the transformer model.
        """
        super().__init__(
            version=version,
            pipeline_type=pipeline_type,
            text_encoder_weight_streaming_budget_percentage=t5_weight_streaming_budget_percentage,
            denoiser_weight_streaming_budget_percentage=transformer_weight_streaming_budget_percentage,
            **kwargs,
        )
        self.guidance_scale = guidance_scale
        self.max_sequence_length = max_sequence_length
        self.calibration_dataset = calibration_dataset  # Currently supported for Flux ControlNet pipelines only

        # Initialize LoRA
        self.lora_loader = None
        if lora_path:
            self.lora_weights = dict()
            self.lora_loader = FLUXLoraLoader(lora_path, lora_weight, lora_scale)
            assert len(lora_path) == len(lora_weight)
            for i, path in enumerate(lora_path):
                self.lora_weights[path] = lora_weight[i]

    @classmethod
    def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> FluxPipeline:
        """Factory method to construct a `FluxPipeline` object from parsed arguments.

        Overrides:
            DiffusionPipeline.FromArgs
        """
        MAX_BATCH_SIZE = 4
        DEVICE = "cuda"
        DO_RETURN_LATENTS = False

        # Resolve all paths.
        dd_path = path_module.resolve_path(
            cls.get_model_names(pipeline_type), args, pipeline_type, cls._get_pipeline_uid(args.version)
        )

        return cls(
            dd_path=dd_path,
            version=args.version,
            pipeline_type=pipeline_type,
            guidance_scale=args.guidance_scale,
            max_sequence_length=args.max_sequence_length,
            bf16=args.bf16,
            calibration_dataset=args.calibration_dataset if hasattr(args, "calibration_dataset") else None,
            low_vram=args.low_vram,
            torch_fallback=args.torch_fallback,
            weight_streaming=args.ws,
            t5_weight_streaming_budget_percentage=args.t5_ws_percentage,
            transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage,
            max_batch_size=MAX_BATCH_SIZE,
            denoising_steps=args.denoising_steps,
            scheduler=args.scheduler,
            lora_scale=args.lora_scale,
            lora_weight=args.lora_weight,
            lora_path=args.lora_path,
            device=DEVICE,
            output_dir=args.output_dir,
            hf_token=args.hf_token,
            verbose=args.verbose,
            nvtx_profile=args.nvtx_profile,
            use_cuda_graph=args.use_cuda_graph,
            framework_model_dir=args.framework_model_dir,
            return_latents=DO_RETURN_LATENTS,
            torch_inference=args.torch_inference,
        )

    @classmethod
    def get_model_names(cls, pipeline_type: PIPELINE_TYPE, controlnet_type: str = None) -> List[str]:
        """Return a list of model names used by this pipeline.

        Overrides:
            DiffusionPipeline.get_model_names
        """
        if pipeline_type.is_img2img():
            return ["clip", "t5", "transformer", "vae", "vae_encoder"]
        else:
            return ["clip", "t5", "transformer", "vae"]

    def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
        if self.fp16:
            raise ValueError(
                "ONNX models can be downloaded only for the following precisions: BF16, FP8, FP4. This pipeline is running in FP16."
            )

        hf_download_path = "-".join([load.get_path(self.version, self.pipeline_type.name), "onnx"])
        model_path = model_config["onnx_opt_path"]
        base_dir = os.path.dirname(os.path.dirname(model_config["onnx_opt_path"]))

        if not os.path.exists(model_path):
            if model_name == "clip":
                dirname = "clip.opt"
            elif model_name == "t5":
                dirname = "t5.opt"
            elif model_name == "transformer":
                if model_config["use_fp4"]:
                    dirname = "transformer.opt/fp4"
                    if self.version == "flux.1-kontext-dev":
                        dirname = "_".join([dirname, "svd32"])
                elif model_config["use_fp8"]:
                    dirname = "transformer.opt/fp8"
                elif self.bf16:
                    dirname = "transformer.opt/bf16"
            elif model_name == "vae":
                dirname = "vae.opt"
            elif model_name == "vae_encoder":
                dirname = "vae_encoder.opt"
            else:
                raise ValueError(f"{model_name} not found in {self.stages}")

            snapshot_download(
                repo_id=hf_download_path,
                allow_patterns=os.path.join(dirname, "*"),
                local_dir=base_dir,
                token=self.hf_token,
            )
            # Rename directory from <model_name>.opt to <model_name>
            saved_dir = os.path.join(base_dir, dirname)
            model_dir = os.path.dirname(model_path)
            os.rename(saved_dir, model_dir)
            # Rename model from model.onnx to model_optimized.onnx
            os.rename(os.path.join(model_dir, "model.onnx"), model_path)

    def is_native_export_supported(self, model_config: dict[str, Any]) -> bool:
        if self.version.startswith("flux.1") and model_config["use_fp4"]:
            # Native export not supported for FP4.
            raise ValueError(
                f"Native FP4 quantization is not supported. No ONNX model found in {model_config['onnx_opt_path']}. Please pass --download-onnx-models."
            )
        if (
            self.version in ["flux.1-dev-canny", "flux.1-dev-depth"]
            and model_config["use_fp8"]
            and not self.calibration_dataset
        ):
            # Native export of FP8 model requires calibration data.
            raise ValueError(
                f"No ONNX model found in {model_config['onnx_opt_path']}. Please pass --download-onnx-models. If you would like to quantize and export natively, please provide calibration data using --calibration-."
            )
        return True

    def _initialize_models(self, framework_model_dir, int8, fp8, fp4):
        # Load text tokenizer(s)
        self.tokenizer = make_tokenizer(
            self.version, self.pipeline_type, self.hf_token, framework_model_dir,
        )
        self.tokenizer2 = make_tokenizer(
            self.version,
            self.pipeline_type,
            self.hf_token,
            framework_model_dir,
            subfolder="tokenizer_2",
            tokenizer_type="t5",
        )

        # Load pipeline models
        models_args = {
            "version": self.version,
            "pipeline": self.pipeline_type,
            "device": self.device,
            "hf_token": self.hf_token,
            "verbose": self.verbose,
            "framework_model_dir": framework_model_dir,
            "max_batch_size": self.max_batch_size,
        }

        self.bf16 = True if int8 or fp8 or fp4 else self.bf16
        self.fp16 = True if not self.bf16 else False
        self.tf32 = True
        if "clip" in self.stages:
            # BF16 CLIP ONNX export fails with ComplexDouble error in newer PyTorch; use FP16.
            self.models["clip"] = CLIPModel(
                **models_args,
                fp16=True,
                tf32=self.tf32,
                bf16=False,
                embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type),
                keep_pooled_output=True,
                subfolder="text_encoder",
            )

        if "t5" in self.stages:
            # Known accuracy issues with FP16
            self.models["t5"] = T5Model(
                **models_args,
                fp16=self.fp16,
                tf32=self.tf32,
                bf16=self.bf16,
                subfolder="text_encoder_2",
                text_maxlen=self.max_sequence_length,
                weight_streaming=self.weight_streaming,
                weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
            )

        if "vae" in self.stages:
            # Accuracy issues with FP16
            self.models["vae"] = VAEModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16)

        self.vae_scale_factor = (
            2 ** (len(self.models["vae"].config["block_out_channels"]))
            if "vae" in self.stages and self.models["vae"] is not None
            else 16
        )
        self.vae_latent_channels = (
            self.models["vae"].config["latent_channels"]
            if "vae" in self.stages and self.models["vae"] is not None
            else 16
        )

        if "vae_encoder" in self.stages:
            self.image_processor = VaeImageProcessor(
                vae_scale_factor=self.vae_scale_factor * 2, vae_latent_channels=self.vae_latent_channels
            )
            # Add kontext_resolution if available (for FluxKontextPipeline)
            vae_encoder_kwargs = {}
            if hasattr(self, "kontext_image") and self.kontext_image is not None:
                self.resize_height, self.resize_width = self._get_resize_dimensions(self.kontext_image)
                vae_encoder_kwargs["kontext_resolution"] = (self.resize_width, self.resize_height)

            self.models["vae_encoder"] = VAEEncoderModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16, **vae_encoder_kwargs)

        if "transformer" in self.stages:
            transformer_kwargs = {
                **models_args,
                "bf16": self.bf16,
                "fp16": self.fp16,
                "int8": int8,
                "fp8": fp8,
                "tf32": self.tf32,
                "text_maxlen": self.max_sequence_length,
                "weight_streaming": self.weight_streaming,
                "weight_streaming_budget_percentage": self.denoiser_weight_streaming_budget_percentage,
            }
            if hasattr(self, "kontext_image") and self.kontext_image is not None:
                transformer_kwargs["kontext_resolution"] = (self.resize_width, self.resize_height)
            self.models["transformer"] = FluxTransformerModel(**transformer_kwargs)

    def encode_image(self, input_image, encoder="vae_encoder"):
        self.profile_start(encoder, color='red')
        cast_to = torch.float16 if self.models[encoder].fp16 else torch.bfloat16 if self.models[encoder].bf16 else torch.float32
        input_image = input_image.to(dtype=cast_to)
        if self.torch_inference or self.torch_fallback[encoder]:
            image_latents = self.torch_models[encoder](input_image)
        else:
            image_latents = self.run_engine(encoder, {'images': input_image})['latent']

        image_latents = self.models[encoder].config["scaling_factor"] * (image_latents - self.models[encoder].config["shift_factor"])
        self.profile_stop(encoder)
        return image_latents

    # Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py#L546
    def prepare_image(
        self,
        image,
        width,
        height,
        batch_size,
        num_images_per_prompt,
        device,
        dtype,
        do_classifier_free_guidance=False,
        guess_mode=False,
    ):
        if isinstance(image, torch.Tensor):
            pass
        else:
            image = self.image_processor.preprocess(image, height=height, width=width)

        image_batch_size = image.shape[0]

        if image_batch_size == 1:
            repeat_by = batch_size
        else:
            # image batch size is the same as prompt batch size
            repeat_by = num_images_per_prompt

        image = image.repeat_interleave(repeat_by, dim=0)

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

        if do_classifier_free_guidance and not guess_mode:
            image = torch.cat([image] * 2)

        return image

    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.1/src/diffusers/pipelines/flux/pipeline_flux.py#L436
    @staticmethod
    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
        """
        Reshapes latents from (B, C, H, W) to (B, H/2, W/2, C*4) as expected by the denoiser
        """
        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

    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.1/src/diffusers/pipelines/flux/pipeline_flux.py#L444
    @staticmethod
    def _unpack_latents(latents, height, width, vae_scale_factor):
        """
        Reshapes denoised latents to the format (B, C, H, W)
        """
        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

    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.1/src/diffusers/pipelines/flux/pipeline_flux.py#L421
    @staticmethod
    def _prepare_latent_image_ids(height, width, dtype, device):
        """
        Prepares latent image indices
        """
        latent_image_ids = torch.zeros(height // 2, width // 2, 3)
        latent_image_ids[..., 1] = (
            latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
        )
        latent_image_ids[..., 2] = (
            latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
        )

        latent_image_id_height, latent_image_id_width, latent_image_id_channels = (
            latent_image_ids.shape
        )

        latent_image_ids = latent_image_ids.reshape(
            latent_image_id_height * latent_image_id_width, latent_image_id_channels
        )

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

    def initialize_latents(
        self,
        batch_size,
        num_channels_latents,
        latent_height,
        latent_width,
        latent_timestep=None,
        image_latents=None,
        latents_dtype=torch.float32,
    ):
        latents_dtype = latents_dtype  # text_embeddings.dtype
        latents_shape = (batch_size, num_channels_latents, latent_height, latent_width)
        latents = torch.randn(
            latents_shape,
            device=self.device,
            dtype=latents_dtype,
            generator=self.generator,
        )

        if image_latents is not None:
            image_latents = torch.cat([image_latents], dim=0).to(latents_dtype)
            latents = self.scheduler.scale_noise(image_latents, latent_timestep, latents)

        latents = self._pack_latents(
            latents, batch_size, num_channels_latents, latent_height, latent_width
        )

        latent_image_ids = self._prepare_latent_image_ids(latent_height, latent_width, latents_dtype, self.device)

        return latents, latent_image_ids

    # Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_img2img.py#L416C1
    def get_timesteps(self, num_inference_steps, strength):
        # get the original timestep using init_timestep
        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 encode_prompt(
        self, prompt, encoder="clip", max_sequence_length=None, pooled_output=False
    ):
        self.profile_start(encoder, color="green")

        tokenizer = self.tokenizer2 if encoder == "t5" else self.tokenizer
        max_sequence_length = (
            tokenizer.model_max_length
            if max_sequence_length is None
            else max_sequence_length
        )

        def tokenize(prompt, max_sequence_length):
            text_input_ids = (
                tokenizer(
                    prompt,
                    padding="max_length",
                    max_length=max_sequence_length,
                    truncation=True,
                    return_overflowing_tokens=False,
                    return_length=False,
                    return_tensors="pt",
                )
                .input_ids.type(torch.int32)
                .to(self.device)
            )

            untruncated_ids = tokenizer(
                prompt, padding="longest", return_tensors="pt"
            ).input_ids.type(torch.int32).to(self.device)
            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]
                )
                warnings.warn(
                    "The following part of your input was truncated because `max_sequence_length` is set to "
                    f"{max_sequence_length} tokens: {removed_text}"
                )

            if self.torch_inference or self.torch_fallback[encoder]:
                outputs = self.torch_models[encoder](
                    text_input_ids, output_hidden_states=False
                )
                text_encoder_output = (
                    outputs[0].clone()
                    if pooled_output == False
                    else outputs.pooler_output.clone()
                )
            else:
                # NOTE: output tensor for the encoder must be cloned because it will be overwritten when called again for prompt2
                outputs = self.run_engine(encoder, {"input_ids": text_input_ids})
                output_name = (
                    "text_embeddings" if not pooled_output else "pooled_embeddings"
                )
                text_encoder_output = outputs[output_name].clone()

            return text_encoder_output

        # Tokenize prompt
        text_encoder_output = tokenize(prompt, max_sequence_length)

        self.profile_stop(encoder)
        return text_encoder_output.to(torch.float16) if self.fp16 else text_encoder_output.to(torch.bfloat16) if self.bf16 else text_encoder_output

    def denoise_latent(
        self,
        latents,
        timesteps,
        text_embeddings,
        pooled_embeddings,
        text_ids,
        latent_image_ids,
        denoiser="transformer",
        guidance=None,
        control_latent=None,
    ):
        do_autocast = self.torch_inference != "" and self.models[denoiser].fp16
        with torch.autocast("cuda", enabled=do_autocast):
            self.profile_start(denoiser, color="blue")

            # handle guidance
            if self.models[denoiser].config["guidance_embeds"] and guidance is None:
                guidance = torch.full(
                    [1], self.guidance_scale, device=self.device, dtype=torch.float32
                )
                guidance = guidance.expand(latents.shape[0])

            for step_index, timestep in enumerate(timesteps):
                # Prepare latents
                latents_input = latents if control_latent is None else torch.cat((latents, control_latent), dim=-1)

                # prepare inputs
                timestep_inp = timestep.expand(latents.shape[0]).to(latents_input.dtype)

                params = {
                    "hidden_states": latents_input,
                    "timestep": timestep_inp / 1000,
                    "pooled_projections": pooled_embeddings,
                    "encoder_hidden_states": text_embeddings,
                    "txt_ids": text_ids.float(),
                    "img_ids": latent_image_ids.float(),
                }
                if guidance is not None:
                    params.update({"guidance": guidance})

                # Predict the noise residual
                if self.torch_inference or self.torch_fallback[denoiser]:
                    noise_pred = self.torch_models[denoiser](**params)["sample"]
                else:
                    noise_pred = self.run_engine(denoiser, params)["latent"]

                latents = self.scheduler.step(
                    noise_pred, timestep, latents, return_dict=False
                )[0]

        self.profile_stop(denoiser)
        return latents.to(dtype=torch.bfloat16) if self.bf16 else latents.to(dtype=torch.float32)

    def decode_latent(self, latents, decoder="vae"):
        self.profile_start(decoder, color="red")
        cast_to = torch.float16 if self.models[decoder].fp16 else torch.bfloat16 if self.models[decoder].bf16 else torch.float32
        latents = latents.to(dtype=cast_to)
        if self.torch_inference or self.torch_fallback[decoder]:
            images = self.torch_models[decoder](latents, return_dict=False)[0]
        else:
            images = self.run_engine(decoder, {"latent": latents})["images"]
        self.profile_stop(decoder)
        return images

    def print_summary(self, denoising_steps, walltime_ms, batch_size):
        print("|-----------------|--------------|")
        print("| {:^15} | {:^12} |".format("Module", "Latency"))
        print("|-----------------|--------------|")
        print(
            "| {:^15} | {:>9.2f} ms |".format(
                "CLIP",
                cudart.cudaEventElapsedTime(
                    self.events["clip"][0], self.events["clip"][1]
                )[1],
            )
        )
        print(
            "| {:^15} | {:>9.2f} ms |".format(
                "T5",
                cudart.cudaEventElapsedTime(self.events["t5"][0], self.events["t5"][1])[
                    1
                ],
            )
        )
        if "vae_encoder" in self.stages:
            print(
                "| {:^15} | {:>9.2f} ms |".format(
                    "VAE-Enc",
                    cudart.cudaEventElapsedTime(
                        self.events["vae_encoder"][0], self.events["vae_encoder"][1]
                    )[1],
                )
            )
        print(
            "| {:^15} | {:>9.2f} ms |".format(
                "Transformer x " + str(denoising_steps),
                cudart.cudaEventElapsedTime(
                    self.events["transformer"][0], self.events["transformer"][1]
                )[1],
            )
        )
        print(
            "| {:^15} | {:>9.2f} ms |".format(
                "VAE-Dec",
                cudart.cudaEventElapsedTime(
                    self.events["vae"][0], self.events["vae"][1]
                )[1],
            )
        )
        print("|-----------------|--------------|")
        print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
        print("|-----------------|--------------|")
        print("Throughput: {:.5f} image/s".format(batch_size * 1000.0 / walltime_ms))

    def _check_integrity(self, images):
        integrity_checker = PixtralContentFilter(self.device)
        for image in images:
            image_ = np.array(image) / 255.0
            image_ = 2 * image_ - 1
            image_ = torch.from_numpy(image_).to(self.device, dtype=torch.float32).permute(0, 3, 1, 2)
            if integrity_checker.test_image(image_):
                raise ValueError("Your image has been flagged. Choose another prompt/image or try again.")

    def save_images(self, prompt, images, check_integrity=False):
        if check_integrity:
            self._check_integrity(images)
        for image in images:
            self.save_image(image, self.pipeline_type.name.lower(), prompt, self.seed)

    def infer(
        self,
        prompt,
        prompt2,
        image_height,
        image_width,
        input_image=None,
        image_strength=1.0,
        control_image=None,
        warmup=False,
        save_image=True,
    ):
        """
        Run the diffusion pipeline.

        Args:
            prompt (str):
                The text prompt to guide image generation.
            prompt2 (str):
                The prompt to be sent to the T5 tokenizer and text encoder
            image_height (int):
                Height (in pixels) of the image to be generated. Must be a multiple of 8.
            image_width (int):
                Width (in pixels) of the image to be generated. Must be a multiple of 8.
            input_image (PIL.Image.Image):
                `Image` representing an image batch to be used as the starting point.
            image_strength (`float`, *optional*, defaults to 1.0):
                Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
                starting point and more noise is added the higher the `strength`. The number of denoising steps depends
                on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
                process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
                essentially ignores `image`.
            control_image (PIL.Image.Image):
                The ControlNet input condition to provide guidance to the `transformer` for generation.
            warmup (bool):
                Indicate if this is a warmup run.
            save_image (bool):
                Save the generated image (if applicable)
        """
        assert len(prompt) == len(prompt2)
        batch_size = len(prompt)

        # Spatial dimensions of latent tensor
        latent_height = 2 * (int(image_height) // self.vae_scale_factor)
        latent_width = 2 * (int(image_width) // self.vae_scale_factor)

        num_inference_steps = self.denoising_steps
        latent_kwargs = {}

        with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
            torch.cuda.synchronize()
            e2e_tic = time.perf_counter()

            num_channels_latents = self.models["transformer"].config["in_channels"] // 4
            if control_image:
                num_channels_latents = self.models["transformer"].config["in_channels"] // 8

                # Prepare control latents
                control_image = self.prepare_image(
                    image=control_image,
                    width=image_width,
                    height=image_height,
                    batch_size=batch_size,
                    num_images_per_prompt=1,
                    device=self.device,
                    dtype=torch.float16 if self.models["vae"].fp16 else torch.bfloat16 if self.models["vae"].bf16 else torch.float32,
                )

                if control_image.ndim == 4:
                    with self.model_memory_manager(["vae_encoder"], low_vram=self.low_vram):
                        control_image = self.encode_image(control_image)

                    height_control_image, width_control_image = control_image.shape[2:]
                    control_image = self._pack_latents(
                        control_image,
                        batch_size,
                        num_channels_latents,
                        height_control_image,
                        width_control_image,
                    )

            # CLIP and T5 text encoder(s)
            with self.model_memory_manager(["clip", "t5"], low_vram=self.low_vram):
                pooled_embeddings = self.encode_prompt(prompt, pooled_output=True)
                text_embeddings = self.encode_prompt(
                    prompt2, encoder="t5", max_sequence_length=self.max_sequence_length
                )
                text_ids = torch.zeros(text_embeddings.shape[1], 3).to(
                    device=self.device, dtype=text_embeddings.dtype
                )

            # Prepare timesteps
            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
            image_seq_len = (latent_height // 2) * (latent_width // 2)
            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 = None
            # TODO: support custom timesteps
            if timesteps is not None:
                if (
                    "timesteps"
                    not in inspect.signature(self.scheduler.set_timesteps).parameters
                ):
                    raise ValueError(
                        f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom"
                        f" timestep schedules. Please check whether you are using the correct scheduler."
                    )
                self.scheduler.set_timesteps(timesteps=timesteps, device=self.device)
                assert self.denoising_steps == len(self.scheduler.timesteps)
            else:
                self.scheduler.set_timesteps(sigmas=sigmas, mu=mu, device=self.device)
            timesteps = self.scheduler.timesteps.to(self.device)
            num_inference_steps = len(timesteps)

            # Pre-process input image and timestep for the img2img pipeline
            if input_image:
                input_image = self.image_processor.preprocess(input_image, height=image_height, width=image_width).to(
                    self.device
                )
                with self.model_memory_manager(["vae_encoder"], low_vram=self.low_vram):
                    image_latents = self.encode_image(input_image)

                timesteps, num_inference_steps = self.get_timesteps(self.denoising_steps, image_strength)
                if num_inference_steps < 1:
                    raise ValueError(
                        f"After adjusting the num_inference_steps by strength parameter: {image_strength}, the number of pipeline"
                        f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline."
                    )
                latent_timestep = timesteps[:1].repeat(batch_size)

                latent_kwargs.update({"image_latents": image_latents, "latent_timestep": latent_timestep})

            # Initialize latents
            latents, latent_image_ids = self.initialize_latents(
                batch_size=batch_size,
                num_channels_latents=num_channels_latents,
                latent_height=latent_height,
                latent_width=latent_width,
                latents_dtype=torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32,
                **latent_kwargs,
            )

            # DiT denoiser
            with self.model_memory_manager(["transformer"], low_vram=self.low_vram):
                latents = self.denoise_latent(
                    latents,
                    timesteps,
                    text_embeddings,
                    pooled_embeddings,
                    text_ids,
                    latent_image_ids,
                    control_latent=control_image,
                )

            # VAE decode latent
            with self.model_memory_manager(["vae"], low_vram=self.low_vram):
                latents = self._unpack_latents(
                    latents, image_height, image_width, self.vae_scale_factor
                )
                latents = (
                    latents / self.models["vae"].config["scaling_factor"]
                ) + self.models["vae"].config["shift_factor"]
                images = self.decode_latent(latents)

            torch.cuda.synchronize()
            e2e_toc = time.perf_counter()

        walltime_ms = (e2e_toc - e2e_tic) * 1000.0
        if not warmup:
            self.print_summary(num_inference_steps, walltime_ms, batch_size)
            if not self.return_latents and save_image:
                # post-process images
                images = (
                    ((images + 1) * 255 / 2)
                    .clamp(0, 255)
                    .detach()
                    .permute(0, 2, 3, 1)
                    .round()
                    .type(torch.uint8)
                    .cpu()
                    .numpy()
                )

        return (latents, walltime_ms) if self.return_latents else (images, walltime_ms)

    def run(
        self,
        prompt,
        prompt2,
        height,
        width,
        batch_count,
        num_warmup_runs,
        use_cuda_graph,
        **kwargs,
    ):
        if self.low_vram and self.use_cuda_graph:
            print("[W] Using low_vram, use_cuda_graph will be disabled")
            self.use_cuda_graph = False
        num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
        if num_warmup_runs > 0:
            print("[I] Warming up ..")
            for _ in range(num_warmup_runs):
                self.infer(prompt, prompt2, height, width, warmup=True, **kwargs)

        images = []
        for _ in range(batch_count):
            print("[I] Running Flux pipeline")
            if self.nvtx_profile:
                cudart.cudaProfilerStart()
            image, _ = self.infer(prompt, prompt2, height, width, warmup=False, **kwargs)
            images.append(image)
            if self.nvtx_profile:
                cudart.cudaProfilerStop()
        return images


class FluxKontextPipeline(FluxPipeline):
    """
    Application showcasing the acceleration of Flux Kontext pipelines using Nvidia TensorRT.
    This pipeline handles the specific logic for the Kontext model variant.
    """

    def __init__(
        self,
        kontext_image,
        version="flux.1-kontext-dev",
        pipeline_type=PIPELINE_TYPE.IMG2IMG,
        guidance_scale=3.5,
        max_sequence_length=512,
        **kwargs,
    ):
        """
        Initializes the Flux Kontext pipeline.

        Args:
            kontext_image (`PIL.Image.Image`):
                Kontext Image to be edited.
            version (`str`, defaults to `flux.1-kontext-dev`)
                Version of the underlying Flux Kontext model.
            guidance_scale (`float`, defaults to 3.5):
                Guidance scale is enabled by setting as > 1.
                Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality.
            max_sequence_length (`int`, defaults to 512):
                Maximum sequence length to use with the `prompt`.
        """
        super().__init__(
            version=version,
            pipeline_type=pipeline_type,
            guidance_scale=guidance_scale,
            max_sequence_length=max_sequence_length,
            **kwargs,
        )
        self.kontext_image = kontext_image
        # WAR to avoid RuntimeError: ScalarType ComplexDouble is an unexpected tensor scalar type during CLIP export
        self.config["clip_torch_fallback"] = True

    @classmethod
    def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> FluxKontextPipeline:
        """Factory method to construct a `FluxKontextPipeline` object from parsed arguments.

        Overrides:
            FluxPipeline.FromArgs
        """
        MAX_BATCH_SIZE = 4
        DEVICE = "cuda"
        DO_RETURN_LATENTS = False

        # Resolve all paths.
        dd_path = path_module.resolve_path(
            cls.get_model_names(pipeline_type), args, pipeline_type, cls._get_pipeline_uid(args.version)
        )

        return cls(
            dd_path=dd_path,
            version=args.version,
            pipeline_type=pipeline_type,
            guidance_scale=args.guidance_scale,
            max_sequence_length=args.max_sequence_length,
            bf16=args.bf16,
            low_vram=args.low_vram,
            torch_fallback=args.torch_fallback,
            weight_streaming=args.ws,
            t5_weight_streaming_budget_percentage=args.t5_ws_percentage,
            transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage,
            max_batch_size=MAX_BATCH_SIZE,
            denoising_steps=args.denoising_steps,
            scheduler=args.scheduler,
            lora_scale=args.lora_scale,
            lora_weight=args.lora_weight,
            lora_path=args.lora_path,
            kontext_image=args.kontext_image if hasattr(args, "kontext_image") else None,
            device=DEVICE,
            output_dir=args.output_dir,
            hf_token=args.hf_token,
            verbose=args.verbose,
            nvtx_profile=args.nvtx_profile,
            use_cuda_graph=args.use_cuda_graph,
            framework_model_dir=args.framework_model_dir,
            return_latents=DO_RETURN_LATENTS,
            torch_inference=args.torch_inference,
        )

    def initialize_latents(
        self,
        batch_size,
        num_channels_latents,
        latent_height,
        latent_width,
        latent_timestep=None,
        image_latents=None,
        latents_dtype=torch.float32,
    ):
        """
        Initialize latents for Kontext pipeline.
        Overrides FluxPipeline.initialize_latents to handle Kontext-specific logic.
        """
        latents_dtype = latents_dtype  # text_embeddings.dtype
        latents_shape = (batch_size, num_channels_latents, latent_height, latent_width)
        latents = torch.randn(
            latents_shape,
            device=self.device,
            dtype=latents_dtype,
            generator=self.generator,
        )

        image_ids = None
        if image_latents is not None:
            image_latents = torch.cat([image_latents], dim=0).to(latents_dtype)
            image_latent_height, image_latent_width = image_latents.shape[2:]
            image_latents = self._pack_latents(
                image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
            )
            image_ids = self._prepare_latent_image_ids(
                image_latent_height, image_latent_width, latents_dtype, self.device
            )
            # image ids are the same as latent ids with the first dimension set to 1 instead of 0
            image_ids[..., 0] = 1

        latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width)

        latent_ids = self._prepare_latent_image_ids(latent_height, latent_width, latents_dtype, self.device)

        latent_image_ids = torch.cat([latent_ids, image_ids], dim=0) if image_ids is not None else latent_ids

        return latents, image_latents, latent_image_ids

    def denoise_latent(
        self,
        latents,
        timesteps,
        text_embeddings,
        pooled_embeddings,
        text_ids,
        latent_image_ids,
        image_latents,
        denoiser="transformer",
        guidance=None,
    ):
        """
        Denoise latents for Kontext pipeline.
        Overrides FluxPipeline.denoise_latent to handle Kontext-specific logic.
        """
        do_autocast = self.torch_inference != "" and self.models[denoiser].fp16
        with torch.autocast("cuda", enabled=do_autocast):
            self.profile_start(denoiser, color="blue")

            # handle guidance
            if self.models[denoiser].config["guidance_embeds"] and guidance is None:
                guidance = torch.full([1], self.guidance_scale, device=self.device, dtype=torch.float32)
                guidance = guidance.expand(latents.shape[0])

            for step_index, timestep in enumerate(timesteps):
                # Kontext-specific: concatenate image_latents along dim=1
                latents_input = torch.cat([latents, image_latents], dim=1)

                # prepare inputs
                timestep_inp = timestep.expand(latents.shape[0]).to(latents_input.dtype)

                params = {
                    "hidden_states": latents_input,
                    "timestep": timestep_inp / 1000,
                    "pooled_projections": pooled_embeddings,
                    "encoder_hidden_states": text_embeddings,
                    "txt_ids": text_ids.float(),
                    "img_ids": latent_image_ids.float(),
                }
                if guidance is not None:
                    params.update({"guidance": guidance})

                # Predict the noise residual
                if self.torch_inference or self.torch_fallback[denoiser]:
                    noise_pred = self.torch_models[denoiser](**params)["sample"]
                else:
                    noise_pred = self.run_engine(denoiser, params)["latent"]

                # Kontext-specific: extract only the latent part of the prediction
                noise_pred = noise_pred[:, : latents.size(1)]

                latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False)[0]

        self.profile_stop(denoiser)
        return latents.to(dtype=torch.bfloat16) if self.bf16 else latents.to(dtype=torch.float32)

    def _get_resize_dimensions(self, input_image):
        """
        Preprocess input image for Kontext pipeline using preferred resolutions.
        Uses FluxKontextUtil to get the standardized list of preferred resolutions.
        """
        multiple_of = self.vae_scale_factor * 2
        resize_height, resize_width = self.image_processor.get_default_height_width(input_image)
        aspect_ratio = resize_width / resize_height
        # Kontext is trained on specific resolutions, using one of them is recommended
        preferred_resolutions = FluxKontextUtil.get_preferred_resolutions()
        _, resize_width, resize_height = min(
            (abs(aspect_ratio - w / h), w, h) for w, h in preferred_resolutions
        )
        resize_width = resize_width // multiple_of * multiple_of
        resize_height = resize_height // multiple_of * multiple_of

        return resize_height, resize_width

    def infer(
        self,
        prompt,
        prompt2,
        image_height,
        image_width,
        image_strength=1.0,
        warmup=False,
        save_image=True,
    ):
        """
        Run the Kontext diffusion pipeline.

        Args:
            prompt (str):
                The text prompt to guide image generation.
            prompt2 (str):
                The prompt to be sent to the T5 tokenizer and text encoder
            image_height (int):
                Height (in pixels) of the image to be generated. Must be a multiple of 8.
            image_width (int):
                Width (in pixels) of the image to be generated. Must be a multiple of 8.
            image_strength (`float`, *optional*, defaults to 1.0):
                Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
                starting point and more noise is added the higher the `strength`. The number of denoising steps depends
                on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
                process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
                essentially ignores `image`.
            warmup (bool):
                Indicate if this is a warmup run.
            save_image (bool):
                Save the generated image (if applicable)
        """
        assert len(prompt) == len(prompt2)
        batch_size = len(prompt)

        # Spatial dimensions of latent tensor
        latent_height = 2 * (int(image_height) // self.vae_scale_factor)
        latent_width = 2 * (int(image_width) // self.vae_scale_factor)

        num_inference_steps = self.denoising_steps
        latent_kwargs = {}

        with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
            torch.cuda.synchronize()
            e2e_tic = time.perf_counter()

            num_channels_latents = self.models["transformer"].config["in_channels"] // 4

            # CLIP and T5 text encoder(s)
            with self.model_memory_manager(["clip", "t5"], low_vram=self.low_vram):
                pooled_embeddings = self.encode_prompt(prompt, pooled_output=True)
                text_embeddings = self.encode_prompt(
                    prompt2, encoder="t5", max_sequence_length=self.max_sequence_length
                )
                text_ids = torch.zeros(text_embeddings.shape[1], 3).to(device=self.device, dtype=text_embeddings.dtype)

            # Prepare timesteps
            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
            image_seq_len = (latent_height // 2) * (latent_width // 2)
            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 = None
            # TODO: support custom timesteps
            if timesteps is not None:
                if "timesteps" not in inspect.signature(self.scheduler.set_timesteps).parameters:
                    raise ValueError(
                        f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom"
                        f" timestep schedules. Please check whether you are using the correct scheduler."
                    )
                self.scheduler.set_timesteps(timesteps=timesteps, device=self.device)
                assert self.denoising_steps == len(self.scheduler.timesteps)
            else:
                self.scheduler.set_timesteps(sigmas=sigmas, mu=mu, device=self.device)
            timesteps = self.scheduler.timesteps.to(self.device)
            num_inference_steps = len(timesteps)

            # Pre-process kontext image and timestep for the img2img pipeline
            if self.kontext_image:
                # Kontext-specific image preprocessing
                kontext_image = self.image_processor.resize(self.kontext_image, self.resize_height, self.resize_width)

                kontext_image = self.image_processor.preprocess(
                    kontext_image, height=self.resize_height, width=self.resize_width
                ).to(self.device)
                with self.model_memory_manager(["vae_encoder"], low_vram=self.low_vram):
                    image_latents = self.encode_image(kontext_image)

                timesteps, num_inference_steps = self.get_timesteps(self.denoising_steps, image_strength)
                if num_inference_steps < 1:
                    raise ValueError(
                        f"After adjusting the num_inference_steps by strength parameter: {image_strength}, the number of pipeline"
                        f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline."
                    )
                latent_timestep = timesteps[:1].repeat(batch_size)

                latent_kwargs.update({"image_latents": image_latents, "latent_timestep": latent_timestep})

            # Initialize latents
            latents, image_latents, latent_ids = self.initialize_latents(
                batch_size=batch_size,
                num_channels_latents=num_channels_latents,
                latent_height=latent_height,
                latent_width=latent_width,
                latents_dtype=torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32,
                **latent_kwargs,
            )

            # DiT denoiser
            with self.model_memory_manager(["transformer"], low_vram=self.low_vram):
                latents = self.denoise_latent(
                    latents,
                    timesteps,
                    text_embeddings,
                    pooled_embeddings,
                    text_ids,
                    latent_ids,
                    image_latents,
                )

            # VAE decode latent
            with self.model_memory_manager(["vae"], low_vram=self.low_vram):
                latents = self._unpack_latents(latents, image_height, image_width, self.vae_scale_factor)
                latents = (latents / self.models["vae"].config["scaling_factor"]) + self.models["vae"].config[
                    "shift_factor"
                ]
                images = self.decode_latent(latents)

            torch.cuda.synchronize()
            e2e_toc = time.perf_counter()

        walltime_ms = (e2e_toc - e2e_tic) * 1000.0
        if not warmup:
            self.print_summary(num_inference_steps, walltime_ms, batch_size)
            if not self.return_latents and save_image:
                # post-process images
                images = (
                    ((images + 1) * 255 / 2)
                    .clamp(0, 255)
                    .detach()
                    .permute(0, 2, 3, 1)
                    .round()
                    .type(torch.uint8)
                    .cpu()
                    .numpy()
                )

        return (latents, walltime_ms) if self.return_latents else (images, walltime_ms)