fix pre-commit

Signed-off-by: Yibin Li <[email protected]>

Author: yibinl-nvidia

examples/models/core/multimodal/requirements-eclair.txt

@@ -1 +1 @@
-timm
+timm

tensorrt_llm/models/enc_dec/model.py

@@ -20,7 +20,8 @@
 import torch
 
 from tensorrt_llm._common import default_net
-from tensorrt_llm._utils import numpy_to_torch, str_dtype_to_torch
+from tensorrt_llm._utils import (numpy_to_torch, pad_vocab_size,
+                                 str_dtype_to_torch)
 from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType,
                                      MLPType, PositionEmbeddingType, Tensor,
                                      assertion, cast, gather_last_token_logits,
@@ -44,8 +45,6 @@
 from tensorrt_llm.parameter import Parameter
 from tensorrt_llm.plugin.plugin import current_all_reduce_helper
 from tensorrt_llm.quantization import QuantMode
-from tensorrt_llm._utils import pad_vocab_size
-
 
 layernorm_map = {
     LayerNormType.LayerNorm: LayerNorm,
@@ -1159,7 +1158,7 @@ def __init__(self, config: PretrainedConfig):
 
         if self.mapping.is_last_pp_rank():
             vocab_size_padded = pad_vocab_size(self.config.vocab_size,
-                                           self.config.mapping.tp_size)
+                                               self.config.mapping.tp_size)
             self.lm_head = ColumnLinear(
                 self.config.hidden_size,
                 vocab_size_padded,

tensorrt_llm/tools/multimodal_builder.py

@@ -25,26 +25,25 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
-from safetensors.torch import save_file
-from safetensors.torch import load_model
+from safetensors.torch import load_model, save_file
 from transformers import CLIPImageProcessor
 
 from ..runtime.session import Session
 
 
 def add_multimodal_arguments(parser):
-    parser.add_argument('--model_type',
-                        type=str,
-                        default=None,
-                        choices=[
-                            'blip2', 'llava', 'llava_next', 'llava_onevision',
-                            'llava_onevision_lmms', 'vila', 'nougat', 'cogvlm',
-                            'fuyu', 'pix2struct', 'neva', 'kosmos-2',
-                            'video-neva', 'phi-3-vision', 'phi-4-multimodal',
-                            'mllama', 'internvl', 'qwen2_vl',
-                            'internlm-xcomposer2', 'qwen2_audio', 'pixtral', 'eclair'
-                        ],
-                        help="Model type")
+    parser.add_argument(
+        '--model_type',
+        type=str,
+        default=None,
+        choices=[
+            'blip2', 'llava', 'llava_next', 'llava_onevision',
+            'llava_onevision_lmms', 'vila', 'nougat', 'cogvlm', 'fuyu',
+            'pix2struct', 'neva', 'kosmos-2', 'video-neva', 'phi-3-vision',
+            'phi-4-multimodal', 'mllama', 'internvl', 'qwen2_vl',
+            'internlm-xcomposer2', 'qwen2_audio', 'pixtral', 'eclair'
+        ],
+        help="Model type")
     parser.add_argument(
         '--model_path',
         type=str,
@@ -1743,20 +1742,33 @@ def forward(self, pixel_values, attention_mask):
         engine_name=f"model.engine",
         dtype=torch.bfloat16)
 
+
 def build_eclair_engine(args):
-    
+
     class RadioWithNeck(torch.nn.Module):
+
         def __init__(self):
             super().__init__()
 
-            self.model_encoder = torch.hub.load("NVlabs/RADIO", "radio_model", version="radio_v2.5-h")
+            self.model_encoder = torch.hub.load("NVlabs/RADIO",
+                                                "radio_model",
+                                                version="radio_v2.5-h")
             self.model_encoder.summary_idxs = torch.tensor(4)
 
             self.conv1 = torch.nn.Conv1d(1280, 1024, 1)
-            self.layer_norm1 = torch.nn.LayerNorm(1024, eps=1e-6, elementwise_affine=True)
-            self.conv2 = torch.nn.Conv2d(1024, 1024, kernel_size=(1, 4), stride=(1, 4), padding=0, bias=False)
-            self.layer_norm2 = torch.nn.LayerNorm(1024, eps=1e-6, elementwise_affine=True)
-        
+            self.layer_norm1 = torch.nn.LayerNorm(1024,
+                                                  eps=1e-6,
+                                                  elementwise_affine=True)
+            self.conv2 = torch.nn.Conv2d(1024,
+                                         1024,
+                                         kernel_size=(1, 4),
+                                         stride=(1, 4),
+                                         padding=0,
+                                         bias=False)
+            self.layer_norm2 = torch.nn.LayerNorm(1024,
+                                                  eps=1e-6,
+                                                  elementwise_affine=True)
+
         @torch.no_grad
         def forward(self, pixel_values):
             _, feature = self.model_encoder(pixel_values)
@@ -1770,26 +1782,29 @@ def forward(self, pixel_values):
             output = output.flatten(-2, -1).permute(0, 2, 1)
             output = self.layer_norm2(output)
             return output
-    
+
     processor = NougatProcessor.from_pretrained(args.model_path)
     model = VisionEncoderDecoderModel.from_pretrained("facebook/nougat-base")
     model.encoder = RadioWithNeck()
     model.decoder.resize_token_embeddings(len(processor.tokenizer))
-    model.config.decoder_start_token_id = processor.tokenizer.eos_token_id # 2
+    model.config.decoder_start_token_id = processor.tokenizer.eos_token_id  # 2
     model.config.pad_token_id = processor.tokenizer.pad_token_id  # 1
     load_model(model, os.path.join(args.model_path, "model.safetensors"))
-    
+
     wrapper = model.encoder.to(args.device)
     # temporary fix due to TRT onnx export bug
     for block in wrapper.model_encoder.model.blocks:
         block.attn.fused_attn = False
-    
-    image = torch.randn((1, 3, 2048, 1648), device=args.device, dtype=torch.float16)
+
+    image = torch.randn((1, 3, 2048, 1648),
+                        device=args.device,
+                        dtype=torch.float16)
     export_onnx(wrapper, image, f'{args.output_dir}/onnx')
     build_trt_engine(
         args.model_type,
         [image.shape[1], image.shape[2], image.shape[3]],  # [3, H, W]
         f'{args.output_dir}/onnx',
         args.output_dir,
         args.max_batch_size,
-        dtype=torch.bfloat16,engine_name='visual_encoder.engine')
+        dtype=torch.bfloat16,
+        engine_name='visual_encoder.engine')