nemo/export/multimodal/build.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# 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 os
import shutil
import tarfile
import tempfile
from pathlib import Path
from time import time
from typing import List

import tensorrt as trt
import torch
import yaml
from omegaconf import OmegaConf
from PIL import Image
from tensorrt_llm._common import check_max_num_tokens
from tensorrt_llm.builder import BuildConfig, Builder
from tensorrt_llm.commands.build import build as build_trtllm
from tensorrt_llm.mapping import Mapping
from tensorrt_llm.models import MLLaMAForCausalLM
from tensorrt_llm.plugin import PluginConfig
from transformers import AutoModel, AutoProcessor, MllamaForConditionalGeneration

from nemo.collections.multimodal.speech_llm.modules.perception_modules import AudioPerceptionModule
from nemo.core.classes.common import typecheck
from nemo.export.tensorrt_llm import TensorRTLLM
from nemo.export.trt_llm.nemo_ckpt_loader.nemo_file import load_nemo_model

from .converter import convert_mllama_nemo_to_hf

logger = trt.Logger(trt.Logger.INFO)


def build_trtllm_engine(
    model_dir: str,
    visual_checkpoint_path: str,
    llm_checkpoint_path: str = None,
    model_type: str = "neva",
    llm_model_type: str = "llama",
    tensor_parallelism_size: int = 1,
    max_input_len: int = 256,
    max_output_len: int = 256,
    max_batch_size: int = 1,
    max_multimodal_len: int = 1024,
    dtype: str = "bfloat16",
    use_lora_plugin: str = None,
    lora_target_modules: List[str] = None,
    max_lora_rank: int = 64,
    lora_ckpt_list: List[str] = None,
):
    """Build TRTLLM engine by nemo export"""
    trt_llm_exporter = TensorRTLLM(model_dir=model_dir, lora_ckpt_list=lora_ckpt_list, load_model=False)
    trt_llm_exporter.export(
        nemo_checkpoint_path=visual_checkpoint_path if llm_checkpoint_path is None else llm_checkpoint_path,
        model_type=llm_model_type,
        tensor_parallelism_size=tensor_parallelism_size,
        max_input_len=max_input_len,
        max_output_len=max_output_len,
        max_seq_len=max_input_len + max_output_len,
        max_batch_size=max_batch_size,
        max_prompt_embedding_table_size=max_multimodal_len,
        dtype=dtype,
        load_model=False,
        use_lora_plugin=use_lora_plugin,
        lora_target_modules=lora_target_modules,
        max_lora_rank=max_lora_rank,
        use_mcore_path=False,
    )


def build_mllama_trtllm_engine(
    model_dir: str,
    hf_model_path: str,
    tensor_parallelism_size: int = 1,
    max_input_len: int = 256,
    max_output_len: int = 256,
    max_batch_size: int = 1,
    max_multimodal_len: int = 1024,
    dtype: str = "bfloat16",
    use_lora_plugin: str = None,
    lora_target_modules: List[str] = None,
    max_lora_rank: int = 64,
    lora_ckpt_list: List[str] = None,
):
    """Build mllama TRTLLM engine from HF"""
    if max_batch_size < 4:
        print(
            "TensorRT LLM may hit a runtime issue with batch size is smaller than 4 on some models." " Force set to 4"
        )
        max_batch_size = 4

    plugin_config = PluginConfig()
    plugin_config.gpt_attention_plugin = "auto"
    plugin_config.gemm_plugin = "auto"
    plugin_config.enable_paged_kv_cache(tokens_per_block=128)
    plugin_config.remove_input_padding = True
    plugin_config.use_paged_context_fmha = True

    max_seq_len = max_input_len + max_output_len
    max_num_tokens, opt_num_tokens = check_max_num_tokens(
        max_num_tokens=None,
        opt_num_tokens=None,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
        max_input_len=max_input_len,
        max_beam_width=1,
        remove_input_padding=True,
        enable_context_fmha=plugin_config.context_fmha,
        tokens_per_block=128,
        multiple_profiles=False,
    )

    build_dict = {
        'max_input_len': max_input_len,
        'max_output_len': max_output_len,
        'max_encoder_input_len': max_multimodal_len,
        'max_batch_size': max_batch_size,
        'max_beam_width': 1,
        'max_seq_len': max_seq_len,
        'max_num_tokens': max_num_tokens,
        'opt_num_tokens': opt_num_tokens,
        'strongly_typed': True,
        'builder_opt': None,
    }
    build_config = BuildConfig.from_dict(build_dict, plugin_config=plugin_config)

    for rank in range(tensor_parallelism_size):
        mapping = Mapping(world_size=tensor_parallelism_size, rank=rank, tp_size=tensor_parallelism_size)
        model = MLLaMAForCausalLM.from_hugging_face(
            hf_model_path,
            dtype,
            mapping=mapping,
        )

        engine = build_trtllm(model, build_config)
        engine.save(model_dir)


def export_visual_wrapper_onnx(
    visual_wrapper, input, output_dir, input_names=['input'], dynamic_axes={'input': {0: 'batch'}}
):
    """Export visual wrapper to ONNX"""
    logger.log(trt.Logger.INFO, "Exporting onnx")
    os.makedirs(f'{output_dir}/onnx', exist_ok=True)
    torch.onnx.export(
        visual_wrapper,
        input,
        f'{output_dir}/onnx/visual_encoder.onnx',
        opset_version=17,
        input_names=input_names,
        output_names=['output'],
        dynamic_axes=dynamic_axes,
    )


def export_perception_wrapper_onnx(
    perception_wrapper,
    input,
    output_dir,
    input_names=['processed_signal', 'processed_signal_length'],
    output_names=['encoded', 'encoded_length'],
    dynamic_axes={
        'processed_signal': {0: 'batch', 2: 'time'},
        'processed_signal_length': {0: 'batch'},
        'encoded': {0: 'batch', 1: 'time'},
        'encoded_length': {0: 'batch'},
    },
):
    """Export perception wrapper to ONNX"""
    logger.log(trt.Logger.INFO, "Exporting onnx")
    os.makedirs(f'{output_dir}/onnx', exist_ok=True)
    torch.onnx.export(
        perception_wrapper,
        input,
        f'{output_dir}/onnx/perception_encoder.onnx',
        opset_version=17,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
    )


def build_trt_engine(
    model_type,
    input_sizes,
    output_dir,
    vision_max_batch_size,
    dtype=torch.bfloat16,
    image_size=None,
    num_frames=None,
    nemo_config=None,
    part_name='visual_encoder',
):
    """Build TRT engine from onnx"""
    onnx_file = '%s/onnx/%s.onnx' % (output_dir, part_name)
    engine_file = '%s/%s.engine' % (output_dir, part_name)
    config_file = '%s/%s' % (output_dir, "config.json")
    nemo_config_file = '%s/%s' % (output_dir, "nemo_config.yaml")

    with open(nemo_config_file, 'w') as f:
        yaml.dump(nemo_config, f)

    logger.log(trt.Logger.INFO, "Building TRT engine for %s" % part_name)

    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    profile = builder.create_optimization_profile()

    config_args = {"precision": str(dtype).split('.')[-1], "model_type": model_type}
    if image_size is not None:
        config_args["image_size"] = image_size
    if num_frames is not None:
        config_args["num_frames"] = num_frames

    config_wrapper = Builder().create_builder_config(**config_args)
    config = config_wrapper.trt_builder_config

    parser = trt.OnnxParser(network, logger)

    with open(onnx_file, 'rb') as model:
        if not parser.parse(model.read(), os.path.abspath(onnx_file)):
            logger.log(trt.Logger.ERROR, "Failed parsing %s" % onnx_file)
            for error in range(parser.num_errors):
                logger.log(trt.Logger.ERROR, parser.get_error(error))
        logger.log(trt.Logger.INFO, "Succeeded parsing %s" % onnx_file)

    # Delete onnx files since we don't need them now
    shutil.rmtree(f'{output_dir}/onnx')

    nBS = -1
    nMinBS = 1
    nOptBS = max(nMinBS, int(vision_max_batch_size / 2))
    nMaxBS = vision_max_batch_size

    inputT = network.get_input(0)

    # input sizes can be a list of ints (e.g., [3, H, W]) when inputs are images,
    # or a list of three int lists (e.g., [[1, 1, 2700], [1, 500, 2700], [1, 4096, 2700]]).
    # or a list of three list of lists
    # (e.g., [{input1: min_shape, input2: min_shape, }, \
    #     {input1: opt_shape, input2: opt_shape}, \
    # {input1: max_shape, input2: max_shape}] )
    assert isinstance(input_sizes, list), "input_sizes must be a list"
    if isinstance(input_sizes[0], int):
        logger.log(trt.Logger.INFO, f"Processed input sizes {input_sizes}")
        inputT.shape = [nBS, *input_sizes]
        min_size = opt_size = max_size = input_sizes
    elif len(input_sizes) == 3 and isinstance(input_sizes[0], list):
        min_size, opt_size, max_size = input_sizes
        logger.log(trt.Logger.INFO, f"Processed min/opt/max input sizes {min_size}/{opt_size}/{max_size}")
    elif len(input_sizes) == 3 and isinstance(input_sizes[0], dict):
        logger.log(trt.Logger.INFO, f"Processed min/opt/max input sizes {input_sizes}")
    else:
        raise ValueError(f"invalid input sizes: {input_sizes}")

    if isinstance(input_sizes[0], dict):
        for i in range(network.num_inputs):
            inputT = network.get_input(i)
            input_name = inputT.name
            min_size = input_sizes[0][input_name]
            opt_size = input_sizes[1][input_name]
            max_size = input_sizes[2][input_name]
            logger.log(trt.Logger.INFO, f"{input_name} min/opt/max input sizes {min_size}/{opt_size}/{max_size}")
            profile.set_shape(input_name, min_size, opt_size, max_size)
    else:
        profile.set_shape(inputT.name, [nMinBS, *min_size], [nOptBS, *opt_size], [nMaxBS, *max_size])

    config.add_optimization_profile(profile)

    t0 = time()
    engine_string = builder.build_serialized_network(network, config)
    t1 = time()
    if engine_string is None:
        raise RuntimeError("Failed building %s" % (engine_file))
    else:
        logger.log(trt.Logger.INFO, "Succeeded building %s in %d s" % (engine_file, t1 - t0))
        with open(engine_file, 'wb') as f:
            f.write(engine_string)

    Builder.save_config(config_wrapper, config_file)


def build_neva_engine(
    model_type: str,
    model_dir: str,
    visual_checkpoint_path: str,
    vision_max_batch_size: int = 1,
):
    """Build neva visual engine"""
    device = torch.device("cuda") if torch.cuda.is_available() else "cpu"

    if os.path.isdir(visual_checkpoint_path):
        # load untar checkpoint
        config_path = os.path.join(visual_checkpoint_path, 'model_config.yaml')
        with open(config_path, 'r') as f:
            nemo_config = yaml.safe_load(f)
        try:
            weights_path = os.path.join(visual_checkpoint_path, 'model_weights.ckpt')
            mp0_weights = torch.load(weights_path, map_location=device)
        except FileNotFoundError:
            weights_path = os.path.join(visual_checkpoint_path, 'mp_rank_00/model_weights.ckpt')
            mp0_weights = torch.load(weights_path, map_location=device)
    else:
        # extract NeMo checkpoint
        with tempfile.TemporaryDirectory() as temp:
            temp_path = Path(temp)
            mp0_weights, nemo_config, _ = load_nemo_model(visual_checkpoint_path, temp_path)

    vision_config = nemo_config["mm_cfg"]["vision_encoder"]

    class DownSampleBlock(torch.nn.Module):
        # pylint: disable=C0115,C0116
        def forward(self, x):
            vit_embeds = x
            h = w = int(vit_embeds.shape[1] ** 0.5)
            vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
            vit_embeds = self.flat_square(vit_embeds)
            vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
            return vit_embeds

        def flat_square(self, x):
            n, w, h, c = x.size()
            if w % 2 == 1:
                x = torch.cat([x, torch.zeros((n, 1, h, c), dtype=x.dtype).to(x.device)], dim=1).contiguous()
                n, w, h, c = x.size()
            if h % 2 == 1:
                x = torch.cat([x, torch.zeros((n, w, 1, c), dtype=x.dtype).to(x.device)], dim=2).contiguous()
                n, w, h, c = x.size()
            x = x.view(n, w, int(h / 2), int(c * 2))
            x = x.permute(0, 2, 1, 3).contiguous()
            x = x.view(n, int(h / 2), int(w / 2), int(c * 4))
            return x

    class VisionEncoderWrapper(torch.nn.Module):
        # pylint: disable=C0115,C0116
        def __init__(self, encoder, connector):
            super().__init__()
            self.encoder = encoder
            self.connector = connector

        def forward(self, images):
            vision_x = self.encoder(pixel_values=images, output_hidden_states=True)
            vision_x = vision_x.hidden_states[-2]
            vision_x = self.connector(vision_x)
            return vision_x

    encoder = AutoModel.from_pretrained(
        vision_config["from_pretrained"],
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation='eager',
    )
    vision_encoder = encoder.vision_model
    hf_config = encoder.config
    dtype = hf_config.torch_dtype

    # connector
    if nemo_config["mm_cfg"]["mm_mlp_adapter_type"] == "mlp2x_gelu":
        vision_connector = torch.nn.Sequential(
            torch.nn.Linear(vision_config["hidden_size"], nemo_config["hidden_size"], bias=True),
            torch.nn.GELU(),
            torch.nn.Linear(nemo_config["hidden_size"], nemo_config["hidden_size"], bias=True),
        ).to(dtype=dtype)

        key_prefix = "model.embedding.word_embeddings.adapter_layer.mm_projector_adapter.mm_projector"
        for layer in range(0, 3, 2):
            vision_connector[layer].load_state_dict(
                {
                    'weight': mp0_weights[f"{key_prefix}.{layer}.weight"].to(dtype),
                    'bias': mp0_weights[f"{key_prefix}.{layer}.bias"].to(dtype),
                }
            )
    elif nemo_config["mm_cfg"]["mm_mlp_adapter_type"] == "linear":
        vision_connector = torch.nn.Linear(vision_config["hidden_size"], nemo_config["hidden_size"], bias=True)
        key_prefix = "model.embedding.word_embeddings.adapter_layer.mm_projector_adapter.mm_projector"
        vision_connector.load_state_dict(
            {
                'weight': mp0_weights[f"{key_prefix}.weight"].to(dtype),
                'bias': mp0_weights[f"{key_prefix}.bias"].to(dtype),
            }
        )
    elif nemo_config["mm_cfg"]["mm_mlp_adapter_type"] == "mlp_downsample":
        vision_connector = torch.nn.Sequential(
            DownSampleBlock(),
            torch.nn.LayerNorm(vision_config["hidden_size"] * 4),
            torch.nn.Linear(vision_config["hidden_size"] * 4, nemo_config["hidden_size"], bias=True),
            torch.nn.GELU(),
            torch.nn.Linear(nemo_config["hidden_size"], nemo_config["hidden_size"], bias=True),
        ).to(dtype=dtype)
        key_prefix = "model.embedding.word_embeddings.adapter_layer.mm_projector_adapter.mm_projector"
        for layer in [1, 2, 4]:
            vision_connector[layer].load_state_dict(
                {
                    'weight': mp0_weights[f"{key_prefix}.{layer}.weight"].to(dtype),
                    'bias': mp0_weights[f"{key_prefix}.{layer}.bias"].to(dtype),
                }
            )

    else:
        raise ValueError(f"Unknown projector type: {nemo_config['mm_cfg']['mm_mlp_adapter_type']}")

    # export the whole wrapper
    lita_num_frames = None
    wrapper = VisionEncoderWrapper(vision_encoder, vision_connector).to(device, dtype)
    if model_type == "lita" or model_type == "vila":
        image_size = hf_config.image_size
        if model_type == "lita":
            lita_num_frames = nemo_config['mm_cfg']['lita']['sample_frames']
    else:
        image_size = hf_config.vision_config.image_size
        if model_type == "vita":
            lita_num_frames = nemo_config['mm_cfg']['lita']['sample_frames']
    dummy_image = torch.empty(
        1, 3, image_size, image_size, dtype=dtype, device=device
    )  # dummy image shape [B, C, H, W]

    export_visual_wrapper_onnx(wrapper, dummy_image, model_dir)
    build_trt_engine(
        model_type,
        [3, image_size, image_size],
        model_dir,
        vision_max_batch_size,
        dtype,
        image_size=image_size,
        num_frames=lita_num_frames if model_type == "lita" or model_type == 'vita' else None,
        nemo_config=nemo_config,
    )


def build_video_neva_engine(
    model_dir: str,
    visual_checkpoint_path: str,
    vision_max_batch_size: int = 1,
):
    """Build video neva visual engine"""
    device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
    # extract NeMo checkpoint
    with tarfile.open(visual_checkpoint_path) as tar:
        nemo_config = yaml.safe_load(tar.extractfile("./model_config.yaml"))
        try:
            # trained without TP
            mp0_weights = torch.load(tar.extractfile("./model_weights.ckpt"), map_location=device)
        except KeyError:
            # trained with TP
            mp0_weights = torch.load(tar.extractfile("./mp_rank_00/model_weights.ckpt"), map_location=device)

    vision_config = nemo_config["mm_cfg"]["vision_encoder"]

    class VisionEncoderWrapper(torch.nn.Module):
        # pylint: disable=C0115,C0116
        def __init__(self, encoder, connector):
            super().__init__()
            self.encoder = encoder
            self.connector = connector

        def forward(self, images):
            b, num_frames, c, h, w = images.shape
            images = images.view(b * num_frames, c, h, w)
            vision_x = self.encoder(pixel_values=images, output_hidden_states=True)  # [(B num_frames), C, H, W]
            vision_x = vision_x.hidden_states[-2]
            vision_x = vision_x[:, 1:]

            # reshape back to [B, num_frames, img_size, hidden_size]
            vision_x = vision_x.view(b, num_frames, -1, vision_x.shape[-1])

            vision_x = self.connector(vision_x)
            return vision_x

    encoder = AutoModel.from_pretrained(
        vision_config["from_pretrained"],
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation='eager',
    )
    vision_encoder = encoder.vision_model
    hf_config = encoder.config
    dtype = hf_config.torch_dtype

    # connector
    assert nemo_config["mm_cfg"]["mm_mlp_adapter_type"] == "linear"
    vision_connector = torch.nn.Linear(vision_config["hidden_size"], nemo_config["hidden_size"], bias=True)

    key_prefix = "model.embedding.word_embeddings.adapter_layer.mm_projector_adapter.mm_projector"
    vision_connector.load_state_dict(
        {
            'weight': mp0_weights[f"{key_prefix}.weight"].to(dtype),
            'bias': mp0_weights[f"{key_prefix}.bias"].to(dtype),
        }
    )

    # export the whole wrapper
    wrapper = VisionEncoderWrapper(vision_encoder, vision_connector).to(device, dtype)
    image_size = hf_config.vision_config.image_size
    num_frames = nemo_config['data']['num_frames']
    dummy_video = torch.empty(1, num_frames, 3, image_size, image_size, dtype=dtype, device=device)  # dummy image
    export_visual_wrapper_onnx(wrapper, dummy_video, model_dir)
    build_trt_engine(
        "video-neva",
        [num_frames, 3, image_size, image_size],  # [num_frames, 3, H, W]
        model_dir,
        vision_max_batch_size,
        dtype,
        image_size=image_size,
        num_frames=num_frames,
    )


def build_perception_engine(
    model_dir: str,
    perception_checkpoint_path: str,
    model_type: str = "salm",
    max_batch_size: int = 1,
):
    """Build perception engine"""
    assert model_type == "salm", f"Invalid model type {model_type}"

    def load_perception_model(perception_checkpoint_path):
        weights = "model_weights.ckpt"
        perception_state_dict = torch.load(os.path.join(perception_checkpoint_path, weights))
        config = "model_config.yaml"
        config = OmegaConf.load(os.path.join(perception_checkpoint_path, config))
        perception = AudioPerceptionModule(cfg=config)
        perception.load_state_dict(perception_state_dict)
        perception.eval()
        return perception

    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    # load perception model
    perception_model = load_perception_model(perception_checkpoint_path)
    feature_extractor = perception_model.preprocessor
    input_signal = torch.randn(1, 1000, dtype=torch.float32)
    input_signal_length = torch.tensor([1000], dtype=torch.int32)

    processed_signal, processed_signal_length = feature_extractor(
        input_signal=input_signal, length=input_signal_length
    )
    processed_signal_length = processed_signal_length.to(torch.int32)
    dump_path = model_dir + "/feature_extractor.ts"  # dump the feature extractor as torchscript
    feature_extractor.export(dump_path, (input_signal, input_signal_length))

    class PerceptionWrapper(torch.nn.Module):
        # pylint: disable=C0115,C0116
        def __init__(self, encoder, modality_adapter, proj):
            super().__init__()
            self.encoder = encoder
            self.modality_adapter = modality_adapter
            self.proj = proj

        @typecheck.disable_checks()
        def forward(self, processed_signal, processed_signal_length):
            encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
            encoded, encoded_len = self.modality_adapter(audio_signal=encoded, length=encoded_len)
            # b, c, t -> b, t, c
            encoded = self.proj(encoded.transpose(1, 2))
            encoded_len = encoded_len.to(torch.int32)
            return encoded, encoded_len

    perception = PerceptionWrapper(perception_model.encoder, perception_model.modality_adapter, perception_model.proj)
    export_perception_wrapper_onnx(perception, (processed_signal, processed_signal_length), model_dir)
    # export the onnx perception model to tensorrt engine
    # 512 -> 5.12 sec, 3072 -> 30.72 sec
    opt_batch_size = max(1, max_batch_size // 2)
    shapes = [
        {"processed_signal": [1, 80, 64], "processed_signal_length": [1]},
        {"processed_signal": [opt_batch_size, 80, 512], "processed_signal_length": [opt_batch_size]},
        {"processed_signal": [max_batch_size, 80, 3072], "processed_signal_length": [max_batch_size]},
    ]
    build_trt_engine(
        model_type,
        shapes,
        model_dir,
        max_batch_size,
        dtype=torch.float16,
        nemo_config=None,
        part_name='perception_encoder',
    )


def build_mllama_visual_engine(
    model_dir: str,
    hf_model_path: str,
    processor_name: str = "meta-llama/Llama-3.2-11B-Vision-Instruct",
    vision_max_batch_size: int = 1,
):
    """Build mllama visual engine"""
    hf_model = MllamaForConditionalGeneration.from_pretrained(hf_model_path, torch_dtype="auto", device_map="auto")
    model_dtype = hf_model.dtype

    class MLLaMAVisionWrapper(torch.nn.Module):
        # pylint: disable=C0115,C0116
        def __init__(self, vision_model, output_proj):
            super().__init__()
            self.vision_model = vision_model
            self.output_proj = output_proj

        def forward(self, pixel_values, aspect_ratio_ids, aspect_ratio_mask):
            out = self.vision_model(pixel_values, aspect_ratio_ids, aspect_ratio_mask).last_hidden_state
            out = self.output_proj(out)
            return out

    wrapper = MLLaMAVisionWrapper(hf_model.vision_model, hf_model.multi_modal_projector)

    processor = AutoProcessor.from_pretrained(processor_name)
    image = Image.new('RGB', [2048, 2688])
    inputs = processor(images=image, return_tensors="pt").to(model_dtype)

    export_visual_wrapper_onnx(
        wrapper,
        tuple([value for _, value in inputs.items()]),
        model_dir,
        input_names=[key for key in inputs],
        dynamic_axes={key: {0: "batch"} for key in inputs},
    )
    shapes = [{k: list(v.shape) for k, v in inputs.items()}] * 3
    shapes[2] = shapes[0].copy()
    for k, v in shapes[2].items():
        shapes[2][k] = [vision_max_batch_size] + v[1:]
    build_trt_engine("mllama", shapes, model_dir, vision_max_batch_size, model_dtype)


def build_visual_engine(
    model_dir: str,
    visual_checkpoint_path: str,
    model_type: str = "neva",
    vision_max_batch_size: int = 1,
):
    """Build visual engine"""
    model_list = ['neva', 'lita', 'vila', 'vita']
    if model_type in model_list:
        build_neva_engine(model_type, model_dir, visual_checkpoint_path, vision_max_batch_size)
    elif model_type == "video-neva":
        build_video_neva_engine(model_dir, visual_checkpoint_path, vision_max_batch_size)
    else:
        raise RuntimeError(f"Invalid model type {model_type}")


def extract_lora_ckpt(
    lora_ckpt: str,
    output_dir: str,
):
    """Extrace lora from checkpoint"""
    if os.path.exists(os.path.join(lora_ckpt, "model_weights.ckpt")):
        model_weight = torch.load(os.path.join(lora_ckpt, "model_weights.ckpt"))
    elif os.path.exists(os.path.join(lora_ckpt, "mp_rank_00", "model_weights.ckpt")):
        model_weight = torch.load(os.path.join(lora_ckpt, "mp_rank_00", "model_weights.ckpt"))
    else:
        raise RuntimeError("Imcompatible lora checkpoint format")

    model_config = os.path.join(lora_ckpt, "model_config.yaml")

    if not os.path.exists(model_config):
        raise RuntimeError("Imcompatible lora checkpoint format")

    llm_lora_weight = {}

    for k, v in model_weight.items():
        if "mm_projector" not in k:
            llm_lora_weight[k] = v

    llm_lora_path = os.path.join(output_dir, "llm_lora.nemo")
    with tempfile.TemporaryDirectory() as tmp_dir:
        llm_weight_path = os.path.join(tmp_dir, "model_weights.ckpt")
        torch.save(llm_lora_weight, llm_weight_path)

        with tarfile.open(llm_lora_path, "w") as tar:
            tar.add(llm_weight_path, arcname="model_weights.ckpt")
            tar.add(model_config, arcname="model_config.yaml")

    return llm_lora_path


def build_mllama_engine(
    model_dir: str,
    checkpoint_path: str,
    processor_name: str = "meta-llama/Llama-3.2-11B-Vision-Instruct",
    vision_max_batch_size: int = 1,
    tensor_parallelism_size: int = 1,
    max_input_len: int = 256,
    max_output_len: int = 256,
    max_batch_size: int = 1,
    max_multimodal_len: int = 1024,
    dtype: str = "bfloat16",
    use_lora_plugin: str = None,
    lora_target_modules: List[str] = None,
    max_lora_rank: int = 64,
    lora_ckpt_list: List[str] = None,
):
    """Build mllama engine"""
    new_state_dict, config = convert_mllama_nemo_to_hf(checkpoint_path, processor_name)

    hf_model = MllamaForConditionalGeneration(config)
    hf_model = hf_model.to(torch.bfloat16)
    hf_model.load_state_dict(new_state_dict)

    with tempfile.TemporaryDirectory() as tmp_dir:
        hf_model_path = os.path.join(tmp_dir, "hf_checkpoint")
        hf_model.save_pretrained(hf_model_path)
        del hf_model, new_state_dict

        build_mllama_visual_engine(
            os.path.join(model_dir, "visual_engine"),
            hf_model_path,
            vision_max_batch_size=vision_max_batch_size,
        )
        build_mllama_trtllm_engine(
            os.path.join(model_dir, "llm_engine"),
            hf_model_path,
            tensor_parallelism_size,
            max_input_len,
            max_output_len,
            max_batch_size,
            max_multimodal_len,
            dtype,
        )