groot/control/tensorrt_utils.py

import torch
import os
import subprocess
import tensorrt as trt
import sys
import atexit
import ctypes
import modelopt.torch.quantization as mtq
from typing import Dict, List, Tuple
import shutil

import numpy as np
import torch


FP8_DEFAULT_CONFIG = {
    "quant_cfg": {
        "*weight_quantizer": {"num_bits": (4, 3), "axis": None},
        "*input_quantizer": {"num_bits": (4, 3), "axis": None},
        "*output_quantizer": {"enable": False},
        "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None},
        "*softmax_quantizer": {
            "num_bits": (4, 3),
            "axis": None,
        },
        "default": {"enable": False},
    },
    "algorithm": "max",
}

NVFP4_DEFAULT_CONFIG = {
    "quant_cfg": {
        "*weight_quantizer": {
            "num_bits": (2, 1),
            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
            "axis": None,
            "enable": True,
        },
        "*input_quantizer": {
            "num_bits": (2, 1),
            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
            "axis": None,
            "enable": True,
        },
        "*output_quantizer": {"enable": False},
        "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None},
        "*softmax_quantizer": {
            "num_bits": (4, 3),
            "axis": None,
        },
        "default": {"enable": False},
    },
    "algorithm": "max",
}

    

def wan_quantize(
    policy,
    quantization_config,
    model_type,
    forward_loop,
):
    """Quantize the VLA model using ModelOpt - simplified to use calc_mse_for_single_trajectory."""

    # Configure quantization - disable problematic layers
    if "quant_cfg" in quantization_config:
        quantization_config["quant_cfg"]["*patch_embedding*"] = {"enable": False}
    #    if model_type == "14B" or model_type == "ar_14B":
    #        # Workaround: until we understand the issue https://nvbugspro.nvidia.com/bug/5612316
    #        quantization_config["quant_cfg"]["*.self_attn.o.*"] = {"enable": False}
    #        quantization_config["quant_cfg"]["*.cross_attn.o.*"] = {"enable": False}

    policy.trained_model.action_head.model = mtq.quantize(
        policy.trained_model.action_head.model, quantization_config, forward_loop=forward_loop
    )
    mtq.print_quant_summary(policy.trained_model.action_head.model)

    return


def wan_trt_quantize_and_load_engine(
    policy,
    cfg,
    onnx_path,
    engine_path,
    model_type,
    forward_loop,
):
    if (
        os.path.exists(os.path.dirname(engine_path))
        and cfg.inference_mode == "trt_build"
    ):
        shutil.rmtree(os.path.dirname(engine_path))

    quantization_config = None
    if cfg.quantize_dtype == "fp8":
        quantization_config = FP8_DEFAULT_CONFIG.copy()
    elif cfg.quantize_dtype == "nvfp4":
        quantization_config = NVFP4_DEFAULT_CONFIG.copy()
    else:
        print(f"Quantization type {cfg.quantize_dtype} not supported. Skipping quantization.")

    if quantization_config is not None and cfg.inference_mode == "trt_build":
        #policy.trained_model.action_head.model.to(torch.float16)
        wan_quantize(
            policy,
            quantization_config,
            model_type=model_type,
            forward_loop=forward_loop,
        )

    if  cfg.inference_mode == "trt_build":
        policy.trained_model.action_head.model.to(torch.float16)

        print("Export model:", policy.trained_model.action_head.model)

        test_inputs = create_wan_test_inputs(policy, device="cuda", model_type=model_type)
        min_shape = None
        max_shape = None
        opt_shape = None

        if model_type == "ar_14B":

            policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt
            dynamic_axes = {
                "kv_cache_packed": {3: "kv_cache_len"},
            }
            min_shape = "kv_cache_packed:40x2x1x880x40x128"
            max_shape = "kv_cache_packed:40x2x1x8800x40x128"
            opt_shape = "kv_cache_packed:40x2x1x7920x40x128"
        elif model_type == "ar_14B_droid":
            policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt
            dynamic_axes = {
                "kv_cache_packed": {3: "kv_cache_len"},
            }
            min_shape = "kv_cache_packed:40x2x1x880x40x128"
            max_shape = "kv_cache_packed:40x2x1x8800x40x128"
            opt_shape = "kv_cache_packed:40x2x1x7920x40x128"
        elif model_type == "ar_5B_n6":
            policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt
            dynamic_axes = {
                "kv_cache_packed": {3: "kv_cache_len"},
            }
            min_shape = "kv_cache_packed:30x2x1x220x24x128"
            max_shape = "kv_cache_packed:30x2x1x3080x24x128"
            opt_shape = "kv_cache_packed:30x2x1x2860x24x128"
        else:
            dynamic_axes = None

        if cfg.quantize_dtype == "nvfp4":
            export_to_onnx_fp4(policy.trained_model.action_head.model, test_inputs, onnx_path, dynamic_axes=dynamic_axes)
        else:
            export_to_onnx(
                policy.trained_model.action_head.model,
                test_inputs,
                onnx_path,
                model_type=model_type,
                quantization_mode=cfg.quantize_dtype,
                dynamic_axes=dynamic_axes,
            )

        build_tensorrt_engine(onnx_path, engine_path, min_shape, max_shape, opt_shape)

    trt_wan_model = load_tensorrt_engine(engine_path, model_type=model_type)
    policy.trained_model.action_head.model = trt_wan_model

def export_to_onnx_fp4(model, test_inputs, onnx_save_path, dynamic_axes=None):
    from modelopt.torch._deploy.utils.torch_onnx import OnnxBytes
    from modelopt.torch._deploy.utils.torch_onnx import get_onnx_bytes_and_metadata

    print("exporting to onnx fp4")
    try:
        onnx_bytes, _ = get_onnx_bytes_and_metadata(model=model, dummy_input=test_inputs, dynamic_axes=dynamic_axes)
        onnx_model = OnnxBytes.from_bytes(onnx_bytes)
    except Exception as e:
        print(f"Error exporting model to ONNX: {e}")
        return
    save_dir = os.path.dirname(os.path.abspath(onnx_save_path))
    os.makedirs(save_dir, exist_ok=True)
    for filename, file_bytes in onnx_model.onnx_model.items():
        file_path = os.path.join(save_dir, filename)
        with open(file_path, "wb") as f:
            f.write(file_bytes)
        print(f"exported onnx to {file_path}")


def export_to_onnx(
    pytorch_model,
    test_inputs,
    onnx_path="tensorrt/wan_model.onnx",
    model_type="5B",
    quantization_mode="fp8",
    dynamic_axes=None,
):
    #
    if model_type == "5B":
        return export_to_onnx_5B(pytorch_model, test_inputs, onnx_path, dynamic_axes)
    elif model_type == "14B":
        return export_to_onnx_14B(pytorch_model, test_inputs, onnx_path, dynamic_axes)
    elif model_type == "ar_14B" or model_type == "ar_14B_droid":
        return export_to_onnx_ar_14B(pytorch_model, test_inputs, onnx_path, dynamic_axes)
    else:
        raise ValueError(f"Model type {model_type} not supported")


def export_to_onnx_ar_14B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx", dynamic_axes=None):
    """Export PyTorch model to ONNX"""
    print("Exporting AR 14B model to ONNX...", onnx_path)

    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
    pytorch_model.eval()
    pytorch_model.to(torch.float16)

    input_names = [
        "x",
        "timestep",
        "context",
        "kv_cache_packed",
        "y",
        "clip_feature",
        "action",
        "timestep_action",
        "state",
    ]
    output_names = ["video_noise_pred", "action_noise_pred"]

    try:
        with torch.no_grad():
            torch.onnx.export(
                pytorch_model,
                test_inputs,
                onnx_path,
                export_params=True,
                opset_version=20,
                do_constant_folding=True,
                input_names=input_names,
                output_names=output_names,
                dynamic_axes=dynamic_axes,
            )
        print(f"  ONNX model exported to: {onnx_path}")
        return onnx_path

    except Exception as e:
        import traceback
        print(f"  ERROR: ONNX export failed. Exception type: {type(e)}")
        print("Traceback:")
        traceback.print_exc()
        return None



def export_to_onnx_5B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx"):
    """Export PyTorch model to ONNX"""
    print("Exporting model to ONNX...")

    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
    pytorch_model.eval()
    pytorch_model.to(torch.float16)

    x, action, timestep, context, state, embodiment_id = test_inputs

    # Define input names for better ONNX graph
    input_names = ["x", "action", "timestep", "context", "state", "embodiment_id"]
    output_names = ["video_noise_pred", "action_noise_pred"]

    try:
        with torch.no_grad():
            torch.onnx.export(
                pytorch_model,
                (x, action, timestep, context, state, embodiment_id),
                onnx_path,
                export_params=True,
                opset_version=20,
                do_constant_folding=True,
                input_names=input_names,
                output_names=output_names,
            )
        print(f"  ONNX model exported to: {onnx_path}")
        return onnx_path

    except Exception as e:
        import traceback
        print(f"  ERROR: ONNX export failed. Exception type: {type(e)}")
        print("Traceback:")
        traceback.print_exc()
        return None


def export_to_onnx_14B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx"):
    """Export PyTorch model to ONNX"""
    print("Exporting model to ONNX...")

    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
    pytorch_model.eval()
    pytorch_model.to(torch.float16)

    x, action, timestep, context, state, embodiment_id, clip_feature, y = test_inputs

    # Define input names for better ONNX graph
    input_names = [
        "x",
        "action",
        "timestep",
        "context",
        "state",
        "embodiment_id",
        "clip_feature",
        "y",
    ]
    output_names = ["video_noise_pred", "action_noise_pred"]

    try:
        with torch.no_grad():
            torch.onnx.export(
                pytorch_model,
                (x, action, timestep, context, state, embodiment_id, clip_feature, y),
                onnx_path,
                export_params=True,
                opset_version=20,
                do_constant_folding=True,
                input_names=input_names,
                output_names=output_names,
            )
        print(f"  ONNX model exported to: {onnx_path}")
        return onnx_path

    except Exception as e:
        import traceback
        print(f"  ERROR: ONNX export failed. Exception type: {type(e)}")
        print("Traceback:")
        traceback.print_exc()
        return None


def build_tensorrt_engine(onnx_path, engine_path="tensorrt/wan_model.trt", min_shape=None, max_shape=None, opt_shape=None):
    """Build TensorRT engine from ONNX using trtexec"""
    print("Building TensorRT engine with trtexec...")

    if not os.path.exists(onnx_path):
        print(f"  ERROR: ONNX file not found: {onnx_path}")
        return None

    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(engine_path), exist_ok=True)

    # Build engine using trtexec (much faster than torch_tensorrt)
    trtexec_bin = shutil.which("trtexec") or "/opt/tensorrt/bin/trtexec"
    cmd = [
        trtexec_bin,
        f"--onnx={onnx_path}",
        f"--saveEngine={engine_path}",
        "--fp8",
        "--fp16",
        "--bf16",
        "--separateProfileRun",
        "--profilingVerbosity=detailed",
        "--memPoolSize=workspace:65536",
        "--dumpProfile",
        "--dumpLayerInfo",
        "--useCudaGraph",
        "--verbose",
    ]

    if min_shape is not None:
        cmd.append(f"--minShapes={min_shape}")
    if max_shape is not None:
        cmd.append(f"--maxShapes={max_shape}")
    if opt_shape is not None:
        cmd.append(f"--optShapes={opt_shape}")

    # Create log file for trtexec output
    log_file = engine_path.replace(".trt", "_build.log")

    try:
        print(f"  Running: {' '.join(cmd)}")
        print(f"  Logging output to: {log_file}")

        with open(log_file, "w") as f:
            result = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT, text=True, timeout=600)

        if result.returncode == 0:
            print(f"  TensorRT engine built successfully: {engine_path}")
            print(f"  Build log saved to: {log_file}")
            return engine_path
        else:
            print(f"  ERROR: trtexec failed with return code {result.returncode}")
            print(f"  Check build log for details: {log_file}")
            # Print last few lines of log file for immediate feedback
            try:
                with open(log_file, "r") as f:
                    lines = f.readlines()
                    if lines:
                        print("  Last few lines from build log:")
                        for line in lines[-10:]:  # Show last 10 lines
                            print(f"    {line.rstrip()}")
            except:
                pass
            return None

    except subprocess.TimeoutExpired:
        print("  ERROR: trtexec timed out after 5 minutes")
        print(f"  Partial build log saved to: {log_file}")
        return None
    except Exception as e:
        print(f"  ERROR: Failed to run trtexec: {e}")
        return None


def torch_type(trt_type):
    mapping = {
        trt.float32: torch.float32,  # Added missing FLOAT mapping
        trt.float16: torch.float16,
        trt.bfloat16: torch.bfloat16,
        trt.int8: torch.int8,
        trt.int32: torch.int32,
        trt.bool: torch.bool,
        trt.uint8: torch.uint8,
        trt.int64: torch.int64,
    }
    if trt_type in mapping:
        return mapping[trt_type]

    raise TypeError(
        f"Could not resolve TensorRT datatype to an equivalent torch datatype. {trt_type}"
    )


class Engine(object):
    def __init__(self, file, plugins=[]):
        super().__init__()

        self.logger = trt.Logger(trt.Logger.ERROR)
        trt.init_libnvinfer_plugins(self.logger, "")

        self.plugins = [ctypes.CDLL(plugin, ctypes.RTLD_GLOBAL) for plugin in plugins]
        self.file = file
        self.load(file)

        def destroy(self):
            del self.execution_context
            del self.handle

        atexit.register(destroy, self)
        self.print()

    def print(self):

        print("============= TRT Engine Detail =============")
        print(f"Engine file: {self.file}")
        print(f"Inputs: {len(self.in_meta)}")
        for ib, item in enumerate(self.in_meta):
            tensor_name, shape, dtype = item[:3]
            print(f"   {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]")

        print(f"Outputs: {len(self.out_meta)}")
        for ib, item in enumerate(self.out_meta):
            tensor_name, shape, dtype = item[:3]
            print(f"   {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]")
        print("=============================================")

    def load(self, file):
        runtime = trt.Runtime(self.logger)

        with open(file, "rb") as f:
            self.handle = runtime.deserialize_cuda_engine(f.read())
            assert (
                self.handle is not None
            ), f"Failed to deserialize the cuda engine from file: {file}"

        self.execution_context = self.handle.create_execution_context()
        self.meta, self.in_meta, self.out_meta = [], [], []
        for tensor_name in self.handle:
            shape = self.handle.get_tensor_shape(tensor_name)
            print(f"Tensor name: {tensor_name}, shape: {shape}")
            dtype = torch_type(self.handle.get_tensor_dtype(tensor_name))
            if self.handle.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
                self.in_meta.append([tensor_name, shape, dtype])
            else:
                self.out_meta.append([tensor_name, shape, dtype])

    def __call__(self, *args, **inputs):
        return self.forward(*args, **inputs)

    def set_runtime_tensor_shape(self, name, shape):
        self.execution_context.set_input_shape(name, shape)

    def forward(self, *args, **kwargs):
        return_list = kwargs.pop("return_list", False)
        reference_tensors = []
        stream = torch.cuda.current_stream()
        for iarg, x in enumerate(args):
            name, shape, dtype = self.in_meta[iarg]
            runtime_shape = self.execution_context.get_tensor_shape(name)
            assert isinstance(x, torch.Tensor), f"Unsupported tensor type: {type(x)}"
            assert runtime_shape == x.shape, f"Invalid input shape: {runtime_shape} != {x.shape}"
            assert (
                dtype == x.dtype
            ), f"Invalid tensor dtype, excepted dtype is {dtype}, but got {x.dtype}"
            assert x.is_cuda, f"Invalid tensor device, excepted device is cuda, but got {x.device}"
            x = x.cuda().contiguous()
            self.execution_context.set_tensor_address(name, x.data_ptr())
            reference_tensors.append(x)

        for name, shape, dtype in self.in_meta:
            if name not in kwargs:
                continue

            runtime_shape = self.execution_context.get_tensor_shape(name)
            x = kwargs[name]
            assert isinstance(x, torch.Tensor), f"Unsupported tensor[{name}] type: {type(x)}"
            assert (
                runtime_shape == x.shape
            ), f"Invalid input[{name}] shape: {x.shape}, but the expected shape is: {runtime_shape}"
            assert (
                dtype == x.dtype
            ), f"Invalid tensor[{name}] dtype, expected dtype is {dtype}, but got {x.dtype}"
            assert (
                x.is_cuda
            ), f"Invalid tensor[{name}] device, expected device is cuda, but got {x.device}"
            x = x.cuda().contiguous()
            self.execution_context.set_tensor_address(name, x.data_ptr())
            reference_tensors.append(x)

        for item in self.out_meta:
            name = item[0]
            runtime_shape = self.execution_context.get_tensor_shape(name)
            output_tensor = torch.zeros(
                *runtime_shape, dtype=item[2], device=reference_tensors[0].device
            )
            self.execution_context.set_tensor_address(name, output_tensor.data_ptr())
            reference_tensors.append(output_tensor)

        self.execution_context.execute_async_v3(stream.cuda_stream)
        stream.synchronize()
        assert len(reference_tensors) == len(self.in_meta) + len(
            self.out_meta
        ), f"Invalid input tensors. The expected I/O tensors are {len(self.in_meta) + len(self.out_meta)}, but got {len(reference_tensors)}"

        if return_list:
            return [
                reference_tensors[len(self.in_meta) + i] for i, item in enumerate(self.out_meta)
            ]
        else:
            return {
                item[0]: reference_tensors[len(self.in_meta) + i]
                for i, item in enumerate(self.out_meta)
            }


class WanTrtModel5B(torch.nn.Module):
    def __init__(self, eng_path: str):
        super().__init__()
        self.engine = Engine(eng_path)

    def forward(
        self,
        x: torch.Tensor,
        action: torch.Tensor,
        timestep: torch.Tensor,
        context: torch.Tensor,
        state: torch.Tensor,
        embodiment_id: torch.Tensor,
    ):

        self.engine.set_runtime_tensor_shape("x", x.shape)
        self.engine.set_runtime_tensor_shape("action", action.shape)
        self.engine.set_runtime_tensor_shape("context", context.shape)
        self.engine.set_runtime_tensor_shape("state", state.shape)

        output = self.engine(
            x=x.to(torch.float16),
            action=action.to(torch.float16),
            timestep=timestep.to(torch.float16),
            context=context.to(torch.float16),
            state=state.to(torch.float16),
            embodiment_id=embodiment_id.to(torch.int32),
        )
        if "out.0" in output: # for nvfp4 model export through modelopt
            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
        else:
            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()


class WanTrtModel14B(torch.nn.Module):
    def __init__(self, eng_path: str):
        super().__init__()
        self.engine = Engine(eng_path)

    def forward(
        self,
        x: torch.Tensor,
        action: torch.Tensor,
        timestep: torch.Tensor,
        context: torch.Tensor,
        state: torch.Tensor,
        embodiment_id: torch.Tensor,
        clip_feature: torch.Tensor,
        y: torch.Tensor,
    ):

        self.engine.set_runtime_tensor_shape("x", x.shape)
        self.engine.set_runtime_tensor_shape("action", action.shape)
        self.engine.set_runtime_tensor_shape("context", context.shape)
        self.engine.set_runtime_tensor_shape("state", state.shape)
        self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape)
        self.engine.set_runtime_tensor_shape("y", y.shape)

        output = self.engine(
            x=x.to(torch.float16),
            action=action.to(torch.float16),
            timestep=timestep.to(torch.float16),
            context=context.to(torch.float16),
            state=state.to(torch.float16),
            embodiment_id=embodiment_id.to(torch.int32),
            clip_feature=clip_feature.to(torch.float16),
            y=y.to(torch.float16),
        )
        if "out.0" in output: # for nvfp4 model export through modelopt
            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
        else:
            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()


class WanTrtModelAr5B(torch.nn.Module):
    """TRT wrapper for ar_5B_n6 model type - uses kv_cache but no clip_feature."""
    def __init__(self, eng_path: str):
        super().__init__()
        self.engine = Engine(eng_path)

    def forward(
        self,
        x,
        timestep,
        context,
        kv_cache: list[torch.Tensor],
        y=None,
        action=None,
        timestep_action=None,
        state=None,
    ):

        kv_cache_packed = torch.stack(kv_cache, dim=0)

        self.engine.set_runtime_tensor_shape("x", x.shape)
        self.engine.set_runtime_tensor_shape("timestep", timestep.shape)
        self.engine.set_runtime_tensor_shape("context", context.shape)
        self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape)
        # self.engine.set_runtime_tensor_shape("y", y.shape)
        self.engine.set_runtime_tensor_shape("action", action.shape)
        self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape)
        self.engine.set_runtime_tensor_shape("state", state.shape)


        output = self.engine(
            x.to(torch.float16),
            timestep.to(torch.float16),
            context.to(torch.float16),
            kv_cache_packed.to(torch.float16),
            # y.to(torch.float16),
            action.to(torch.float16),
            timestep_action.to(torch.float16),
            state.to(torch.float16),
        )

        if "out.0" in output: # for nvfp4 model export through modelopt
            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
        else:
            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()


class WanTrtModelAr14B(torch.nn.Module):
    def __init__(self, eng_path: str):
        super().__init__()
        self.engine = Engine(eng_path)

    def forward(
        self,
        x,
        timestep,
        context,
        kv_cache: list[torch.Tensor],
        y=None,
        clip_feature=None,
        action=None,
        timestep_action=None,
        state=None,
    ):

        kv_cache_packed = torch.stack(kv_cache, dim=0)

        self.engine.set_runtime_tensor_shape("x", x.shape)
        self.engine.set_runtime_tensor_shape("timestep", timestep.shape)
        self.engine.set_runtime_tensor_shape("context", context.shape)
        self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape)
        self.engine.set_runtime_tensor_shape("y", y.shape)
        self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape)
        self.engine.set_runtime_tensor_shape("action", action.shape)
        self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape)
        self.engine.set_runtime_tensor_shape("state", state.shape)


        output = self.engine(
            x.to(torch.float16),
            timestep.to(torch.float16),
            context.to(torch.float16),
            kv_cache_packed.to(torch.float16),
            y.to(torch.float16),
            clip_feature.to(torch.float16),
            action.to(torch.float16),
            timestep_action.to(torch.float16),
            state.to(torch.float16),
        )

        if "out.0" in output: # for nvfp4 model export through modelopt
            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
        else:
            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()

def load_tensorrt_engine(engine_path="tensorrt/wan_model.trt", model_type="5B"):
    """Load TensorRT engine"""
    if model_type == "5B":
        trt_inference = WanTrtModel5B(engine_path)
    elif model_type == "ar_5B_n6" or model_type == "ar_5B":
        trt_inference = WanTrtModelAr5B(engine_path)
    elif model_type == "14B":
        trt_inference = WanTrtModel14B(engine_path)
    elif model_type == "ar_14B" or model_type == "ar_14B_droid":
        trt_inference = WanTrtModelAr14B(engine_path)
    else:
        raise ValueError(f"Model type {model_type} not supported")
    return trt_inference


def create_wan_test_inputs(policy, device="cuda", model_type="5B"):
    # Get dtype from model parameters
    dtype = torch.float16

    # Use hardcoded dimensions from the original working version of the script
    if model_type == "5B":
        x = torch.randn(1, 48, 13, 22, 40, dtype=dtype, device=device)
        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
        timestep = torch.randn(1, dtype=dtype, device=device)
        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
        timestep_action = torch.randn(1, 48, dtype=dtype, device=device)
        seq_len = torch.tensor(440, dtype=torch.int32, device=device)
        return x, action, timestep, context, state, embodiment_id, timestep_action, seq_len
    elif model_type == "ar_5B_n6":
        # ar_5B_n6 uses _forward_inference_trt which requires kv_cache_packed
        # Shape from dynamic_axes: kv_cache_packed:30x2x1x220x24x128
        # Note: 5B model doesn't use clip_feature (unlike 14B), but still needs y
        x = torch.randn(1, 48, 2, 22, 40, dtype=dtype, device=device)
        timestep = torch.randn(1, 2, dtype=dtype, device=device)
        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
        # y = torch.randn(1, 52, 2, 22, 40, dtype=dtype, device=device)  # y is required by _forward_inference_trt
        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
        timestep_action = torch.randn(1, 48, dtype=dtype, device=device)
        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
    
        num_heads = 24
        head_dim = 128
        num_layers = 30
        B = 1
    
        kv_cache = []
        for _ in range(num_layers):
            kv_cache.append(
                torch.zeros([2, B, 13*220, num_heads, head_dim], dtype=dtype, device=device)
            )
    
        kv_cache_packed = torch.stack(kv_cache, dim=0)
        # Return order matches _forward_inference_trt signature: x, timestep, context, kv_cache_packed, y, action, timestep_action, state
        return (x, timestep, context, kv_cache_packed, action, timestep_action, state)
    elif model_type == "14B":
        x = torch.randn(1, 16, 13, 44, 80, dtype=dtype, device=device)
        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
        timestep = torch.randn(1, dtype=dtype, device=device)
        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
        clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device)
        y = torch.randn(1, 20, 13, 44, 80, dtype=dtype, device=device)
        return x, action, timestep, context, state, embodiment_id, clip_feature, y
    elif model_type == "ar_14B":
        clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device)
        y = torch.randn(1, 20, 2, 44, 80, dtype=dtype, device=device)
        timestep_action = torch.randn(1, 48, dtype=dtype, device=device)
        x = torch.randn(1, 16, 2, 44, 80, dtype=dtype, device=device)
        timestep = torch.randn(1, 2, dtype=dtype, device=device)
        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
        seq_len = torch.tensor(1760, dtype=torch.int32, device=device)
        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
    
        num_heads = 40 
        head_dim = 5120 // num_heads
        num_layers = 40 
        B = 1
    
        kv_cache = []
        for _ in range(num_layers):
            kv_cache.append(
                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
            )
    
        crossattn_k_cache = []
        for _ in range(num_layers):
            crossattn_k_cache.append(
                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
            )
        kv_cache_packed = torch.stack(kv_cache, dim=0)
        crossattn_packed = torch.stack(crossattn_k_cache, dim=0)
        return (x, timestep, context, kv_cache_packed, y, clip_feature, action, timestep_action, state)
    elif model_type == "ar_14B_droid":
        clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device)
        y = torch.randn(1, 20, 2, 44, 80, dtype=dtype, device=device)
        timestep_action = torch.randn(1, 24, dtype=dtype, device=device)
        x = torch.randn(1, 16, 2, 44, 80, dtype=dtype, device=device)
        timestep = torch.randn(1, 2, dtype=dtype, device=device)
        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
        seq_len = torch.tensor(1760, dtype=torch.int32, device=device)
        action = torch.randn(1, 24, 32, dtype=dtype, device=device)
        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
    
        num_heads = 40 
        head_dim = 5120 // num_heads
        num_layers = 40 
        B = 1
    
        kv_cache = []
        for _ in range(num_layers):
            kv_cache.append(
                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
            )
    
        crossattn_k_cache = []
        for _ in range(num_layers):
            crossattn_k_cache.append(
                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
            )
        kv_cache_packed = torch.stack(kv_cache, dim=0)
        crossattn_packed = torch.stack(crossattn_k_cache, dim=0)
        return (x, timestep, context, kv_cache_packed, y, clip_feature, action, timestep_action, state)