| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import itertools |
| 5 | from functools import partial |
| 6 | |
| 7 | import pytest |
| 8 | from PIL import Image |
| 9 | from pqdm.threads import pqdm |
| 10 | |
| 11 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 12 | from vllm.multimodal.parse import ImageSize |
| 13 | from vllm.multimodal.processing import BaseMultiModalProcessor |
| 14 | from vllm.tokenizers.hf import maybe_make_thread_pool |
| 15 | |
| 16 | from ...utils import build_model_context |
| 17 | |
| 18 | |
| 19 | def _validate_image_max_tokens_one( |
| 20 | processor: BaseMultiModalProcessor, |
| 21 | max_tokens: int, |
| 22 | failed_size_excs: list[tuple[ImageSize, Exception]], |
| 23 | image_size: ImageSize, |
| 24 | ) -> None: |
| 25 | info = processor.info |
| 26 | feature_size = info.get_num_image_tokens( |
| 27 | image_width=image_size.width, image_height=image_size.height |
| 28 | ) |
| 29 | |
| 30 | try: |
| 31 | assert feature_size <= max_tokens, f"{feature_size} <= {max_tokens}" |
| 32 | except Exception as exc: |
| 33 | failed_size_excs.append((image_size, exc)) |
| 34 | |
| 35 | |
| 36 | @pytest.mark.skip( |
| 37 | "This test takes around 5 minutes to run. Comment this out to run it manually." |
| 38 | ) |
| 39 | @pytest.mark.parametrize("model_id", ["llava-hf/llava-v1.6-mistral-7b-hf"]) |
| 40 | def test_processor_max_tokens(model_id): |
| 41 | ctx = build_model_context( |
| 42 | model_id, |
| 43 | mm_processor_kwargs=None, |
| 44 | limit_mm_per_prompt={"image": 1}, |
| 45 | ) |
| 46 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 47 | info = processor.info |
| 48 | |
| 49 | seen_aspect_ratios = set[float]() |
| 50 | image_sizes = list[ImageSize]() |
| 51 | |
| 52 | # The aspect ratio of the grid layout is between 1 and 2 |
| 53 | # NOTE: Assumes that feature size calculation is the same if we |
| 54 | # swap the width and height of the image |
| 55 | for w, h in itertools.product(range(32, 4096), repeat=2): |
| 56 | aspect_ratio = w / h |
| 57 | if 1 <= aspect_ratio <= 2 and aspect_ratio not in seen_aspect_ratios: |
| 58 | image_sizes.append(ImageSize(w, h)) |
| 59 | seen_aspect_ratios.add(aspect_ratio) |
| 60 | |
| 61 | failed_size_excs = list[tuple[ImageSize, Exception]]() |
| 62 | |
| 63 | validate_one = partial( |
| 64 | _validate_image_max_tokens_one, |
| 65 | processor, |
| 66 | info.get_max_image_tokens(), # type: ignore |
| 67 | failed_size_excs, |
| 68 | ) |
| 69 | pqdm(image_sizes, validate_one, n_jobs=8, desc="Validating image sizes") |
| 70 | |
| 71 | if failed_size_excs: |
| 72 | msg = "Found failing image sizes:" + "\n========\n".join( |
| 73 | f"[{size}]\n{exc}" for size, exc in failed_size_excs |
| 74 | ) |
| 75 | raise AssertionError(msg) |
| 76 | |
| 77 | |
| 78 | def _validate_image_prompt_replacements_one( |
| 79 | processor: BaseMultiModalProcessor, |
| 80 | num_imgs: int, |
| 81 | failed_size_excs: list[tuple[ImageSize, Exception]], |
| 82 | image_size: ImageSize, |
| 83 | ) -> None: |
| 84 | prompt = "<image>" * num_imgs |
| 85 | image = Image.new("RGB", size=image_size) |
| 86 | mm_data = {"image": [image] * num_imgs} |
| 87 | |
| 88 | try: |
| 89 | # The processor will throw an error if there is a mismatch |
| 90 | # in the prompt replacements |
| 91 | processed_inputs = processor( |
| 92 | prompt, |
| 93 | mm_items=processor.info.parse_mm_data(mm_data), |
| 94 | hf_processor_mm_kwargs={}, |
| 95 | ) |
| 96 | |
| 97 | image_placeholders = processed_inputs["mm_placeholders"]["image"] |
| 98 | assert len(image_placeholders) == num_imgs |
| 99 | |
| 100 | first_placeholder = image_placeholders[0] |
| 101 | |
| 102 | # NOTE: There is a BOS token |
| 103 | assert first_placeholder.offset == 1 |
| 104 | assert ( |
| 105 | first_placeholder.length |
| 106 | == (len(processed_inputs["prompt_token_ids"]) - 1) // num_imgs |
| 107 | ) |
| 108 | |
| 109 | except Exception as exc: |
| 110 | failed_size_excs.append((image_size, exc)) |
| 111 | |
| 112 | |
| 113 | def _test_image_prompt_replacements( |
| 114 | processor, |
| 115 | *, |
| 116 | num_imgs: int, |
| 117 | image_sizes: list[ImageSize], |
| 118 | ) -> None: |
| 119 | """Ensure LlavaNextMultiModalProcessor |
| 120 | handles prompt replacement properly for input images. |
| 121 | """ |
| 122 | failed_size_excs = list[tuple[ImageSize, Exception]]() |
| 123 | |
| 124 | validate_one = partial( |
| 125 | _validate_image_prompt_replacements_one, |
| 126 | processor, |
| 127 | num_imgs, |
| 128 | failed_size_excs, |
| 129 | ) |
| 130 | pqdm(image_sizes, validate_one, n_jobs=8, desc="Validating image sizes") |
| 131 | |
| 132 | if failed_size_excs: |
| 133 | msg = "Found failing image sizes:" + "\n========\n".join( |
| 134 | f"[{size}]\n{exc}" for size, exc in failed_size_excs |
| 135 | ) |
| 136 | raise AssertionError(msg) |
| 137 | |
| 138 | |
| 139 | @pytest.mark.parametrize("model_id", ["llava-hf/llava-v1.6-mistral-7b-hf"]) |
| 140 | @pytest.mark.parametrize("num_imgs", [1, 2]) |
| 141 | def test_processor_prompt_replacements_regression(model_id, num_imgs): |
| 142 | ctx = build_model_context( |
| 143 | model_id, |
| 144 | mm_processor_kwargs=None, |
| 145 | limit_mm_per_prompt={"image": num_imgs}, |
| 146 | ) |
| 147 | |
| 148 | # Avoid tokenizer already borrowed error |
| 149 | maybe_make_thread_pool(ctx.tokenizer) |
| 150 | |
| 151 | processor = MULTIMODAL_REGISTRY.create_processor( |
| 152 | ctx.model_config, |
| 153 | tokenizer=ctx.tokenizer, |
| 154 | ) |
| 155 | |
| 156 | image_ratios = [ |
| 157 | (171, 152), |
| 158 | (184, 161), |
| 159 | (198, 176), |
| 160 | (333, 296), |
| 161 | (369, 328), |
| 162 | (488, 183), |
| 163 | (2560, 1669), |
| 164 | ] |
| 165 | image_sizes = [ |
| 166 | size for w, h in image_ratios for size in [ImageSize(w, h), ImageSize(h, w)] |
| 167 | ] |
| 168 | |
| 169 | _test_image_prompt_replacements( |
| 170 | processor, |
| 171 | num_imgs=num_imgs, |
| 172 | image_sizes=image_sizes, |
| 173 | ) |
| 174 | |
| 175 | |
| 176 | @pytest.mark.skip( |
| 177 | "This test takes around 2 hours to run. Comment this out to run it manually." |
| 178 | ) |
| 179 | @pytest.mark.parametrize("model_id", ["llava-hf/llava-v1.6-mistral-7b-hf"]) |
| 180 | @pytest.mark.parametrize("num_imgs", [1]) |
| 181 | def test_processor_prompt_replacements_all(model_id, num_imgs): |
| 182 | ctx = build_model_context( |
| 183 | model_id, |
| 184 | mm_processor_kwargs=None, |
| 185 | limit_mm_per_prompt={"image": num_imgs}, |
| 186 | ) |
| 187 | |
| 188 | # Avoid tokenizer already borrowed error |
| 189 | maybe_make_thread_pool(ctx.tokenizer) |
| 190 | |
| 191 | processor = MULTIMODAL_REGISTRY.create_processor( |
| 192 | ctx.model_config, |
| 193 | tokenizer=ctx.tokenizer, |
| 194 | ) |
| 195 | |
| 196 | seen_aspect_ratios = set[float]() |
| 197 | image_sizes = list[ImageSize]() |
| 198 | |
| 199 | # The aspect ratio of the grid layout is between 1 and 2 |
| 200 | # NOTE: Assumes that feature size calculation is the same if we |
| 201 | # swap the width and height of the image |
| 202 | for w, h in itertools.product(range(64, 1024), repeat=2): |
| 203 | aspect_ratio = w / h |
| 204 | if 1 <= aspect_ratio <= 2 and aspect_ratio not in seen_aspect_ratios: |
| 205 | image_sizes.append(ImageSize(w, h)) |
| 206 | seen_aspect_ratios.add(aspect_ratio) |
| 207 | |
| 208 | _test_image_prompt_replacements( |
| 209 | processor, |
| 210 | num_imgs=num_imgs, |
| 211 | image_sizes=image_sizes, |
| 212 | ) |