| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Mapping |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | from PIL import Image as PILImage |
| 9 | from transformers.models.gemma4 import Gemma4AudioFeatureExtractor |
| 10 | |
| 11 | from vllm.exceptions import VLLMValidationError |
| 12 | from vllm.model_executor.models.gemma4_mm import ( |
| 13 | Gemma4ImagePixelInputs, |
| 14 | _dummy_audio_num_samples, |
| 15 | ) |
| 16 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 17 | from vllm.multimodal.inputs import MultiModalFieldConfig |
| 18 | |
| 19 | from ....conftest import ImageTestAssets |
| 20 | from ...utils import build_model_context |
| 21 | |
| 22 | # The Unified model ID for testing purposes |
| 23 | GEMMA4_UNIFIED_MODEL_ID = "google/gemma-4-12B-it" |
| 24 | |
| 25 | |
| 26 | def test_gemma4_unified_image_schema_accepts_variable_patch_counts(): |
| 27 | Gemma4ImagePixelInputs( |
| 28 | pixel_values=[ |
| 29 | torch.randn(10080, 768), |
| 30 | torch.randn(2520, 768), |
| 31 | ], |
| 32 | pixel_position_ids=[ |
| 33 | torch.zeros(10080, 2, dtype=torch.long), |
| 34 | torch.zeros(2520, 2, dtype=torch.long), |
| 35 | ], |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | def test_gemma4_unified_image_batching_keeps_variable_patch_counts_unstacked(): |
| 40 | field = MultiModalFieldConfig.batched("image").field |
| 41 | elems = field.build_elems( |
| 42 | "image", |
| 43 | "pixel_values", |
| 44 | [torch.randn(10080, 768), torch.randn(2520, 768)], |
| 45 | ) |
| 46 | |
| 47 | reduced = field.reduce_data(list(elems)) |
| 48 | |
| 49 | assert isinstance(reduced, list) |
| 50 | assert [tensor.shape for tensor in reduced] == [ |
| 51 | torch.Size([10080, 768]), |
| 52 | torch.Size([2520, 768]), |
| 53 | ] |
| 54 | |
| 55 | |
| 56 | @pytest.mark.parametrize( |
| 57 | "image_width,image_height,max_soft_tokens", |
| 58 | [ |
| 59 | (900, 3, 280), |
| 60 | (3, 900, 280), |
| 61 | (900, 3, 70), |
| 62 | (4000, 2, 1120), |
| 63 | ], |
| 64 | ) |
| 65 | @pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) |
| 66 | def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens( |
| 67 | model_id: str, |
| 68 | image_width: int, |
| 69 | image_height: int, |
| 70 | max_soft_tokens: int, |
| 71 | ): |
| 72 | """Verify ``_compute_num_soft_tokens`` caps output at ``max_soft_tokens``.""" |
| 73 | ctx = build_model_context( |
| 74 | model_id, |
| 75 | mm_processor_kwargs={"do_pan_and_scan": True}, |
| 76 | limit_mm_per_prompt={"image": 1}, |
| 77 | ) |
| 78 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 79 | |
| 80 | num_soft_tokens = processor.info._compute_num_soft_tokens( |
| 81 | image_width=image_width, |
| 82 | image_height=image_height, |
| 83 | max_soft_tokens=max_soft_tokens, |
| 84 | ) |
| 85 | |
| 86 | assert num_soft_tokens <= max_soft_tokens, ( |
| 87 | f"_compute_num_soft_tokens returned {num_soft_tokens} for " |
| 88 | f"image_width={image_width}, image_height={image_height}, " |
| 89 | f"max_soft_tokens={max_soft_tokens} — exceeds the cap." |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | @pytest.mark.parametrize( |
| 94 | ("mm_processor_kwargs", "expected_image_tokens"), |
| 95 | [ |
| 96 | ({}, 280), |
| 97 | ({"max_soft_tokens": 70}, 70), |
| 98 | ({"max_soft_tokens": 280}, 280), |
| 99 | ({"max_soft_tokens": 1120}, 1120), |
| 100 | ({"images_kwargs": {"max_soft_tokens": 560}}, 560), |
| 101 | ({"images_kwargs": None}, 280), |
| 102 | ({"images_kwargs": "not-a-dict"}, 280), |
| 103 | ], |
| 104 | ) |
| 105 | @pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) |
| 106 | def test_get_mm_max_tokens_per_item_respects_configured_max_soft_tokens( |
| 107 | model_id: str, |
| 108 | mm_processor_kwargs: dict[str, object], |
| 109 | expected_image_tokens: int, |
| 110 | ): |
| 111 | ctx = build_model_context( |
| 112 | model_id, |
| 113 | mm_processor_kwargs=mm_processor_kwargs, |
| 114 | limit_mm_per_prompt={"image": 1, "video": 1}, |
| 115 | ) |
| 116 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 117 | |
| 118 | tokens = processor.info.get_mm_max_tokens_per_item( |
| 119 | seq_len=ctx.model_config.max_model_len, |
| 120 | mm_counts={"image": 1, "video": 1}, |
| 121 | ) |
| 122 | |
| 123 | assert tokens is not None |
| 124 | assert tokens["image"] == expected_image_tokens |
| 125 | assert tokens["video"] == 32 * (70 + 2 + 6) |
| 126 | |
| 127 | |
| 128 | @pytest.mark.parametrize( |
| 129 | ("limit_mm_per_prompt", "expected_video_tokens"), |
| 130 | [ |
| 131 | ({"video": 1}, 32 * (70 + 2 + 6)), |
| 132 | ({"video": {"count": 1}}, 32 * (70 + 2 + 6)), |
| 133 | ({"video": {"count": 1, "num_frames": 1}}, 1 * (70 + 2 + 6)), |
| 134 | ({"video": {"count": 1, "num_frames": 8}}, 8 * (70 + 2 + 6)), |
| 135 | ({"video": {"count": 1, "num_frames": 32}}, 32 * (70 + 2 + 6)), |
| 136 | ({"video": {"count": 1, "num_frames": 40}}, 32 * (70 + 2 + 6)), |
| 137 | ], |
| 138 | ) |
| 139 | @pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) |
| 140 | def test_get_mm_max_tokens_per_item_respects_configured_video_num_frames( |
| 141 | model_id: str, |
| 142 | limit_mm_per_prompt: Mapping[str, int | Mapping[str, int]], |
| 143 | expected_video_tokens: int, |
| 144 | ): |
| 145 | ctx = build_model_context( |
| 146 | model_id, |
| 147 | limit_mm_per_prompt=limit_mm_per_prompt, |
| 148 | ) |
| 149 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 150 | |
| 151 | tokens = processor.info.get_mm_max_tokens_per_item( |
| 152 | seq_len=ctx.model_config.max_model_len, |
| 153 | mm_counts={"video": 1}, |
| 154 | ) |
| 155 | |
| 156 | assert tokens is not None |
| 157 | assert tokens["image"] == 280 |
| 158 | assert tokens["video"] == expected_video_tokens |
| 159 | |
| 160 | |
| 161 | @pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) |
| 162 | def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str): |
| 163 | ctx = build_model_context( |
| 164 | model_id, |
| 165 | mm_processor_kwargs={"images_kwargs": {"max_soft_tokens": 560}}, |
| 166 | limit_mm_per_prompt={"image": 1}, |
| 167 | ) |
| 168 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 169 | image = PILImage.new("RGB", (1000, 1000), color="white") |
| 170 | image_size = image.size |
| 171 | mm_items = processor.info.parse_mm_data({"image": image}) |
| 172 | |
| 173 | prompt_update = processor._get_prompt_updates(mm_items, {}, {})[0] |
| 174 | replacement = prompt_update.resolve(0).content.full |
| 175 | expected = processor.info.get_image_repl( |
| 176 | image_width=image_size[0], |
| 177 | image_height=image_size[1], |
| 178 | processor=processor.info.get_hf_processor(), |
| 179 | max_soft_tokens=560, |
| 180 | ).full |
| 181 | |
| 182 | assert replacement == expected |
| 183 | |
| 184 | |
| 185 | @pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) |
| 186 | def test_limit_mm_per_prompt( |
| 187 | image_assets: ImageTestAssets, |
| 188 | model_id: str, |
| 189 | ): |
| 190 | """Test that limit_mm_per_prompt restricts multiple images correctly.""" |
| 191 | ctx = build_model_context( |
| 192 | model_id, |
| 193 | mm_processor_kwargs={}, |
| 194 | limit_mm_per_prompt={"image": 1}, |
| 195 | ) |
| 196 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 197 | |
| 198 | prompt = "<image><image>" |
| 199 | images = [asset.pil_image for asset in image_assets][:2] |
| 200 | if len(images) < 2: |
| 201 | images = [images[0], images[0]] |
| 202 | |
| 203 | mm_data = {"image": images} |
| 204 | |
| 205 | with pytest.raises(VLLMValidationError, match="At most 1 image"): |
| 206 | processor( |
| 207 | prompt, |
| 208 | mm_items=processor.info.parse_mm_data(mm_data), |
| 209 | hf_processor_mm_kwargs={}, |
| 210 | ) |
| 211 | |
| 212 | |
| 213 | def test_dummy_audio_length_falls_back_for_the_unified_extractor(): |
| 214 | """The unified audio feature extractor does not define ``fft_length``. |
| 215 | |
| 216 | ``Gemma4DummyInputsBuilder.get_dummy_mm_data`` read that attribute |
| 217 | unconditionally, so memory profiling raised ``AttributeError`` for every |
| 218 | unified Gemma 4 deployment in which audio is the modality with the largest |
| 219 | per-item token count -- which is the case as soon as ``audio_seq_length`` |
| 220 | is raised above the video budget. |
| 221 | """ |
| 222 | |
| 223 | class _UnifiedFeatureExtractor: |
| 224 | sampling_rate = 16000 |
| 225 | |
| 226 | class _UnifiedProcessor: |
| 227 | feature_extractor = _UnifiedFeatureExtractor() |
| 228 | audio_seq_length = 3000 |
| 229 | audio_ms_per_token = 40 |
| 230 | |
| 231 | # 3000 tokens * 40 ms = 120 s at 16 kHz |
| 232 | assert _dummy_audio_num_samples(_UnifiedProcessor()) == 120 * 16000 |
| 233 | |
| 234 | |
| 235 | def test_dummy_audio_length_still_uses_fft_length_for_the_tower_extractor(): |
| 236 | """The mel/tower extractor keeps its existing behaviour.""" |
| 237 | |
| 238 | class _TowerFeatureExtractor(Gemma4AudioFeatureExtractor): |
| 239 | def __init__(self): # noqa: D107 - bypass the real feature-extractor init |
| 240 | self.fft_length = 512 |
| 241 | self.sampling_rate = 16000 |
| 242 | |
| 243 | class _TowerProcessor: |
| 244 | feature_extractor = _TowerFeatureExtractor() |
| 245 | audio_seq_length = 750 |
| 246 | audio_ms_per_token = 40 |
| 247 | |
| 248 | assert _dummy_audio_num_samples(_TowerProcessor()) == 512 |