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

import argparse
import gc
import html
import os
import pathlib
import random
import re
import time
from typing import Any, Callable, Dict, List, Optional, Union

import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
from diffusers.utils.torch_utils import randn_tensor
from PIL import Image
from tqdm.auto import tqdm

try:
    import ftfy
    FTFY_AVAILABLE = True
except ImportError:
    FTFY_AVAILABLE = False

import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
import demo_diffusion.path as path_module
from demo_diffusion.model import (
    T5Model,
    WanTransformerModel,
    AutoencoderKLWanModel,
    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)

# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L78
def basic_clean(text):
    if FTFY_AVAILABLE:
        text = ftfy.fix_text(text)
    text = html.unescape(html.unescape(text))
    return text.strip()

# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L84
def whitespace_clean(text):
    text = re.sub(r"\s+", " ", text)
    text = text.strip()
    return text

# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L90
def prompt_clean(text):
    text = whitespace_clean(basic_clean(text))
    return text


class WanPipeline(DiffusionPipeline):
    """
    Application showcasing the acceleration of Wan 2.2 T2V pipeline using Nvidia TensorRT.
    """
    
    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
    
    def __init__(
        self,
        dd_path,
        version='wan2.2-t2v-a14b',
        pipeline_type=PIPELINE_TYPE.TXT2VID,
        boundary_ratio: float = 0.875,
        guidance_scale: float = 4.0,
        guidance_scale_2: float = 3.0,
        t5_weight_streaming_budget_percentage=None,
        transformer_weight_streaming_budget_percentage=None,
        **kwargs
    ):
        """
        Initializes the Wan T2V 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 [wan2.2-t2v-a14b]
            pipeline_type (PIPELINE_TYPE):
                Type of current pipeline (TXT2VID)
            boundary_ratio (float, defaults to 0.875):
                Ratio of total timesteps to use as the boundary for switching between transformers
                in two-stage denoising. Transformer handles high-noise stages (timesteps >= boundary)
                and transformer_2 handles low-noise stages (timesteps < boundary). 
                Wan 2.2 T2V always uses two-stage denoising.
            guidance_scale (float):
                Guidance scale for high-noise stage (transformer). Wan default: 4.0
            guidance_scale_2 (float):
                Guidance scale for low-noise stage (transformer_2). Wan default: 3.0
            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 models.
        """
        super().__init__(
            dd_path=dd_path,
            version=version,
            pipeline_type=pipeline_type,
            scheduler="UniPC",
            bf16=True,
            text_encoder_weight_streaming_budget_percentage=t5_weight_streaming_budget_percentage,
            denoiser_weight_streaming_budget_percentage=transformer_weight_streaming_budget_percentage,
            **kwargs
        )

        # Validate boundary_ratio (required for Wan 2.2 two-stage denoising)
        if boundary_ratio is None or not (0.0 < boundary_ratio < 1.0):
            raise ValueError(
                f"`boundary_ratio` must be between 0.0 and 1.0, got {boundary_ratio}"
            )

        self.boundary_ratio = boundary_ratio
        self.guidance_scale = guidance_scale
        self.guidance_scale_2 = guidance_scale_2

        self.vae_scale_factor_temporal = 4
        self.vae_scale_factor_spatial = 8
        
        self.opt_image_height = 720
        self.opt_image_width = 1280
        self.opt_num_frames = 81
        self.max_sequence_length = 512

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

        MAX_BATCH_SIZE = 1  # Wan always uses batch size 1
        DEVICE = "cuda"

        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,
            boundary_ratio=args.boundary_ratio,
            denoising_steps=args.denoising_steps,
            guidance_scale=args.guidance_scale,
            guidance_scale_2=args.guidance_scale_2,
            t5_weight_streaming_budget_percentage=args.t5_ws_percentage if hasattr(args, 't5_ws_percentage') else None,
            transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage if hasattr(args, 'transformer_ws_percentage') else None,
            max_batch_size=MAX_BATCH_SIZE,
            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,
            low_vram=args.low_vram,
            torch_inference=args.torch_inference,
            torch_fallback=args.torch_fallback if hasattr(args, 'torch_fallback') else None,
            weight_streaming=args.ws if hasattr(args, 'ws') else False,
        )
    
    @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
        """
        return ["text_encoder", "transformer", "transformer_2", "vae_decoder"]
    
    def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
        raise NotImplementedError(
            "Pre-exported Wan ONNX models are not available for download. "
            "Export ONNX models locally using the provided export script."
        )
    
    def _initialize_models(self, framework_model_dir, int8=False, fp8=False, fp4=False):
        self.tokenizer = make_tokenizer(
            self.version,
            self.pipeline_type,
            self.hf_token,
            framework_model_dir,
            subfolder='tokenizer',
            tokenizer_type='t5'
        )

        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': 1
        }

        if "text_encoder" in self.stages:
            self.models['text_encoder'] = T5Model(
                **models_args,
                fp16=False,
                bf16=True,
                text_maxlen=self.max_sequence_length,
                weight_streaming=self.weight_streaming,
                weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
                use_attention_mask=True,
            )

        if "transformer" in self.stages:
            self.models['transformer'] = WanTransformerModel(
                **models_args,
                subfolder='transformer',
                fp16=False,
                bf16=True,
                text_maxlen=self.max_sequence_length,
                num_frames=self.opt_num_frames,
                height=self.opt_image_height,
                width=self.opt_image_width,
                weight_streaming=self.weight_streaming,
                weight_streaming_budget_percentage=self.denoiser_weight_streaming_budget_percentage,
            )

        if "transformer_2" in self.stages:
            self.models['transformer_2'] = WanTransformerModel(
                **models_args,
                subfolder='transformer_2',
                fp16=False,
                bf16=True,
                text_maxlen=self.max_sequence_length,
                num_frames=self.opt_num_frames,
                height=self.opt_image_height,
                width=self.opt_image_width,
                weight_streaming=self.weight_streaming,
                weight_streaming_budget_percentage=self.denoiser_weight_streaming_budget_percentage,
            )

        if "vae_decoder" in self.stages:
            self.models['vae_decoder'] = AutoencoderKLWanModel(
                **models_args,
            )

            self.config['vae_decoder_torch_fallback'] = True

    def load_resources(self, image_height, image_width, batch_size, seed):
        """Override to create additional 'denoise' event for combined transformer timing."""
        super().load_resources(image_height, image_width, batch_size, seed)
        # additional event for combined denoising timing (both transformers)
        self.events['denoise'] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]

    def print_summary(self, denoising_steps, walltime_ms, batch_size, num_frames):
        print("|----------------------|--------------|")
        print("| {:^20} | {:^12} |".format("Module", "Latency"))
        print("|----------------------|--------------|")
        
        # calculate transformer timings from combined denoise event
        total_denoise_time = cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1]
        transformer_steps_map = {
            'transformer': self.transformer_steps if (self.transformer_steps > 0 and self.transformer_2_steps > 0) else denoising_steps,
            'transformer_2': self.transformer_2_steps if self.transformer_2_steps > 0 else 0
        }
        
        for stage in self.stages:
            if stage in transformer_steps_map and transformer_steps_map[stage] > 0:
                steps = transformer_steps_map[stage]
                time_ms = total_denoise_time * (steps / denoising_steps)
                stage_label = f"{stage} x {steps}"
            elif stage in transformer_steps_map:
                continue # skip transformer_2 if unused
            else:
                time_ms = cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1]
                stage_label = stage
            
            print("| {:^20} | {:>9.2f} ms |".format(stage_label, time_ms))
        
        print("|----------------------|--------------|")
        print("| {:^20} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
        print("|----------------------|--------------|")
        print("Throughput: {:.2f} videos/min ({} frames)".format(batch_size * 60000.0 / walltime_ms, num_frames))

    def save_video(self, frames, pipeline, prompt, seed):
        if isinstance(prompt, list):
            prompt_prefix = ''.join(set([p.replace(' ','_')[:10] for p in prompt]))
        else:
            prompt_prefix = prompt.replace(' ','_')[:10]
        
        seed_str = str(seed) if seed is not None else 'random'
        precision = 'bf16' if self.bf16 else 'fp16' if self.fp16 else 'fp32'
        video_name_prefix = '-'.join([pipeline, prompt_prefix, precision, seed_str, str(random.randint(1000,9999))])
        video_name_suffix = 'torch' if self.torch_inference else 'trt'
        video_path = video_name_prefix+'-'+video_name_suffix+'.gif'
        full_path = os.path.join(self.output_dir, video_path)
        print(f"Saving video to: {full_path}")
        frames[0].save(full_path, save_all=True, optimize=False, append_images=frames[1:], loop=0)

    # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L198
    def encode_prompt(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        do_classifier_free_guidance: bool = True,
        num_videos_per_prompt: int = 1,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        max_sequence_length: int = 226,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        r"""
        Encodes the prompt into text encoder hidden states.
        
        Implementation modeled from diffusers Wan pipeline, adapted for TensorRT.

        Args:
            prompt (`str` or `List[str]`, *optional*):
                prompt to be encoded
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
                Whether to use classifier free guidance or not.
            num_videos_per_prompt (`int`, *optional*, defaults to 1):
                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
            prompt_embeds (`torch.Tensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.Tensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            max_sequence_length (`int`, defaults to `226`):
                Maximum sequence length for text encoder.
            device: (`torch.device`, *optional*):
                torch device
            dtype: (`torch.dtype`, *optional*):
                torch dtype
        """
        self.profile_start('text_encoder', color='green')
        
        device = device or self._execution_device

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

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

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

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

            negative_prompt_embeds = self._get_t5_prompt_embeds(
                prompt=negative_prompt,
                num_videos_per_prompt=num_videos_per_prompt,
                max_sequence_length=max_sequence_length,
                device=device,
                dtype=dtype,
            )

        self.profile_stop('text_encoder')
        return prompt_embeds, negative_prompt_embeds

    def denoise_latents(
        self,
        latents: torch.Tensor,
        prompt_embeds: torch.Tensor,
        negative_prompt_embeds: Optional[torch.Tensor],
        timesteps: torch.Tensor,
        guidance_scale: float,
        guidance_scale_2: float,
        transformer_dtype: torch.dtype,
        num_warmup_steps: int,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        callback_on_step_end: Optional[Callable] = None,
        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
        warmup: bool = False,
    ) -> torch.Tensor:
        boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
        self.profile_start('denoise', color='blue')
        
        timestep_stages = []
        for i, t in enumerate(timesteps):
            if t >= boundary_timestep:
                timestep_stages.append((i, t, 'transformer', guidance_scale))
            else:
                timestep_stages.append((i, t, 'transformer_2', guidance_scale_2))
        
        stage_groups = []
        if timestep_stages:
            current_group = {
                'transformer': timestep_stages[0][2],
                'guidance_scale': timestep_stages[0][3],
                'timesteps': [(timestep_stages[0][0], timestep_stages[0][1])]
            }
            
            for i, t, transformer_name, gs in timestep_stages[1:]:
                if transformer_name == current_group['transformer']:
                    current_group['timesteps'].append((i, t))
                else:
                    stage_groups.append(current_group)
                    current_group = {
                        'transformer': transformer_name,
                        'guidance_scale': gs,
                        'timesteps': [(i, t)]
                    }
            stage_groups.append(current_group)
        
        self.transformer_steps = sum(len(g['timesteps']) for g in stage_groups if g['transformer'] == 'transformer')
        self.transformer_2_steps = sum(len(g['timesteps']) for g in stage_groups if g['transformer'] == 'transformer_2')
        
        with tqdm(total=len(timesteps)) as progress_bar:
            for stage_group in stage_groups:
                transformer_name = stage_group['transformer']
                current_guidance_scale = stage_group['guidance_scale']
                
                with self.model_memory_manager([transformer_name], low_vram=self.low_vram):
                    for step_index, t in stage_group['timesteps']:
                        latent_model_input = latents.to(transformer_dtype)
                        timestep = t.expand(latent_model_input.shape[0])

                        if self.torch_inference or self.torch_fallback[transformer_name]:
                            current_model = self.torch_models[transformer_name]
                            
                            noise_pred_cond = current_model(
                                hidden_states=latent_model_input,
                                timestep=timestep,
                                encoder_hidden_states=prompt_embeds,
                                attention_kwargs=attention_kwargs,
                                return_dict=False,
                            )[0]

                            if self.do_classifier_free_guidance:
                                noise_pred_uncond = current_model(
                                    hidden_states=latent_model_input,
                                    timestep=timestep,
                                    encoder_hidden_states=negative_prompt_embeds,
                                    attention_kwargs=attention_kwargs,
                                    return_dict=False,
                                )[0]

                                noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_cond - noise_pred_uncond)
                            else:
                                noise_pred = noise_pred_cond
                        else:
                            if self.do_classifier_free_guidance:
                                params_cond = {
                                    "hidden_states": latent_model_input,
                                    "timestep": timestep,
                                    "encoder_hidden_states": prompt_embeds,
                                }
                                
                                # conditional engine call
                                output_cond = self.run_engine(transformer_name, params_cond)['denoised_latents']
                                
                                noise_pred_cond = output_cond.clone()
                                
                                params_uncond = {
                                    "hidden_states": latent_model_input,
                                    "timestep": timestep,
                                    "encoder_hidden_states": negative_prompt_embeds,
                                }
                                
                                # unconditional engine call
                                output_uncond = self.run_engine(transformer_name, params_uncond)['denoised_latents']
                                
                                noise_pred_uncond = output_uncond.clone()
                                
                                # Apply classifier-free guidance
                                noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_cond - noise_pred_uncond)
                            else:
                                # No CFG
                                params = {
                                    "hidden_states": latent_model_input,
                                    "timestep": timestep,
                                    "encoder_hidden_states": prompt_embeds,
                                }
                                noise_pred = self.run_engine(transformer_name, params)['denoised_latents']

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

                        if callback_on_step_end is not None:
                            callback_kwargs = {}
                            for k in callback_on_step_end_tensor_inputs:
                                callback_kwargs[k] = locals()[k]
                            callback_outputs = callback_on_step_end(self, step_index, t, callback_kwargs)

                            latents = callback_outputs.pop("latents", latents)
                            prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                            negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)

                        if step_index == len(timesteps) - 1 or ((step_index + 1) > num_warmup_steps and (step_index + 1) % self.scheduler.order == 0):
                            progress_bar.update()
            
        self.profile_stop('denoise')
        return latents

    # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L324
    def prepare_latents(
        self,
        batch_size: int,
        num_channels_latents: int = 16,
        height: int = 720,
        width: int = 1280,
        num_frames: int = 81,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if latents is not None:
            return latents.to(device=device, dtype=dtype)

        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        shape = (
            batch_size,
            num_channels_latents,
            num_latent_frames,
            int(height) // self.vae_scale_factor_spatial,
            int(width) // self.vae_scale_factor_spatial,
        )
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        if generator is None and hasattr(self, 'generator'):
            generator = self.generator

        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        return latents

    # Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L279
    def check_inputs(
        self,
        prompt,
        negative_prompt,
        height,
        width,
        prompt_embeds=None,
        negative_prompt_embeds=None,
        callback_on_step_end_tensor_inputs=None,
        guidance_scale_2=None,
    ):
        if height % 16 != 0 or width % 16 != 0:
            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")

        if callback_on_step_end_tensor_inputs is not None:
            pass

        if prompt is not None and prompt_embeds is not None:
            raise ValueError(
                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
                " only forward one of the two."
            )
        elif negative_prompt is not None and negative_prompt_embeds is not None:
            raise ValueError(
                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
                " only forward one of the two."
            )
        elif prompt is None and prompt_embeds is None:
            raise ValueError(
                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
            )
        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
        elif negative_prompt is not None and (
            not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
        ):
            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")

    # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L157
    def _get_t5_prompt_embeds(
        self,
        prompt: Union[str, List[str]] = None,
        num_videos_per_prompt: int = 1,
        max_sequence_length: int = 226,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
        encoder: str = "text_encoder",
    ):
        device = device or self._execution_device
        dtype = dtype or (self.torch_models[encoder].dtype if self.torch_fallback.get(encoder) and encoder in self.torch_models else torch.float32)

        prompt = [prompt] if isinstance(prompt, str) else prompt
        prompt = [prompt_clean(u) for u in prompt]
        batch_size = len(prompt)

        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=max_sequence_length,
            truncation=True,
            add_special_tokens=True,
            return_attention_mask=True,
            return_tensors="pt",
        )
        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
        seq_lens = mask.gt(0).sum(dim=1).long()

        if self.torch_inference or self.torch_fallback[encoder]:
            outputs = self.torch_models[encoder](text_input_ids.to(device), mask.to(device))
            prompt_embeds = outputs.last_hidden_state.clone()
        else:
            outputs = self.run_engine(encoder, {
                "input_ids": text_input_ids.to(device),
                "attention_mask": mask.to(device)
            })
            prompt_embeds = outputs['text_embeddings'].clone()
        
        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
        
        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
        prompt_embeds = torch.stack(
            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
        )

        _, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)

        return prompt_embeds

    @property
    def _execution_device(self):
        return self.device

    @property
    def do_classifier_free_guidance(self):
        return self.guidance_scale > 1.0

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

    def decode_latents(self, latents, num_frames):
        self.profile_start('vae_decoder', color='red')
        
        vae_config = self.models['vae_decoder'].config
        z_dim = vae_config.get("z_dim", 16)

        vae_dtype = torch.float32
        latents = latents.to(vae_dtype)
        
        latents_mean = (
            torch.tensor(vae_config.get("latents_mean"))
            .view(1, z_dim, 1, 1, 1)
            .to(latents.device, latents.dtype)
        )
        latents_std = 1.0 / torch.tensor(vae_config.get("latents_std")).view(1, z_dim, 1, 1, 1).to(
            latents.device, latents.dtype
        )
        
        latents = latents / latents_std + latents_mean

        frames = self.torch_models['vae_decoder'].decode(latents, return_dict=False)[0]
        
        self.profile_stop('vae_decoder')
        return frames

    def postprocess(self, video: torch.Tensor, output_type: str = "pil"):
        # Convert [F, C, H, W] -> [F, H, W, C]
        video = video.permute(0, 2, 3, 1)
        # Convert to list of PIL Images
        video = (video + 1.0) / 2.0
        video = torch.clamp(video, 0.0, 1.0)
        video = (video * 255.0).to(torch.uint8).cpu().numpy()
        pil_frames = [Image.fromarray(frame) for frame in video]
        return pil_frames

    def infer(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        height: int = 720,
        width: int = 1280,
        num_frames: int = 81,
        num_inference_steps: int = 40,
        num_videos_per_prompt: int = 1,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        output_type: str = "pil",
        attention_kwargs: Optional[Dict[str, Any]] = None,
        callback_on_step_end: Optional[Union[Callable, Any]] = None,
        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
        max_sequence_length: int = 512,
        warmup: bool = False,
        save_video: bool = True,
    ):
        """
        Run the Wan text-to-video diffusion pipeline.
        """

        self.check_inputs(
            prompt,
            negative_prompt,
            height,
            width,
            prompt_embeds,
            negative_prompt_embeds,
            callback_on_step_end_tensor_inputs,
            self.guidance_scale_2,
        )

        if num_frames % self.vae_scale_factor_temporal != 1:
            print(f"[W] `num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number.")
            num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
        num_frames = max(num_frames, 1)

        device = self._execution_device
        batch_size = 1

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

            with self.model_memory_manager(["text_encoder"], low_vram=self.low_vram):
                prompt_embeds, negative_prompt_embeds = self.encode_prompt(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    do_classifier_free_guidance=self.do_classifier_free_guidance,
                    num_videos_per_prompt=num_videos_per_prompt,
                    prompt_embeds=prompt_embeds,
                    negative_prompt_embeds=negative_prompt_embeds,
                    max_sequence_length=max_sequence_length,
                    device=device,
                )

            transformer_dtype = torch.bfloat16
            prompt_embeds = prompt_embeds.to(transformer_dtype)
            if negative_prompt_embeds is not None:
                negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)

            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
            timesteps = self.scheduler.timesteps

            num_channels_latents = self.models["transformer"].config.get("in_channels", 16)
            latents = self.prepare_latents(
                batch_size * num_videos_per_prompt,
                num_channels_latents,
                height,
                width,
                num_frames,
                torch.float32,
                device,
                generator,
                latents,
            )

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

            latents = self.denoise_latents(
                latents=latents,
                prompt_embeds=prompt_embeds,
                negative_prompt_embeds=negative_prompt_embeds,
                timesteps=timesteps,
                guidance_scale=self.guidance_scale,
                guidance_scale_2=self.guidance_scale_2,
                transformer_dtype=transformer_dtype,
                num_warmup_steps=num_warmup_steps,
                attention_kwargs=attention_kwargs,
                callback_on_step_end=callback_on_step_end,
                callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
                warmup=warmup,
            )

            with self.model_memory_manager(["vae_decoder"], low_vram=self.low_vram):
                video_raw = self.decode_latents(latents, num_frames)
            video = image_module.tensor2vid(video_raw, self, output_type="pil")

            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, num_frames)
            if save_video:
                self.save_video(video[0], self.pipeline_type.name.lower(), prompt, self.seed)

        return video, walltime_ms

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

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