#
# 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 random
import time
import warnings
from typing import Any, List

import numpy as np
import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
from diffusers.video_processor import VideoProcessor
from flux.content_filters import PixtralContentFilter
from tqdm import tqdm

from demo_diffusion import path as path_module
from demo_diffusion.model import (
    AutoencoderKLWanEncoderModel,
    AutoencoderKLWanModel,
    CosmosTransformerModel,
    T5Model,
    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)


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

    def __init__(
        self,
        version="cosmos-predict2-2b",
        pipeline_type=PIPELINE_TYPE.TXT2IMG,
        guidance_scale=6.0,
        max_sequence_length=512,
        t5_weight_streaming_budget_percentage=None,
        transformer_weight_streaming_budget_percentage=None,
        **kwargs,
    ):
        """
        Initializes the Cosmos pipeline.

        Args:
            version (`str`, defaults to `cosmos-1.0-7B`)
                Version of the underlying Cosmos 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.do_classifier_free_guidance = self.guidance_scale > 1

        # WAR ONNX export error: Exporting the operator 'aten::_upsample_nearest_exact2d' to ONNX opset version 19 is not supported
        self.config["vae_torch_fallback"] = True
        self.config["vae_encoder_torch_fallback"] = True

    @classmethod
    def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> CosmosPipeline:
        """Factory method to construct a `CosmosPipeline` 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,
            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,
            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_video2world():
            return ["vae_encoder", "t5", "transformer", "vae"]
        return ["t5", "transformer", "vae"]

    def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
        raise ValueError("ONNX models download is not supported for the Cosmos Pipeline")

    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, 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.fp16 = True if not self.bf16 else False
        self.tf32 = True
        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,
                text_maxlen=self.max_sequence_length,
                use_attention_mask=True,
            )

        if "transformer" in self.stages:
            self.models["transformer"] = CosmosTransformerModel(
                **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 "vae" in self.stages:
            self.models["vae"] = AutoencoderKLWanModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16)

        if "vae_encoder" in self.stages:
            self.models["vae_encoder"] = AutoencoderKLWanEncoderModel(
                **models_args, fp16=False, tf32=self.tf32, bf16=self.bf16
            )

        self.vae_scale_factor_temporal = (
            2 ** sum(self.models["vae"].config["temperal_downsample"])
            if "vae" in self.stages and self.models["vae"] is not None
            else 4
        )
        self.vae_scale_factor_spatial = (
            2 ** len(self.models["vae"].config["temperal_downsample"])
            if "vae" in self.stages and self.models["vae"] is not None
            else 8
        )

        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)

    def encode_video(self, video):
        self.profile_start("vae_encoder", color="red")
        cast_to = (
            torch.float16
            if self.models["vae_encoder"].fp16
            else torch.bfloat16 if self.models["vae_encoder"].bf16 else torch.float32
        )
        video = video.to(dtype=cast_to)
        if self.torch_inference:
            image_latents = self.torch_models["vae_encoder"](video)
        else:
            image_latents = self.run_engine("vae_encoder", {"images": video})["latent"]
        self.profile_stop("vae_encoder")
        return image_latents

    def initialize_latents_text2image(
        self,
        batch_size,
        num_channels_latents,
        num_latent_frames,
        latent_height,
        latent_width,
        latents_dtype=torch.float32,
    ):
        latents_shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
        latents = torch.randn(
            latents_shape,
            device=self.device,
            dtype=latents_dtype,
            generator=self.generator,
        )

        return latents * self.scheduler.config.sigma_max

    def initialize_latents_video2world(
        self,
        video,
        batch_size,
        num_channels_latents,
        num_frames,
        latent_height,
        latent_width,
        latents_dtype=torch.float32,
        do_classifier_free_guidance=False,
    ):
        num_cond_frames = video.size(2)
        if num_cond_frames >= num_frames:
            # Take the last `num_frames` frames for conditioning
            num_cond_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
            video = video[:, :, -num_frames:]
        else:
            num_cond_latent_frames = (num_cond_frames - 1) // self.vae_scale_factor_temporal + 1
            num_padding_frames = num_frames - num_cond_frames
            last_frame = video[:, :, -1:]
            padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1)
            video = torch.cat([video, padding], dim=2)

        # Encode video
        with self.model_memory_manager(["vae_encoder"], low_vram=self.low_vram):
            video_latents = self.encode_video(
                video=video,
            )

        latents_mean = (
            torch.tensor(self.models["vae"].config["latents_mean"])
            .view(1, self.models["vae"].config["z_dim"], 1, 1, 1)
            .to(self.device, latents_dtype)
        )
        latents_std = (
            torch.tensor(self.models["vae"].config["latents_std"])
            .view(1, self.models["vae"].config["z_dim"], 1, 1, 1)
            .to(self.device, latents_dtype)
        )
        init_latents = (video_latents - latents_mean) / latents_std * self.scheduler.config.sigma_data

        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)

        latents = torch.randn(
            shape,
            device=self.device,
            dtype=latents_dtype,
            generator=self.generator,
        )

        latents = latents * self.scheduler.config.sigma_max

        padding_shape = (batch_size, 1, num_latent_frames, latent_height, latent_width)
        ones_padding = latents.new_ones(padding_shape)
        zeros_padding = latents.new_zeros(padding_shape)

        cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1)
        cond_indicator[:, :, :num_cond_latent_frames] = 1.0
        cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding

        uncond_indicator = uncond_mask = None
        if do_classifier_free_guidance:
            uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1)
            uncond_indicator[:, :, :num_cond_latent_frames] = 1.0
            uncond_mask = uncond_indicator * ones_padding + (1 - uncond_indicator) * zeros_padding

        return latents, init_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask

    # 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 _duplicate_text_embeddings(self, batch_size, text_embeddings, num_outputs_per_prompt):
        # duplicate text embeddings for each generation per prompt, using mps friendly method
        _, seq_len, _ = text_embeddings.shape
        text_embeddings = text_embeddings.repeat(1, num_outputs_per_prompt, 1)
        text_embeddings = text_embeddings.view(batch_size * num_outputs_per_prompt, seq_len, -1)
        return text_embeddings

    def _prepare_timesteps(self, num_inference_steps):
        """Prepare timesteps for the scheduler."""
        sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
        sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype)
        accept_sigmas = "sigmas" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())
        if not accept_sigmas:
            raise ValueError(
                f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom"
                f" sigmas schedules. Please check whether you are using the correct scheduler."
            )
        self.scheduler.set_timesteps(sigmas=sigmas, device=self.device)
        timesteps = self.scheduler.timesteps
        num_inference_steps = len(timesteps)
        if self.scheduler.config.get("final_sigmas_type", "zero") == "sigma_min":
            # Replace the last sigma (which is zero) with the minimum sigma value
            self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
        return timesteps, num_inference_steps

    def _encode_text_prompts(self, prompt, negative_prompt, batch_size, num_outputs_per_prompt):
        """Encode text prompts using T5 encoder."""
        with self.model_memory_manager(["t5"], low_vram=self.low_vram):
            text_embeddings = self.encode_prompt(prompt)
            text_embeddings = self._duplicate_text_embeddings(batch_size, text_embeddings, num_outputs_per_prompt)
            negative_text_embeddings = None
            if self.do_classifier_free_guidance:
                negative_text_embeddings = self.encode_prompt(negative_prompt)
                negative_text_embeddings = self._duplicate_text_embeddings(
                    batch_size, negative_text_embeddings, num_outputs_per_prompt
                )
        return text_embeddings, negative_text_embeddings

    def _get_latents_normalization_params(self, device, dtype):
        """Get latents normalization parameters from VAE config."""
        latents_mean = (
            torch.tensor(self.models["vae"].config["latents_mean"])
            .view(1, self.models["vae"].config["z_dim"], 1, 1, 1)
            .to(device, dtype)
        )
        latents_std = (
            torch.tensor(self.models["vae"].config["latents_std"])
            .view(1, self.models["vae"].config["z_dim"], 1, 1, 1)
            .to(device, dtype)
        )
        return latents_mean, latents_std

    def _normalize_and_decode_latents(self, latents, is_video2world=False):
        """Normalize latents and decode using VAE."""
        latents_mean, latents_std = self._get_latents_normalization_params(latents.device, latents.dtype)

        if is_video2world:
            # For video2world: latents * std / sigma_data + mean
            latents = latents * latents_std / self.scheduler.config.sigma_data + latents_mean
        else:
            # For text2image: latents / (1/std) / sigma_data + mean
            latents_std_inv = 1.0 / latents_std
            latents = latents / latents_std_inv / self.scheduler.config.sigma_data + latents_mean

        with self.model_memory_manager(["vae"], low_vram=self.low_vram):
            video = self.decode_latent(latents)

        return video

    def encode_prompt(self, prompt, encoder="t5"):
        self.profile_start(encoder, color="green")

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

            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(self.device)
            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                text_input_ids, untruncated_ids
            ):
                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.max_sequence_length - 1 : -1])
                warnings.warn(
                    "The following part of your input was truncated because `max_sequence_length` is set to "
                    f"{self.max_sequence_length} tokens: {removed_text}"
                )

            if self.torch_inference or self.torch_fallback[encoder]:
                text_encoder_output = self.torch_models[encoder](
                    text_input_ids, attention_mask=attention_mask
                ).last_hidden_state
            else:
                # NOTE: output tensor for the encoder must be cloned because it will be overwritten when called again for prompt2
                text_encoder_output = self.run_engine(
                    encoder, {"input_ids": text_input_ids, "attention_mask": attention_mask}
                )["text_embeddings"]

            lengths = attention_mask.sum(dim=1).cpu()
            for i, length in enumerate(lengths):
                text_encoder_output[i, length:] = 0
            return text_encoder_output

        # Tokenize prompt
        text_encoder_output = tokenize(prompt)

        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.to(torch.float32)
        )

    def denoise_latent(
        self,
        latents,
        timesteps,
        text_embeddings,
        negative_text_embeddings,
        padding_mask,
        denoiser="transformer",
    ):
        do_autocast = self.torch_inference != "" and self.models[denoiser].fp16
        with torch.autocast("cuda", enabled=do_autocast, dtype=torch.float32):
            self.profile_start(denoiser, color="blue")

            for step_index, timestep in tqdm(enumerate(timesteps), total=len(timesteps), desc="Denoising"):
                # Prepare latents
                cast_to = (
                    torch.float16
                    if self.models[denoiser].fp16
                    else torch.bfloat16 if self.models[denoiser].bf16 else torch.float32
                )
                current_sigma = self.scheduler.sigmas[step_index]
                current_t = current_sigma / (current_sigma + 1)
                c_in = 1 - current_t
                c_skip = 1 - current_t
                c_out = -current_t
                timestep_inp = current_t.expand(latents.shape[0]).to(cast_to)  # [B, 1, T, 1, 1]
                latents_input = (latents * c_in).to(cast_to)
                # prepare inputs
                params = {
                    "hidden_states": latents_input,
                    "timestep": timestep_inp,
                    "encoder_hidden_states": text_embeddings,
                    "padding_mask": padding_mask,
                }

                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"].clone()

                noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(cast_to)
                if self.do_classifier_free_guidance:
                    params = {
                        "hidden_states": latents_input,
                        "timestep": timestep_inp,
                        "encoder_hidden_states": negative_text_embeddings,
                        "padding_mask": padding_mask,
                    }

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

                    noise_pred_uncond = (c_skip * latents + c_out * noise_pred_uncond.float()).to(cast_to)
                    noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond)

                noise_pred = (latents - noise_pred) / current_sigma
                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 denoise_latent_video2world(
        self,
        latents,
        timesteps,
        text_embeddings,
        negative_text_embeddings,
        padding_mask,
        fps,
        cond_mask,
        uncond_mask,
        t_conditioning,
        cond_indicator,
        conditioning_latents,
        uncond_indicator,
        unconditioning_latents,
        denoiser="transformer",
    ):
        do_autocast = self.torch_inference != "" and self.models[denoiser].fp16
        with torch.autocast("cuda", enabled=do_autocast, dtype=torch.float32):
            self.profile_start(denoiser, color="blue")

            for step_index, timestep in tqdm(enumerate(timesteps), total=len(timesteps), desc="Denoising"):
                # Prepare latents
                cast_to = (
                    torch.float16
                    if self.models[denoiser].fp16
                    else torch.bfloat16 if self.models[denoiser].bf16 else torch.float32
                )
                current_sigma = self.scheduler.sigmas[step_index]
                current_t = current_sigma / (current_sigma + 1)
                c_in = 1 - current_t
                c_skip = 1 - current_t
                c_out = -current_t
                timestep_inp = current_t.view(1, 1, 1, 1, 1).expand(
                    latents.size(0), -1, latents.size(2), -1, -1
                )  # [B, 1, T, 1, 1]
                latents_input = latents * c_in
                latents_input = (cond_indicator * conditioning_latents + (1 - cond_indicator) * latents_input).to(
                    cast_to
                )
                timestep_inp = (cond_indicator * t_conditioning + (1 - cond_indicator) * timestep_inp).to(cast_to)

                # prepare inputs
                params = {
                    "hidden_states": latents_input,
                    "timestep": timestep_inp,
                    "encoder_hidden_states": text_embeddings,
                    "padding_mask": padding_mask,
                    "fps": fps,
                    "condition_mask": cond_mask,
                }

                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"].clone()

                noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(cast_to)
                noise_pred = cond_indicator * conditioning_latents + (1 - cond_indicator) * noise_pred
                if self.do_classifier_free_guidance:
                    latents_input = latents * c_in
                    latents_input = (
                        uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * latents_input
                    ).to(cast_to)
                    timestep_inp = (uncond_indicator * t_conditioning + (1 - uncond_indicator) * timestep_inp).to(
                        cast_to
                    )
                    params = {
                        "hidden_states": latents_input,
                        "timestep": timestep_inp,
                        "encoder_hidden_states": negative_text_embeddings,
                        "padding_mask": padding_mask,
                        "fps": fps,
                        "condition_mask": uncond_mask,
                    }

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

                    noise_pred_uncond = (c_skip * latents + c_out * noise_pred_uncond.float()).to(cast_to)
                    noise_pred_uncond = (
                        uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * noise_pred_uncond
                    )
                    noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond)

                noise_pred = (latents - noise_pred) / current_sigma
                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]:
            video = self.torch_models[decoder](latents, return_dict=False)[0]
        else:
            video = self.run_engine(decoder, {"latent": latents})["frames"]

        self.profile_stop(decoder)
        return video

    def post_process_video(self, video):
        # Post-process video
        video = self.video_processor.postprocess_video(video, output_type="np")
        video = (video * 255).astype(np.uint8)
        video_batch = []
        for vid in video:
            # vid = self.safety_checker.check_video_safety(vid)
            video_batch.append(vid)
        video = np.stack(video_batch).astype(np.float32) / 255.0 * 2 - 1
        video = torch.from_numpy(video).permute(0, 4, 1, 2, 3)
        video = self.video_processor.postprocess_video(video, output_type="pil")

        if self.pipeline_type.is_video2world():
            return video

        image = [batch[0] for batch in video]
        if isinstance(video, torch.Tensor):
            image = torch.stack(image)
        elif isinstance(video, np.ndarray):
            image = np.stack(image)

        return image

    def _finalize_generation(self, video, walltime_ms, num_inference_steps, batch_size, warmup, save_output):
        """Handle post-processing, saving, and performance reporting."""
        if not warmup:
            self.print_summary(num_inference_steps, walltime_ms, batch_size)
            if not self.return_latents and save_output:
                # post-process video
                processed_output = self.post_process_video(video)

                # save output
                if self.pipeline_type.is_video2world():
                    return (processed_output[0], walltime_ms)
                return (np.array(processed_output), 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 save_video(
        self,
        prompt,
        videos,
        check_integrity=False,
    ):
        for frames in videos:
            if check_integrity:
                self._check_integrity([frames])
            prompt_prefix = "".join(set([prompt[i].replace(" ", "_")[:10] for i in range(len(prompt))]))
            video_name_prefix = "-".join(
                [self.pipeline_type.name.lower(), "fp16", str(self.seed), str(random.randint(1000, 9999))]
            )
            video_name_suffix = "torch" if self.torch_inference else "trt"
            video_path = prompt_prefix + "-" + video_name_prefix + "-" + video_name_suffix + ".gif"
            print(f"Saving video to: {video_path}")
            frames[0].save(
                os.path.join(self.output_dir, video_path),
                save_all=True,
                optimize=False,
                append_images=frames[1:],
                loop=0,
            )

    def print_summary(self, denoising_steps, walltime_ms, batch_size):
        print("|-----------------|--------------|")
        print("| {:^15} | {:^12} |".format("Module", "Latency"))
        print("|-----------------|--------------|")
        for stage in self.stages:
            print(
                "| {:^15} | {:>9.2f} ms |".format(
                    stage + " x " + str(denoising_steps) if stage == "transformer" else stage,
                    cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
                )
            )
        print("|-----------------|--------------|")
        print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
        print("|-----------------|--------------|")
        print("Throughput: {:.2f} image/s".format(batch_size * 1000.0 / walltime_ms))

    def generate_image(
        self,
        prompt,
        negative_prompt,
        image_height,
        image_width,
        num_frames=1,
        num_images_per_prompt=1,
        save_image=True,
        warmup=False,
    ):
        batch_size = len(prompt)

        # Spatial dimensions of latent tensor
        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        latent_height = image_height // self.vae_scale_factor_spatial
        latent_width = image_width // self.vae_scale_factor_spatial

        num_inference_steps = self.denoising_steps

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

            # Prepare timesteps
            timesteps, num_inference_steps = self._prepare_timesteps(num_inference_steps)

            # T5 text encoder
            text_embeddings, negative_text_embeddings = self._encode_text_prompts(
                prompt, negative_prompt, batch_size, num_images_per_prompt
            )

            num_channels_latents = self.models["transformer"].config["in_channels"]
            latents_dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32

            # Initialize latents
            latents = self.initialize_latents_text2image(
                batch_size=batch_size,
                num_channels_latents=num_channels_latents,
                num_latent_frames=num_latent_frames,
                latent_height=latent_height,
                latent_width=latent_width,
                latents_dtype=latents_dtype,
            )
            padding_mask = latents.new_zeros(1, 1, image_height, image_width, dtype=latents_dtype)

            # denoiser
            with self.model_memory_manager(["transformer"], low_vram=self.low_vram):
                latents = self.denoise_latent(
                    latents,
                    timesteps,
                    text_embeddings,
                    negative_text_embeddings,
                    padding_mask,
                )

            # VAE decode latent
            video = self._normalize_and_decode_latents(latents, is_video2world=False)

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

        walltime_ms = (e2e_toc - e2e_tic) * 1000.0
        return self._finalize_generation(
            video,
            walltime_ms,
            num_inference_steps,
            batch_size,
            warmup,
            save_image,
        )

    def generate_video(
        self,
        prompt,
        negative_prompt,
        image_height,
        image_width,
        input_image=None,
        input_video=None,
        num_frames=1,
        fps=16,
        num_videos_per_prompt=1,
        sigma_conditioning=0.0001,
        save_video=True,
        warmup=False,
    ):
        batch_size = len(prompt)

        # Spatial dimensions of latent tensor
        latent_height = image_height // self.vae_scale_factor_spatial
        latent_width = image_width // self.vae_scale_factor_spatial

        num_inference_steps = self.denoising_steps

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

            # Prepare timesteps
            timesteps, num_inference_steps = self._prepare_timesteps(num_inference_steps)

            # T5 text encoder
            text_embeddings, negative_text_embeddings = self._encode_text_prompts(
                prompt, negative_prompt, batch_size, num_videos_per_prompt
            )

            num_channels_latents = self.models["transformer"].config["in_channels"] - 1
            latents_dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32

            # Process input conditioning
            if input_image is not None:
                video = (
                    self.video_processor.preprocess(input_image, image_height, image_width)
                    .unsqueeze(2)
                    .to(device=self.device, dtype=latents_dtype)
                )
            elif input_video is not None:
                video = self.video_processor.preprocess_video(input_video, image_height, image_width).to(
                    device=self.device, dtype=latents_dtype
                )
            else:
                raise ValueError("Video2world pipeline requires either input_image or input_video to be provided")

            # Initialize latents
            latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = (
                self.initialize_latents_video2world(
                    video,
                    batch_size=batch_size,
                    num_channels_latents=num_channels_latents,
                    num_frames=num_frames,
                    latent_height=latent_height,
                    latent_width=latent_width,
                    latents_dtype=latents_dtype,
                    do_classifier_free_guidance=self.do_classifier_free_guidance,
                )
            )
            unconditioning_latents = None

            cond_mask = cond_mask.to(latents_dtype)
            if self.do_classifier_free_guidance:
                uncond_mask = uncond_mask.to(latents_dtype)
                unconditioning_latents = conditioning_latents

            padding_mask = latents.new_zeros(1, 1, image_height, image_width, dtype=latents_dtype)
            sigma_conditioning = torch.tensor(sigma_conditioning, dtype=torch.float32, device=self.device)
            t_conditioning = sigma_conditioning / (sigma_conditioning + 1)

            # denoiser
            with self.model_memory_manager(["transformer"], low_vram=self.low_vram):
                latents = self.denoise_latent_video2world(
                    latents,
                    timesteps,
                    text_embeddings,
                    negative_text_embeddings,
                    padding_mask,
                    fps,
                    cond_mask,
                    uncond_mask,
                    t_conditioning,
                    cond_indicator,
                    conditioning_latents,
                    uncond_indicator,
                    unconditioning_latents,
                )

            # VAE decode latent
            video = self._normalize_and_decode_latents(latents, is_video2world=True)

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

        walltime_ms = (e2e_toc - e2e_tic) * 1000.0
        return self._finalize_generation(
            video,
            walltime_ms,
            num_inference_steps,
            batch_size,
            warmup,
            save_video,
        )

    def infer(
        self,
        prompt,
        negative_prompt,
        image_height,
        image_width,
        input_image=None,
        input_video=None,
        num_frames=1,
        fps=16,
        num_images_per_prompt=1,
        num_videos_per_prompt=1,
        sigma_conditioning=0.0001,
        warmup=False,
        save_output=True,
    ):
        """
        Run the diffusion pipeline.

        Args:
            prompt (str):
                The text prompt to guide image generation.
            negative_prompt (str):
                The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            input_image (image):
                Input image used to initialize the latents.
            input_video (video):
                Input video used to initialize the latents.
            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.
            num_frames (int):
                The number of frames in the generated video.
            fps (int):
                The frames per second of the generated video.
            num_images_per_prompt (int):
                The number of images to generate per prompt.
            num_videos_per_prompt (int):
                The number of videos to generate per prompt.
            sigma_conditioning (`float`, defaults to `0.0001`):
                The sigma value used for scaling conditioning latents. Ideally, it should not be changed or should be
                set to a small value close to zero.
            warmup (bool):
                Indicate if this is a warmup run.
            save_output (bool):
                Save the generated image or video (if applicable)
        """
        if self.pipeline_type.is_txt2img():
            return self.generate_image(
                prompt,
                negative_prompt,
                image_height,
                image_width,
                num_frames,
                num_images_per_prompt,
                save_output,
                warmup,
            )
        elif self.pipeline_type.is_video2world():
            return self.generate_video(
                prompt,
                negative_prompt,
                image_height,
                image_width,
                input_image,
                input_video,
                num_frames,
                fps,
                num_videos_per_prompt,
                sigma_conditioning,
                save_output,
                warmup,
            )
        else:
            raise ValueError(f"Invalid pipeline type: {self.pipeline_type}")

    def run(
        self,
        prompt,
        negative_prompt,
        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, negative_prompt, height, width, warmup=True, **kwargs)

        outputs = []
        for _ in range(batch_count):
            print("[I] Running Cosmos pipeline")
            if self.nvtx_profile:
                cudart.cudaProfilerStart()
            output, _ = self.infer(prompt, negative_prompt, height, width, warmup=False, **kwargs)
            outputs.append(output)
            if self.nvtx_profile:
                cudart.cudaProfilerStop()

        return outputs