#
# 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 abc
import argparse
import gc
import json
import os
import pathlib
import sys
from abc import ABC, abstractmethod
from typing import Any, List

import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
import nvtx
import torch
from cuda.bindings import runtime as cudart
from diffusers import (
    DDIMScheduler,
    DDPMScheduler,
    DDPMWuerstchenScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    FlowMatchEulerDiscreteScheduler,
    LCMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UniPCMultistepScheduler,
)
from torch.utils.data import DataLoader

import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
    make_scheduler,
    merge_loras,
    unload_torch_model,
)
from demo_diffusion.pipeline.calibrate import load_calib_prompts
from demo_diffusion.pipeline.model_memory_manager import ModelMemoryManager
from demo_diffusion.pipeline.type import PIPELINE_TYPE
from demo_diffusion.utils_modelopt import (
    SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG,
    SD_FP8_FP16_DEFAULT_CONFIG,
    SD_FP8_FP32_DEFAULT_CONFIG,
    PromptImageDataset,
    SameSizeSampler,
    check_lora,
    custom_collate,
    filter_func,
    filter_func_no_proj_out,
    fp8_mha_disable,
    generate_fp8_scales,
    get_int8_config,
    infinite_dataloader,
    quantize_lvl,
    set_fmha,
    set_quant_precision,
)


class DiffusionPipeline(ABC):
    """
    Application showcasing the acceleration of Stable Diffusion pipelines using NVidia TensorRT.
    """
    VALID_DIFFUSION_PIPELINES = (
        "1.4",
        "dreamshaper-7",
        "xl-1.0",
        "xl-turbo",
        "svd-xt-1.1",
        "sd3",
        "3.5-medium",
        "3.5-large",
        "cascade",
        "flux.1-dev",
        "flux.1-dev-canny",
        "flux.1-dev-depth",
        "flux.1-schnell",
        "flux.1-kontext-dev",
        "wan2.2-t2v-a14b",
        "cosmos-predict2-2b-text2image",
        "cosmos-predict2-14b-text2image",
        "cosmos-predict2-2b-video2world",
        "cosmos-predict2-14b-video2world",
    )
    SCHEDULER_DEFAULTS = {
        "1.4": "PNDM",
        "dreamshaper-7": "PNDM",
        "xl-1.0": "Euler",
        "xl-turbo": "EulerA",
        "3.5-large": "FlowMatchEuler",
        "3.5-medium": "FlowMatchEuler",
        "svd-xt-1.1": "Euler",
        "cascade": "DDPMWuerstchen",
        "flux.1-dev": "FlowMatchEuler",
        "flux.1-dev-canny": "FlowMatchEuler",
        "flux.1-dev-depth": "FlowMatchEuler",
        "flux.1-schnell": "FlowMatchEuler",
        "flux.1-kontext-dev": "FlowMatchEuler",
        "wan2.2-t2v-a14b": "UniPC",
        "cosmos-predict2-2b-text2image": "FlowMatchEuler",
        "cosmos-predict2-14b-text2image": "FlowMatchEuler",
        "cosmos-predict2-2b-video2world": "FlowMatchEuler",
        "cosmos-predict2-14b-video2world": "FlowMatchEuler",
    }

    def __init__(
        self,
        dd_path,
        version="1.4",
        pipeline_type=PIPELINE_TYPE.TXT2IMG,
        bf16=False,
        max_batch_size=16,
        denoising_steps=30,
        scheduler=None,
        device="cuda",
        output_dir=".",
        hf_token=None,
        verbose=False,
        nvtx_profile=False,
        use_cuda_graph=False,
        framework_model_dir="pytorch_model",
        return_latents=False,
        low_vram=False,
        torch_inference="",
        torch_fallback=None,
        weight_streaming=False,
        text_encoder_weight_streaming_budget_percentage=None,
        denoiser_weight_streaming_budget_percentage=None,
        controlnet=None,
    ):
        """
        Initializes the Diffusion pipeline.

        Args:
            dd_path (load_module.DDPath): DDPath object that contains all paths used in DemoDiffusion.
            version (str):
                The version of the pipeline. Should be one of the values listed in DiffusionPipeline.VALID_DIFFUSION_PIPELINES.
            pipeline_type (PIPELINE_TYPE):
                Task performed by the current pipeline. Should be one of PIPELINE_TYPE.__members__.
            max_batch_size (int):
                Maximum batch size for dynamic batch engine.
            bf16 (`bool`, defaults to False):
                Whether to run the pipeline in BFloat16 precision.
            denoising_steps (int):
                The number of denoising steps.
                More denoising steps usually lead to a higher quality image at the expense of slower inference.
            scheduler (str):
                The scheduler to guide the denoising process. Must be one of the values listed in DiffusionPipeline.SCHEDULER_DEFAULTS.values().
            device (str):
                PyTorch device to run inference. Default: 'cuda'.
            output_dir (str):
                Output directory for log files and image artifacts.
            hf_token (str):
                HuggingFace User Access Token to use for downloading Stable Diffusion model checkpoints.
            verbose (bool):
                Enable verbose logging.
            nvtx_profile (bool):
                Insert NVTX profiling markers.
            use_cuda_graph (bool):
                Use CUDA graph to capture engine execution and then launch inference.
            framework_model_dir (str):
                cache directory for framework checkpoints.
            return_latents (bool):
                Skip decoding the image and return latents instead.
            low_vram (bool):
                [FLUX only] Optimize for low VRAM usage, possibly at the expense of inference performance. Disabled by default.
            torch_inference (str):
                Run inference with PyTorch (using specified compilation mode) instead of TensorRT. The compilation mode specified should be one of ['eager', 'reduce-overhead', 'max-autotune'].
            torch_fallback (str):
                [FLUX only] Comma separated list of models to be inferenced using PyTorch instead of TRT. For example --torch-fallback t5,transformer. If --torch-inference set, this parameter will be ignored.
            weight_streaming (`bool`, defaults to False):
                Whether to enable weight streaming during TensorRT engine build.
            text_encoder_ws_budget_percentage (`int`, defaults to None):
                Weight streaming budget as a percentage of the size of total streamable weights for the text encoder model.
            denoiser_weight_streaming_budget_percentage (`int`, defaults to None):
                Weight streaming budget as a percentage of the size of total streamable weights for the denoiser model.
            controlnet (str, defaults to None):
                Type of ControlNet to use for the pipeline.
        """
        self.bf16 = bf16
        self.dd_path = dd_path

        self.denoising_steps = denoising_steps
        self.max_batch_size = max_batch_size

        self.framework_model_dir = framework_model_dir
        self.output_dir = output_dir
        for directory in [self.framework_model_dir, self.output_dir]:
            if not os.path.exists(directory):
                print(f"[I] Create directory: {directory}")
                pathlib.Path(directory).mkdir(parents=True)

        self.hf_token = hf_token
        self.device = device
        self.verbose = verbose
        self.nvtx_profile = nvtx_profile

        self.version = version
        self.pipeline_type = pipeline_type
        self.return_latents = return_latents

        self.low_vram = low_vram
        self.weight_streaming = weight_streaming
        self.text_encoder_weight_streaming_budget_percentage = text_encoder_weight_streaming_budget_percentage
        self.denoiser_weight_streaming_budget_percentage = denoiser_weight_streaming_budget_percentage

        self.stages = self.get_model_names(self.pipeline_type, controlnet)
        # config to store additional info
        self.config = {}
        if torch_fallback:
            assert type(torch_fallback) is list
            for model_name in torch_fallback:
                if model_name not in self.stages:
                    raise ValueError(f'Model "{model_name}" set in --torch-fallback does not exist')
                self.config[model_name.replace("-", "_") + "_torch_fallback"] = True
                print(f"[I] Setting torch_fallback for {model_name} model.")

        if not scheduler:
            scheduler = 'UniPC' if self.pipeline_type.is_controlnet() and not self.version == "3.5-large" else self.SCHEDULER_DEFAULTS.get(version, 'DDIM')
            print(f"[I] Autoselected scheduler: {scheduler}")

        scheduler_class_map = {
            "DDIM" : DDIMScheduler,
            "DDPM" : DDPMScheduler,
            "EulerA" : EulerAncestralDiscreteScheduler,
            "Euler" : EulerDiscreteScheduler,
            "LCM" : LCMScheduler,
            "LMSD" : LMSDiscreteScheduler,
            "PNDM" : PNDMScheduler,
            "UniPC" : UniPCMultistepScheduler,
            "DDPMWuerstchen" : DDPMWuerstchenScheduler,
            "FlowMatchEuler": FlowMatchEulerDiscreteScheduler,
        }
        try:
            scheduler_class = scheduler_class_map[scheduler]
        except KeyError:
            raise ValueError(
                f"Unsupported scheduler {scheduler}.  Should be one of {list(scheduler_class_map.keys())}."
            )
        self.scheduler = make_scheduler(scheduler_class, version, pipeline_type, hf_token, framework_model_dir)

        self.torch_inference = torch_inference
        if self.torch_inference:
            torch._inductor.config.conv_1x1_as_mm = True
            torch._inductor.config.coordinate_descent_tuning = True
            torch._inductor.config.epilogue_fusion = False
            torch._inductor.config.coordinate_descent_check_all_directions = True
        self.use_cuda_graph = use_cuda_graph

        # initialized in load_engines()
        self.models = {}
        self.torch_models = {}
        self.engine = {}
        self.shape_dicts = {}
        self.shared_device_memory = None
        self.lora_loader = None

        # initialized in load_resources()
        self.events = {}
        self.generator = None
        self.markers = {}
        self.seed = None
        self.stream = None
        self.tokenizer = None

    def model_memory_manager(self, model_names, low_vram=False):
        return ModelMemoryManager(self, model_names, low_vram)

    @classmethod
    @abc.abstractmethod
    def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> DiffusionPipeline:
        """Factory method to construct a concrete pipeline object from parsed arguments."""
        raise NotImplementedError("FromArgs cannot be called from the abstract base class.")

    @classmethod
    @abc.abstractmethod
    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."""
        raise NotImplementedError("get_model_names cannot be called from the abstract base class.")

    @classmethod
    def _get_pipeline_uid(cls, version: str) -> str:
        """Return the unique ID of this pipeline.

        This is typically used to determine the default path for things like engine files, artifacts caches, etc.
        """
        return f"{cls.__name__}_{version}"

    def profile_start(self, name, color="blue", domain=None):
        if self.nvtx_profile:
            self.markers[name] = nvtx.start_range(message=name, color=color, domain=domain)
        if name in self.events:
            cudart.cudaEventRecord(self.events[name][0], 0)

    def profile_stop(self, name):
        if name in self.events:
            cudart.cudaEventRecord(self.events[name][1], 0)
        if self.nvtx_profile:
            nvtx.end_range(self.markers[name])

    def load_resources(self, image_height, image_width, batch_size, seed):
        # Initialize noise generator
        if seed is not None:
            self.seed = seed
            self.generator = torch.Generator(device="cuda").manual_seed(seed)

        # Create CUDA events and stream
        for stage in self.stages:
            self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
        self.stream = cudart.cudaStreamCreate()[1]

        # Allocate TensorRT I/O buffers
        if not self.torch_inference:
            for model_name, obj in self.models.items():
                if self.torch_fallback[model_name]:
                    continue
                shape_dict = obj.get_shape_dict(
                    batch_size, image_height, image_width, 
                    **({'num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {})
                )
                
                self.shape_dicts[model_name] = shape_dict
                if not self.low_vram:
                    self.engine[model_name].allocate_buffers(shape_dict=shape_dict, device=self.device)

    @abstractmethod
    def _initialize_models(self, *args, **kwargs):
        raise NotImplementedError("Please Implement the _initialize_models method")

    def _prepare_model_configs(
        self,
        enable_refit,
        int8,
        fp8,
        fp4
    ):
        model_names = self.models.keys()
        self.torch_fallback = dict(zip(model_names, [self.torch_inference or self.config.get(model_name.replace('-','_')+'_torch_fallback', False) for model_name in model_names]))

        configs = {}
        for model_name in model_names:
            # Initialize config
            do_engine_refit = enable_refit and not self.pipeline_type.is_sd_xl_refiner() and any(model_name.startswith(prefix) for prefix in ("unet", "transformer"))
            do_lora_merge = not enable_refit and self.lora_loader and any(model_name.startswith(prefix) for prefix in ("unet", "transformer"))

            config = {
                "do_engine_refit": do_engine_refit,
                "do_lora_merge": do_lora_merge,
                "use_int8": False,
                "use_fp8": False,
                'use_fp4': False,
            }

            # TODO: Move this to when arguments are first being validated in dd_argparse.py
            # 8-bit/4-bit precision inference
            if int8:
                assert self.pipeline_type.is_sd_xl_base() or self.version in [
                    "1.4",
                ], "int8 quantization only supported for SDXL and SD1.4 pipeline"
                if (self.pipeline_type.is_sd_xl() and model_name == "unetxl") or (model_name == "unet"):
                    config["use_int8"] = True

            elif fp8:
                assert (
                    self.pipeline_type.is_sd_xl()
                    or self.version in ["1.4"]
                    or self.version.startswith("flux.1")
                    or self.version.startswith("3.5-large")
                ), "fp8 quantization only supported for SDXL, SD1.4, SD3.5-large and FLUX pipelines"
                if (
                    (self.pipeline_type.is_sd_xl() and model_name == "unetxl")
                    or (self.version.startswith("flux.1") and model_name == "transformer")
                    or (
                        self.version.startswith("3.5-large")
                        and ("transformer" in model_name or "controlnet" in model_name)
                    )
                    or (model_name == "unet")
                ):
                    config["use_fp8"] = True
            elif fp4:
                config['use_fp4'] = True

            # Setup paths
            config["onnx_path"] = self.dd_path.model_name_to_unoptimized_onnx_path[model_name]
            config["onnx_opt_path"] = self.dd_path.model_name_to_optimized_onnx_path[model_name]
            config["engine_path"] = self.dd_path.model_name_to_engine_path[model_name]
            config["weights_map_path"] = (
                self.dd_path.model_name_to_weights_map_path[model_name] if config["do_engine_refit"] else None
            )
            config["state_dict_path"] = self.dd_path.model_name_to_quantized_model_state_dict_path[model_name]
            config["refit_weights_path"] = self.dd_path.model_name_to_refit_weights_path[model_name]

            configs[model_name] = config

        return configs

    def _calibrate_and_save_model(
            self,
            pipeline,
            model,
            model_config,
            quantization_level,
            quantization_percentile,
            quantization_alpha,
            calibration_size,
            calib_batch_size,
            enable_lora_merge = False,
            **kwargs):
        print(f"[I] Calibrated weights not found, generating {model_config['state_dict_path']}")

        # TODO check size > calibration_size
        def do_calibrate(pipeline, calibration_prompts, **kwargs):
            for i_th, prompts in enumerate(calibration_prompts):
                if i_th >= kwargs["calib_size"]:
                    return
                if kwargs["model_id"] in ("flux.1-dev", "flux.1-schnell"):
                    common_args = {
                        "prompt": prompts,
                        "prompt_2": prompts,
                        "num_inference_steps": kwargs["n_steps"],
                        "height": kwargs.get("height", 1024),
                        "width": kwargs.get("width", 1024),
                        "guidance_scale": 3.5,
                        "max_sequence_length": 512 if kwargs["model_id"] == "flux.1-dev" else 256,
                    }
                else:
                    common_args = {
                        "prompt": prompts,
                        "num_inference_steps": kwargs["n_steps"],
                        "negative_prompt": ["normal quality, low quality, worst quality, low res, blurry, nsfw, nude"]
                        * len(prompts),
                    }

                pipeline(**common_args).images

        def do_calibrate_img2img(pipeline, dataloader, **kwargs):
            for i_th, (img_conds, prompts) in enumerate(dataloader):
                if i_th >= kwargs["calib_size"]:
                    return

                common_args = {
                    "prompt": list(prompts),
                    "control_image": img_conds,
                    "num_inference_steps": kwargs["n_steps"],
                    "height": img_conds.size(2),
                    "width": img_conds.size(3),
                    "generator": torch.Generator().manual_seed(42),
                    "guidance_scale": 3.5,
                    "max_sequence_length": 512,
                }
                pipeline(**common_args).images

        if self.version in ("flux.1-dev-depth", "flux.1-dev-canny"):
            dataset = PromptImageDataset(
                root_dir=self.calibration_dataset,
            )

            dataloader = DataLoader(
                dataset,
                batch_size=calib_batch_size,
                shuffle=False,
                num_workers=0,
                sampler=SameSizeSampler(dataset=dataset, batch_size=calib_batch_size),
                collate_fn=custom_collate,
            )
        else:
            root_dir = os.path.dirname(os.path.abspath(sys.modules["__main__"].__file__))
            calibration_file = os.path.join(root_dir, "calibration_data", "calibration-prompts.txt")
            calibration_prompts = load_calib_prompts(calib_batch_size, calibration_file)

        def forward_loop(model):
            if self.version not in ("sd3", "flux.1-dev", "flux.1-schnell", "flux.1-dev-depth", "flux.1-dev-canny"):
                pipeline.unet = model
            else:
                pipeline.transformer = model

            if self.version in ("flux.1-dev-depth", "flux.1-dev-canny"):
                do_calibrate_img2img(
                    pipeline=pipeline,
                    dataloader=infinite_dataloader(dataloader),
                    calib_size=calibration_size // calib_batch_size,
                    n_steps=self.denoising_steps,
                    model_id=self.version,
                )
            else:
                do_calibrate(
                    pipeline=pipeline,
                    calibration_prompts=calibration_prompts,
                    calib_size=calibration_size // calib_batch_size,
                    n_steps=self.denoising_steps,
                    model_id=self.version,
                    **kwargs
                )

        print(f"[I] Performing calibration for {calibration_size} steps.")
        if model_config['use_int8']:
            quant_config = get_int8_config(
                model,
                quantization_level,
                quantization_alpha,
                quantization_percentile,
                self.denoising_steps
            )
        elif model_config['use_fp8']:
            if self.version.startswith("flux.1"):
                quant_config = SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG
            else:
                quant_config = SD_FP8_FP16_DEFAULT_CONFIG

        # Handle LoRA
        if enable_lora_merge:
            assert self.lora_loader is not None
            model = merge_loras(model, self.lora_loader)

        check_lora(model)

        if self.version.startswith("flux.1"):
            set_quant_precision(quant_config, "BFloat16")
        mtq.quantize(model, quant_config, forward_loop)
        mto.save(model, model_config['state_dict_path'])

    def _get_quantized_model(
            self,
            obj,
            model_config,
            quantization_level,
            quantization_percentile,
            quantization_alpha,
            calibration_size,
            calib_batch_size,
            enable_lora_merge = False,
            **kwargs):
        pipeline = obj.get_pipeline()
        is_flux = self.version.startswith("flux.1")
        model = pipeline.unet if self.version not in ("sd3", "flux.1-dev", "flux.1-schnell", "flux.1-dev-depth", "flux.1-dev-canny") else pipeline.transformer
        if model_config['use_fp8'] and quantization_level == 4.0:
            set_fmha(model, is_flux=is_flux)

        if not os.path.exists(model_config['state_dict_path']):
            self._calibrate_and_save_model(
                pipeline,
                model,
                model_config,
                quantization_level,
                quantization_percentile,
                quantization_alpha,
                calibration_size,
                calib_batch_size,
                enable_lora_merge,
                **kwargs)
        else:
            mto.restore(model, model_config['state_dict_path'])

        if not os.path.exists(model_config['onnx_path']):
            quantize_lvl(self.version, model, quantization_level)
            if self.version.startswith("flux.1"):
                mtq.disable_quantizer(model, filter_func_no_proj_out)
            else:
                mtq.disable_quantizer(model, filter_func)
            if model_config['use_fp8'] and not self.version.startswith("flux.1"):
                generate_fp8_scales(model)
            if quantization_level == 4.0:
                fp8_mha_disable(model, quantized_mha_output=False) # Remove Q/DQ after BMM2 in MHA
        else:
            model = None

        return model

    @abstractmethod
    def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
        """Download pre-exported ONNX Models"""
        raise NotImplementedError("Please Implement the download_onnx_models method")

    def is_native_export_supported(self, model_config: dict[str, Any]) -> bool:
        """Check if pipeline supports native ONNX export"""
        # Native export is supported by default
        return True

    @staticmethod
    def _fix_bf16_resize_nodes(onnx_opt_path):
        """Cast Resize node I/O for strongly-typed TRT engine builds.
        TRT does not support BF16 for the Resize operator, so inputs are
        cast to FP32 and outputs are cast back to BF16."""
        import onnx
        import onnx_graphsurgeon as gs
        from demo_diffusion.model import load
        from demo_diffusion.utils_modelopt import cast_resize_io

        onnx_graph = onnx.load(onnx_opt_path, load_external_data=True)
        graph = gs.import_onnx(onnx_graph)

        resize_nodes = [n for n in graph.nodes if n.op == "Resize"]
        if not resize_nodes:
            return

        print(f"[I] Fixing {len(resize_nodes)} BF16 Resize node(s) in downloaded model: {onnx_opt_path}")
        cast_resize_io(graph, output_dtype=onnx.TensorProto.BFLOAT16)
        graph.cleanup().toposort()
        onnx_graph = gs.export_onnx(graph)

        if load.onnx_graph_needs_external_data(onnx_graph):
            onnx.save_model(
                onnx_graph,
                onnx_opt_path,
                save_as_external_data=True,
                all_tensors_to_one_file=True,
                convert_attribute=False,
            )
        else:
            onnx.save(onnx_graph, onnx_opt_path)

    def _export_onnx(
        self,
        obj,
        model_name,
        model_config,
        opt_image_height,
        opt_image_width,
        static_shape,
        onnx_opset,
        quantization_level,
        quantization_percentile,
        quantization_alpha,
        calibration_size,
        calib_batch_size,
        onnx_export_only,
        download_onnx_models,
    ):
        # With onnx_export_only True, the export still happens even if the TRT engine exists. However, it will not re-run the export if the onnx exists.
        do_export_onnx = (not os.path.exists(model_config['engine_path']) or onnx_export_only) and not os.path.exists(model_config['onnx_opt_path'])
        do_export_weights_map = model_config['weights_map_path'] and not os.path.exists(model_config['weights_map_path'])

        # If ONNX export is required, either download ONNX models or check if the pipeline supports native ONNX export
        if do_export_onnx:
            if download_onnx_models:
                self.download_onnx_models(model_name, model_config)
                # Fix Resize nodes for strongly-typed TRT builds.
                # Downloaded models bypass optimize(), so apply the fix here.
                if obj.bf16:
                    self._fix_bf16_resize_nodes(model_config['onnx_opt_path'])
                do_export_onnx = False
            else:
                self.is_native_export_supported(model_config)

        dynamo = True if (self.pipeline_type.is_video2world() and model_name == "transformer") or (self.pipeline_type.is_txt2vid() and (model_name in ["transformer", "transformer_2"])) or (self.version.startswith("flux.1") and model_name == "transformer" and obj.fp16) else False
        
        export_kwargs = {
            "static_shape": static_shape,
            "dynamo": dynamo,
            **({'opt_num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {})
        }
        
        if do_export_onnx or do_export_weights_map:
            if not model_config['use_int8'] and not model_config['use_fp8']:
                obj.export_onnx(
                    model_config["onnx_path"],
                    model_config["onnx_opt_path"],
                    onnx_opset,
                    opt_image_height,
                    opt_image_width,
                    enable_lora_merge=model_config["do_lora_merge"],
                    lora_loader=self.lora_loader,
                    **export_kwargs,
                )
            else:
                print(f"[I] Generating quantized ONNX model: {model_config['onnx_path']}")
                quantized_model = self._get_quantized_model(
                    obj,
                    model_config,
                    quantization_level,
                    quantization_percentile,
                    quantization_alpha,
                    calibration_size,
                    calib_batch_size,
                    height=opt_image_width,
                    width=opt_image_width,
                    enable_lora_merge=model_config["do_lora_merge"],
                )
                obj.export_onnx(
                    model_config["onnx_path"],
                    model_config["onnx_opt_path"],
                    onnx_opset,
                    opt_image_height,
                    opt_image_width,
                    custom_model=quantized_model,
                    **export_kwargs,
                )

        # FIXME do_export_weights_map needs ONNX graph
        if do_export_weights_map:
            print(f"[I] Saving weights map: {model_config['weights_map_path']}")
            obj.export_weights_map(model_config['onnx_opt_path'], model_config['weights_map_path'])

    def _build_engine(self, obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache):
        update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
        tf32amp = obj.tf32
        weight_streaming = getattr(obj, 'weight_streaming', False)
        precision_constraints = 'none'
        input_profile = obj.get_input_profile(
            opt_batch_size, opt_image_height, opt_image_width, 
            static_batch=static_batch, static_shape=static_shape,
            **({'num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {})
        )
        
        engine.build(
            model_config["onnx_opt_path"],
            tf32=tf32amp,
            input_profile=input_profile,
            enable_refit=model_config["do_engine_refit"],
            enable_all_tactics=enable_all_tactics,
            timing_cache=timing_cache,
            update_output_names=update_output_names,
            weight_streaming=weight_streaming,
            verbose=self.verbose,
            builder_optimization_level=optimization_level,
            precision_constraints=precision_constraints,
        )

    def _refit_engine(self, obj, model_name, model_config):
        assert model_config['weights_map_path']
        with open(model_config['weights_map_path'], 'r') as fp_wts:
            print(f"[I] Loading weights map: {model_config['weights_map_path']} ")
            [weights_name_mapping, weights_shape_mapping] = json.load(fp_wts)

            if not os.path.exists(model_config['refit_weights_path']):
                model = merge_loras(obj.get_model(), self.lora_loader)
                refit_weights, updated_weight_names = engine_module.get_refit_weights(
                    model.state_dict(), model_config["onnx_opt_path"], weights_name_mapping, weights_shape_mapping
                )
                print(f"[I] Saving refit weights: {model_config['refit_weights_path']}")
                torch.save((refit_weights, updated_weight_names), model_config["refit_weights_path"])
                unload_torch_model(model)
            else:
                print(f"[I] Loading refit weights: {model_config['refit_weights_path']}")
                refit_weights, updated_weight_names = torch.load(model_config['refit_weights_path'])
            self.engine[model_name].refit(refit_weights, updated_weight_names)

    def _load_torch_models(self):
        # Load torch models
        for model_name, obj in self.models.items():
            if self.torch_fallback[model_name]:
                self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference)
                if self.low_vram:
                    self.torch_models[model_name] = self.torch_models[model_name].to('cpu')
                    torch.cuda.empty_cache()

    def load_engines(
        self,
        framework_model_dir,
        onnx_opset,
        opt_batch_size,
        opt_image_height,
        opt_image_width,
        optimization_level=3,
        static_batch=False,
        static_shape=True,
        enable_refit=False,
        enable_all_tactics=False,
        timing_cache=None,
        int8=False,
        fp8=False,
        fp4=False,
        quantization_level=2.5,
        quantization_percentile=1.0,
        quantization_alpha=0.8,
        calibration_size=32,
        calib_batch_size=2,
        onnx_export_only=False,
        download_onnx_models=False,
    ):
        """
        Build and load engines for TensorRT accelerated inference.
        Export ONNX models first, if applicable.

        Args:
            framework_model_dir (str):
                Directory to store the framework model ckpt.
            onnx_opset (int):
                ONNX opset version to export the models.
            opt_batch_size (int):
                Batch size to optimize for during engine building.
            opt_image_height (int):
                Image height to optimize for during engine building. Must be a multiple of 8.
            opt_image_width (int):
                Image width to optimize for during engine building. Must be a multiple of 8.
            optimization_level (int):
                Optimization level to build the TensorRT engine with.
            static_batch (bool):
                Build engine only for specified opt_batch_size.
            static_shape (bool):
                Build engine only for specified opt_image_height & opt_image_width. Default = True.
            enable_refit (bool):
                Build engines with refit option enabled.
            enable_all_tactics (bool):
                Enable all tactic sources during TensorRT engine builds.
            timing_cache (str):
                Path to the timing cache to speed up TensorRT build.
            int8 (bool):
                Whether to quantize to int8 format or not (SDXL, SD15 and SD21 only).
            fp8 (bool):
                Whether to quantize to fp8 format or not (SDXL, SD15 and SD21 only).
            quantization_level (float):
                Controls which layers to quantize. 1: CNN, 2: CNN+FFN, 2.5: CNN+FFN+QKV, 3: CNN+FC
            quantization_percentile (float):
                Control quantization scaling factors (amax) collecting range, where the minimum amax in
                range(n_steps * percentile) will be collected. Recommendation: 1.0
            quantization_alpha (float):
                The alpha parameter for SmoothQuant quantization used for linear layers.
                Recommendation: 0.8 for SDXL
            calibration_size (int):
                The number of steps to use for calibrating the model for quantization.
                Recommendation: 32, 64, 128 for SDXL
            calib_batch_size (int):
                The batch size to use for calibration. Defaults to 2.
            onnx_export_only (bool):
                Whether only export onnx without building the TRT engine.
            download_onnx_models (bool):
                Download pre-exported ONNX models
        """
        
        self._initialize_models(framework_model_dir, int8, fp8, fp4)

        model_configs = self._prepare_model_configs(enable_refit, int8, fp8, fp4)

        # Export models to ONNX
        for model_name, obj in self.models.items():
            if self.torch_fallback[model_name]:
                continue
            self._export_onnx(
                obj,
                model_name,
                model_configs[model_name],
                opt_image_height,
                opt_image_width,
                static_shape,
                onnx_opset,
                quantization_level,
                quantization_percentile,
                quantization_alpha,
                calibration_size,
                calib_batch_size,
                onnx_export_only,
                download_onnx_models,
            )

            # Release temp GPU memory during onnx export to avoid OOM.
            gc.collect()
            torch.cuda.empty_cache()

        if onnx_export_only:
            return

        # Build TensorRT engines
        for model_name, obj in self.models.items():
            if self.torch_fallback[model_name]:
                continue

            model_config = model_configs[model_name]
            engine = engine_module.Engine(model_config["engine_path"])
            if not os.path.exists(model_config['engine_path']):
                self._build_engine(obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache)
            self.engine[model_name] = engine

        # Load and refit TensorRT engines
        for model_name, obj in self.models.items():
            if self.torch_fallback[model_name]:
                continue
            model_config = model_configs[model_name]

            # For non low_vram case, the engines will remain in GPU memory from now on.
            assert self.engine[model_name].engine is None
            if not self.low_vram:
                weight_streaming = getattr(obj, 'weight_streaming', False)
                weight_streaming_budget_percentage = getattr(obj, 'weight_streaming_budget_percentage', None)
                self.engine[model_name].load(weight_streaming, weight_streaming_budget_percentage)

            if model_config['do_engine_refit'] and self.lora_loader:
                # For low_vram, using on-demand load and unload for refit.
                if self.low_vram:
                    assert self.engine[model_name].engine is None
                    self.engine[model_name].load()
                self._refit_engine(obj, model_name, model_config)
                if self.low_vram:
                    self.engine[model_name].unload()

        # Load PyTorch models if torch-inference mode is enabled
        self._load_torch_models()

        # Reclaim GPU memory from torch cache
        torch.cuda.empty_cache()

    def calculate_max_device_memory(self):
        max_device_memory = 0
        for model_name, engine in self.engine.items():
            if self.low_vram:
                engine.load()
            max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
            if self.low_vram:
                engine.unload()
        return max_device_memory

    def get_device_memory_sizes(self):
        device_memory_sizes = {}
        for model_name, engine in self.engine.items():
            engine.load()
            device_memory_sizes[model_name] = engine.engine.device_memory_size_v2
            engine.unload()
        return device_memory_sizes

    def activate_engines(self, shared_device_memory=None):
        if shared_device_memory is None:
            max_device_memory = self.calculate_max_device_memory()
            _, shared_device_memory = cudart.cudaMalloc(max_device_memory)
        self.shared_device_memory = shared_device_memory
        # Load and activate TensorRT engines
        if not self.low_vram:
            for engine in self.engine.values():
                engine.activate(device_memory=self.shared_device_memory)

    def run_engine(self, model_name, feed_dict):
        engine = self.engine[model_name]
        # CUDA graphs should be disabled when low_vram is enabled.
        if self.low_vram:
            assert self.use_cuda_graph == False
        return engine.infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph)

    def teardown(self):
        for e in self.events.values():
            cudart.cudaEventDestroy(e[0])
            cudart.cudaEventDestroy(e[1])

        for engine in self.engine.values():
            engine.deallocate_buffers()
            engine.deactivate()
            engine.unload(verbose=False)
            del engine

        if self.shared_device_memory:
            cudart.cudaFree(self.shared_device_memory)

        for torch_model in self.torch_models.values():
            torch_model.to("cpu")
            del torch_model

        cudart.cudaStreamDestroy(self.stream)
        del self.stream

        gc.collect()
        torch.cuda.empty_cache()

    def initialize_latents(self, batch_size, unet_channels, latent_height, latent_width, latents_dtype=torch.float32):
        latents_shape = (batch_size, unet_channels, latent_height, latent_width)
        latents = torch.randn(latents_shape, device=self.device, dtype=latents_dtype, generator=self.generator)
        # Scale the initial noise by the standard deviation required by the scheduler
        latents = latents * self.scheduler.init_noise_sigma
        return latents

    def save_image(self, images, pipeline, prompt, seed):
        # Save image
        prompt_prefix = ''.join(set([prompt[i].replace(' ','_')[:10] for i in range(len(prompt))]))
        image_name_prefix = '-'.join([pipeline, prompt_prefix, str(seed)])
        image_name_suffix = 'torch' if self.torch_inference else 'trt'
        image_module.save_image(images, self.output_dir, image_name_prefix, image_name_suffix)

    @abstractmethod
    def print_summary(self):
        """Print a summary of the pipeline's configuration."""
        raise NotImplementedError("Please Implement the print_summary method")

    @abstractmethod
    def infer(self):
        """Perform inference using the pipeline."""
        raise NotImplementedError("Please Implement the infer method")

    @abstractmethod
    def run(self):
        """Run the pipeline."""
        raise NotImplementedError("Please Implement the run method")
