[VLM][Core] Support profiling with multiple multi-modal inputs per prompt (#7126)

tests/multimodal/test_mapper.py

@@ -1,15 +1,22 @@
+from contextlib import nullcontext
+
 import numpy as np
 import pytest
 from transformers import CLIPImageProcessor, LlavaNextImageProcessor
 
-from vllm.config import ModelConfig
-from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.config import ModelConfig, MultiModalConfig
+from vllm.multimodal import MultiModalRegistry
 from vllm.multimodal.utils import rescale_image_size
 
 
+@pytest.fixture
+def mm_registry():
+    return MultiModalRegistry()
+
+
 @pytest.mark.parametrize("dtype", ["half", "float"])
 @pytest.mark.parametrize("size_factor", [0.25, 0.5, 1.0])
-def test_clip_image_processor(image_assets, dtype, size_factor):
+def test_clip_image_processor(image_assets, mm_registry, dtype, size_factor):
     MODEL_NAME = "llava-hf/llava-1.5-7b-hf"
 
     hf_processor = CLIPImageProcessor.from_pretrained(MODEL_NAME)
@@ -24,6 +31,9 @@ def test_clip_image_processor(image_assets, dtype, size_factor):
         dtype=dtype,
         revision=None,
     )
+    mm_config = MultiModalConfig(limit_per_prompt={"image": 1})
+
+    mm_registry.init_mm_limits_per_prompt(model_config, mm_config)
 
     for asset in image_assets:
         image = rescale_image_size(asset.pil_image, size_factor)
@@ -32,7 +42,7 @@ def test_clip_image_processor(image_assets, dtype, size_factor):
             image,
             return_tensors="pt",
         )
-        vllm_result = MULTIMODAL_REGISTRY.map_input(
+        vllm_result = mm_registry.map_input(
             model_config,
             {"image": image},
         )
@@ -48,7 +58,8 @@ def test_clip_image_processor(image_assets, dtype, size_factor):
 
 @pytest.mark.parametrize("dtype", ["half", "float"])
 @pytest.mark.parametrize("size_factor", [0.25, 0.5, 1.0])
-def test_llava_next_image_processor(image_assets, dtype, size_factor):
+def test_llava_next_image_processor(image_assets, mm_registry, dtype,
+                                    size_factor):
     MODEL_NAME = "llava-hf/llava-v1.6-vicuna-7b-hf"
 
     hf_processor = LlavaNextImageProcessor.from_pretrained(MODEL_NAME)
@@ -63,6 +74,9 @@ def test_llava_next_image_processor(image_assets, dtype, size_factor):
         dtype=dtype,
         revision=None,
     )
+    mm_config = MultiModalConfig(limit_per_prompt={"image": 1})
+
+    mm_registry.init_mm_limits_per_prompt(model_config, mm_config)
 
     for asset in image_assets:
         image = rescale_image_size(asset.pil_image, size_factor)
@@ -71,7 +85,7 @@ def test_llava_next_image_processor(image_assets, dtype, size_factor):
             image,
             return_tensors="pt",
         )
-        vllm_result = MULTIMODAL_REGISTRY.map_input(
+        vllm_result = mm_registry.map_input(
             model_config,
             {"image": image},
         )
@@ -83,3 +97,61 @@ def test_llava_next_image_processor(image_assets, dtype, size_factor):
 
             assert hf_arr.shape == vllm_arr.shape, f"Failed for key={key}"
             assert np.allclose(hf_arr, vllm_arr), f"Failed for key={key}"
+
+
+@pytest.mark.parametrize(
+    ("num_images", "limit", "is_valid"),
+    [(0, 0, True), (0, 1, True), (1, 0, False), (1, 1, True), (1, 2, True),
+     (2, 1, False), (2, 2, True)],
+)
+def test_mm_limits(image_assets, mm_registry, num_images, limit, is_valid):
+    MODEL_NAME = "llava-hf/llava-1.5-7b-hf"
+
+    model_config = ModelConfig(
+        model=MODEL_NAME,
+        tokenizer=MODEL_NAME,
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        seed=0,
+        dtype="half",
+        revision=None,
+    )
+    mm_config = MultiModalConfig(limit_per_prompt={"image": limit})
+
+    mm_registry.init_mm_limits_per_prompt(model_config, mm_config)
+
+    image = image_assets[0].pil_image
+    if num_images == 0:
+        mm_inputs = {}
+    elif num_images == 1:
+        mm_inputs = {"image": image}
+    else:
+        mm_inputs = {"image": [image] * num_images}
+
+    with nullcontext() if is_valid else pytest.raises(ValueError):
+        mm_registry.map_input(model_config, mm_inputs)
+
+
+# NOTE: We don't test zero images since the HF processor doesn't support it
+@pytest.mark.parametrize("num_images", [1, 2])
+def test_image_mapper_multi(image_assets, mm_registry, num_images):
+    MODEL_NAME = "llava-hf/llava-1.5-7b-hf"
+
+    model_config = ModelConfig(
+        model=MODEL_NAME,
+        tokenizer=MODEL_NAME,
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        seed=0,
+        dtype="half",
+        revision=None,
+    )
+    mm_config = MultiModalConfig(limit_per_prompt={"image": num_images})
+
+    mm_registry.init_mm_limits_per_prompt(model_config, mm_config)
+
+    image = image_assets[0].pil_image
+    mm_inputs = {"image": [image] * num_images}
+
+    mapped_inputs = mm_registry.map_input(model_config, mm_inputs)
+    assert len(mapped_inputs["pixel_values"]) == num_images