add additional arguments for mm data

Signed-off-by: John Calderon <jcalderon@nvidia.com>

Author: johncalesp

tensorrt_llm/_torch/models/modeling_qwen2vl.py

@@ -234,18 +234,17 @@ def get_dummy_text(self, input_seq_len: int):
             for _ in range(input_seq_len)
         ])
 
-    def get_dummy_images(self,
-                         max_width: int,
-                         max_height: int,
-                         num_images: int = 1):
+    def get_dummy_images(self, max_width: int, max_height: int,
+                         num_images: int):
         image = Image.new("RGB", (max_width, max_height), color=255)
         return [image] * num_images
 
-    def get_dummy_prompt(self, input_seq_len: int):
+    def get_dummy_prompt(self, input_seq_len: int, mm_data: dict):
+        num_images = mm_data.get("image", 0)
         text = self.get_dummy_text(input_seq_len)
         images = self.get_dummy_images(
-            max_width=3584,
-            max_height=3584)  #sqrt of max_pixels value (12845056)
+            max_width=3584, max_height=3584,
+            num_images=num_images)  #w, h is sqrt of max_pixels value (12845056)
         return default_multimodal_input_loader(
             tokenizer=self.tokenizer,
             model_dir=self.model_path,

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -163,7 +163,8 @@ def _create_dummy_mm_context_request(
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "ViT's encoder")
             return requests
-        text_prompt = input_processor.get_dummy_prompt(input_seq_len)
+        text_prompt = input_processor.get_dummy_prompt(input_seq_len,
+                                                       {'image': 1})
         max_beam_width = self._max_beam_width
         input_processor_with_hash = create_input_processor_with_hash(
             input_processor)

tensorrt_llm/inputs/registry.py

@@ -47,7 +47,7 @@ class BaseDummyInputsBuilder:
     Base class for generating dummy inputs. Specially for profiling
     """
 
-    def get_dummy_prompt(self):
+    def get_dummy_prompt(self, input_seq_len: int, mm_data: dict):
         raise NotImplementedError(
             "Please ensure this method is implemented in your inherited class")