| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Any, TypedDict |
| 5 | |
| 6 | import numpy.typing as npt |
| 7 | import pytest |
| 8 | import torch |
| 9 | from PIL import Image |
| 10 | |
| 11 | from vllm.multimodal.image import rescale_image_size |
| 12 | from vllm.multimodal.video import rescale_video_size, sample_frames_from_video |
| 13 | |
| 14 | from ....conftest import ( |
| 15 | IMAGE_ASSETS, |
| 16 | VIDEO_ASSETS, |
| 17 | PromptImageInput, |
| 18 | PromptVideoInput, |
| 19 | VllmRunner, |
| 20 | ) |
| 21 | from ...utils import check_logprobs_close |
| 22 | |
| 23 | |
| 24 | @pytest.fixture(scope="function", autouse=True) |
| 25 | def enable_pickle(monkeypatch): |
| 26 | """`LLM.apply_model` requires pickling a function.""" |
| 27 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 28 | |
| 29 | |
| 30 | models = ["Qwen/Qwen2-VL-2B-Instruct"] |
| 31 | target_dtype = "half" |
| 32 | IMAGE_SIZE_FACTOR_GROUPS = ( |
| 33 | (0.5,), |
| 34 | (0.5, 0.5), |
| 35 | (0.25, 0.5, 0.5), |
| 36 | ) |
| 37 | VIDEO_SIZE_FACTOR_GROUPS = ( |
| 38 | (0.5,), |
| 39 | (0.5, 0.5), |
| 40 | (0.25, 0.25, 0.5), |
| 41 | ) |
| 42 | |
| 43 | IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" |
| 44 | VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>" |
| 45 | MODEL_HIDDEN_SIZE = 1536 |
| 46 | |
| 47 | |
| 48 | def qwen2_vl_chat_template(*query): |
| 49 | return f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{''.join(query)}<|im_end|><|im_start|>assistant\n" # noqa: E501 |
| 50 | |
| 51 | |
| 52 | IMAGE_PROMPTS = IMAGE_ASSETS.prompts( |
| 53 | { |
| 54 | "stop_sign": qwen2_vl_chat_template( |
| 55 | IMAGE_PLACEHOLDER, |
| 56 | "What is the biggest text's content in this image?", |
| 57 | ), |
| 58 | "cherry_blossom": qwen2_vl_chat_template( |
| 59 | IMAGE_PLACEHOLDER, |
| 60 | "What is the season shown in this image? ", |
| 61 | "Reply with a short sentence (no more than 20 words)", |
| 62 | ), |
| 63 | } |
| 64 | ) |
| 65 | |
| 66 | VIDEO_PROMPTS = VIDEO_ASSETS.prompts( |
| 67 | { |
| 68 | "baby_reading": qwen2_vl_chat_template( |
| 69 | VIDEO_PLACEHOLDER, |
| 70 | "Describe this video with a short sentence ", |
| 71 | "(no more than 20 words)", |
| 72 | ), |
| 73 | } |
| 74 | ) |
| 75 | |
| 76 | MULTIIMAGE_PROMPT = qwen2_vl_chat_template( |
| 77 | IMAGE_PLACEHOLDER, |
| 78 | IMAGE_PLACEHOLDER, |
| 79 | "Describe these two images separately. ", |
| 80 | "For each image, reply with a short sentence ", |
| 81 | "(no more than 10 words).", |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | class Qwen2VLPromptImageEmbeddingInput(TypedDict): |
| 86 | image_embeds: torch.Tensor |
| 87 | image_grid_thw: torch.Tensor |
| 88 | |
| 89 | |
| 90 | class Qwen2VLPromptVideoEmbeddingInput(TypedDict): |
| 91 | video_embeds: torch.Tensor |
| 92 | video_grid_thw: torch.Tensor |
| 93 | |
| 94 | |
| 95 | def batch_make_image_embeddings( |
| 96 | image_batches: list[Image.Image | list[Image.Image]], |
| 97 | processor, |
| 98 | llm: VllmRunner, |
| 99 | ) -> list[Qwen2VLPromptImageEmbeddingInput]: |
| 100 | """Batched image embeddings for Qwen2-VL. |
| 101 | |
| 102 | This will infer all images' embeddings in a single batch, |
| 103 | and split the result according to input batches. |
| 104 | |
| 105 | image_batches: |
| 106 | - Single-image batches: `list[Image.Image]` |
| 107 | - Multiple-image batches: `list[list[Image.Image]]]` |
| 108 | |
| 109 | returns: `list[Qwen2VLPromptImageEmbeddingInput]` |
| 110 | """ |
| 111 | image_batches_: list[Any] = image_batches[:] |
| 112 | |
| 113 | # convert single-image batches to multiple-image batches |
| 114 | for idx in range(len(image_batches_)): |
| 115 | if not isinstance(image_batches_[idx], list): |
| 116 | image_batches_[idx] = [image_batches_[idx]] |
| 117 | |
| 118 | assert isinstance(image_batches_[idx], list) |
| 119 | |
| 120 | # append all images into a list (as a batch) |
| 121 | images: list[Image.Image] = [] |
| 122 | for image_batch in image_batches_: |
| 123 | images += image_batch |
| 124 | |
| 125 | # image to pixel values |
| 126 | image_processor = processor.image_processor |
| 127 | |
| 128 | mm_config = llm.get_llm().model_config.multimodal_config |
| 129 | mm_kwargs = mm_config.merge_mm_processor_kwargs({}) |
| 130 | preprocess_result = image_processor.preprocess( |
| 131 | images=images, return_tensors="pt", **mm_kwargs |
| 132 | ).data |
| 133 | pixel_values = preprocess_result["pixel_values"] |
| 134 | image_grid_thw = preprocess_result["image_grid_thw"] |
| 135 | |
| 136 | # pixel values to embeddings & grid_thws |
| 137 | def get_image_embeds(model): |
| 138 | with torch.no_grad(): |
| 139 | visual = model.visual |
| 140 | |
| 141 | pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype) |
| 142 | return visual(pixel_values_on_device, grid_thw=image_grid_thw).cpu() |
| 143 | |
| 144 | image_embeds = torch.concat(llm.apply_model(get_image_embeds)) |
| 145 | |
| 146 | # split into original batches |
| 147 | result: list[Qwen2VLPromptImageEmbeddingInput] = [] |
| 148 | image_counter = 0 |
| 149 | embed_counter = 0 |
| 150 | for image_batch in image_batches_: |
| 151 | cur_batch_image_count = len(image_batch) |
| 152 | merge_size = image_processor.merge_size |
| 153 | cur_batch_embed_len = sum( |
| 154 | grid_thw.prod(-1) // merge_size // merge_size |
| 155 | for grid_thw in image_grid_thw[ |
| 156 | image_counter : image_counter + cur_batch_image_count |
| 157 | ] |
| 158 | ) |
| 159 | |
| 160 | result.append( |
| 161 | { |
| 162 | "image_embeds": image_embeds[ |
| 163 | embed_counter : embed_counter + cur_batch_embed_len |
| 164 | ], |
| 165 | "image_grid_thw": image_grid_thw[ |
| 166 | image_counter : image_counter + cur_batch_image_count |
| 167 | ], |
| 168 | } |
| 169 | ) |
| 170 | |
| 171 | embed_counter += cur_batch_embed_len |
| 172 | image_counter += cur_batch_image_count |
| 173 | |
| 174 | # ensure we don't lose any images or embeddings |
| 175 | assert embed_counter == image_embeds.size(0) |
| 176 | assert image_counter == image_grid_thw.size(0) |
| 177 | assert len(image_batches) == len(result) |
| 178 | |
| 179 | return result |
| 180 | |
| 181 | |
| 182 | def batch_make_video_embeddings( |
| 183 | video_batches: PromptVideoInput, processor, llm: VllmRunner |
| 184 | ) -> list[Qwen2VLPromptVideoEmbeddingInput]: |
| 185 | """Batched video embeddings for Qwen2-VL. |
| 186 | |
| 187 | A NDArray represents a single video's all frames. |
| 188 | |
| 189 | This will infer all videos' embeddings in a single batch, |
| 190 | and split the result according to input batches. |
| 191 | |
| 192 | video_batches: |
| 193 | - Single-video batches: `list[NDArray]` |
| 194 | - Multiple-video batches: `list[list[NDArray]]` |
| 195 | """ |
| 196 | video_batches_: list[Any] = video_batches[:] |
| 197 | |
| 198 | for idx in range(len(video_batches_)): |
| 199 | if not isinstance(video_batches_[idx], list): |
| 200 | single_video_batch: list[npt.NDArray] = [video_batches_[idx]] |
| 201 | video_batches_[idx] = single_video_batch |
| 202 | |
| 203 | assert isinstance(video_batches_[idx], list) |
| 204 | |
| 205 | # append all videos into a list (as a batch) |
| 206 | videos: list[npt.NDArray] = [] |
| 207 | for video_batch in video_batches_: |
| 208 | videos += video_batch |
| 209 | |
| 210 | # video to pixel values |
| 211 | video_processor = processor.video_processor |
| 212 | |
| 213 | mm_config = llm.get_llm().model_config.multimodal_config |
| 214 | mm_kwargs = mm_config.merge_mm_processor_kwargs({}) |
| 215 | preprocess_result = video_processor.preprocess( |
| 216 | videos=videos, return_tensors="pt", **mm_kwargs |
| 217 | ).data |
| 218 | pixel_values = preprocess_result["pixel_values_videos"] |
| 219 | video_grid_thw = preprocess_result["video_grid_thw"] |
| 220 | |
| 221 | # pixel values to embeddings & grid_thws |
| 222 | def get_image_embeds(model): |
| 223 | with torch.no_grad(): |
| 224 | visual = model.visual |
| 225 | |
| 226 | pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype) |
| 227 | return visual(pixel_values_on_device, grid_thw=video_grid_thw).cpu() |
| 228 | |
| 229 | video_embeds = torch.concat(llm.apply_model(get_image_embeds)) |
| 230 | |
| 231 | # split into original batches |
| 232 | result: list[Qwen2VLPromptVideoEmbeddingInput] = [] |
| 233 | video_counter = 0 |
| 234 | embed_counter = 0 |
| 235 | for video_batch in video_batches_: |
| 236 | cur_batch_video_count = len(video_batch) |
| 237 | merge_size = video_processor.merge_size |
| 238 | cur_batch_embed_len = sum( |
| 239 | grid_thw.prod(-1) // merge_size // merge_size |
| 240 | for grid_thw in video_grid_thw[ |
| 241 | video_counter : video_counter + cur_batch_video_count |
| 242 | ] |
| 243 | ) |
| 244 | |
| 245 | result.append( |
| 246 | { |
| 247 | "video_embeds": video_embeds[ |
| 248 | embed_counter : embed_counter + cur_batch_embed_len |
| 249 | ], |
| 250 | "video_grid_thw": video_grid_thw[ |
| 251 | video_counter : video_counter + cur_batch_video_count |
| 252 | ], |
| 253 | } |
| 254 | ) |
| 255 | |
| 256 | embed_counter += cur_batch_embed_len |
| 257 | video_counter += cur_batch_video_count |
| 258 | |
| 259 | # ensure we don't lose any videos or embeddings |
| 260 | assert embed_counter == video_embeds.size(0) |
| 261 | assert video_counter == video_grid_thw.size(0) |
| 262 | assert len(video_batches) == len(result) |
| 263 | |
| 264 | return result |
| 265 | |
| 266 | |
| 267 | def run_embedding_input_test( |
| 268 | vllm_runner: type[VllmRunner], |
| 269 | inputs: list[tuple[list[str], PromptImageInput, PromptVideoInput]], |
| 270 | model: str, |
| 271 | *, |
| 272 | dtype: str, |
| 273 | max_tokens: int, |
| 274 | num_logprobs: int, |
| 275 | mm_limit: int, |
| 276 | tensor_parallel_size: int, |
| 277 | distributed_executor_backend: str | None = None, |
| 278 | ): |
| 279 | """Inference result should be the same between |
| 280 | original image/video input and image/video embeddings input. |
| 281 | """ |
| 282 | from transformers import AutoProcessor |
| 283 | |
| 284 | processor = AutoProcessor.from_pretrained(model) |
| 285 | |
| 286 | # max_model_len should be greater than image_feature_size |
| 287 | with vllm_runner( |
| 288 | model, |
| 289 | runner="generate", |
| 290 | max_model_len=4000, |
| 291 | max_num_seqs=3, |
| 292 | dtype=dtype, |
| 293 | limit_mm_per_prompt={"image": mm_limit, "video": mm_limit}, |
| 294 | tensor_parallel_size=tensor_parallel_size, |
| 295 | distributed_executor_backend=distributed_executor_backend, |
| 296 | default_torch_num_threads=1, |
| 297 | enable_mm_embeds=True, |
| 298 | ) as vllm_model: |
| 299 | outputs_per_case_for_original_input = [ |
| 300 | vllm_model.generate_greedy_logprobs( |
| 301 | prompts, |
| 302 | max_tokens, |
| 303 | num_logprobs=num_logprobs, |
| 304 | images=images or None, |
| 305 | videos=videos or None, |
| 306 | ) |
| 307 | for prompts, images, videos in inputs |
| 308 | ] |
| 309 | |
| 310 | outputs_per_case_for_embeddings_input = [ |
| 311 | vllm_model.generate_greedy_logprobs( |
| 312 | prompts, |
| 313 | max_tokens, |
| 314 | num_logprobs=num_logprobs, |
| 315 | images=batch_make_image_embeddings(images, processor, vllm_model) |
| 316 | if images |
| 317 | else None, |
| 318 | videos=batch_make_video_embeddings(videos, processor, vllm_model) |
| 319 | if videos |
| 320 | else None, |
| 321 | ) |
| 322 | for prompts, images, videos in inputs |
| 323 | ] |
| 324 | |
| 325 | for outputs_for_original_input, outputs_for_embeddings_input in zip( |
| 326 | outputs_per_case_for_original_input, outputs_per_case_for_embeddings_input |
| 327 | ): |
| 328 | check_logprobs_close( |
| 329 | outputs_0_lst=outputs_for_original_input, |
| 330 | outputs_1_lst=outputs_for_embeddings_input, |
| 331 | name_0="original_input", |
| 332 | name_1="embeddings_input", |
| 333 | ) |
| 334 | |
| 335 | |
| 336 | @pytest.mark.core_model |
| 337 | @pytest.mark.parametrize("model", models) |
| 338 | @pytest.mark.parametrize("dtype", [target_dtype]) |
| 339 | @pytest.mark.parametrize("max_tokens", [128]) |
| 340 | @pytest.mark.parametrize("num_logprobs", [10]) |
| 341 | def test_qwen2_vl_image_embeddings_input( |
| 342 | vllm_runner, |
| 343 | image_assets, |
| 344 | model, |
| 345 | dtype, |
| 346 | max_tokens, |
| 347 | num_logprobs, |
| 348 | monkeypatch, |
| 349 | ) -> None: |
| 350 | images = [asset.pil_image for asset in image_assets] |
| 351 | |
| 352 | inputs_per_case: list[tuple[list[str], PromptImageInput, PromptVideoInput]] = [ |
| 353 | ( |
| 354 | [prompt for _ in size_factors], |
| 355 | [rescale_image_size(image, factor) for factor in size_factors], |
| 356 | [], |
| 357 | ) |
| 358 | for size_factors in IMAGE_SIZE_FACTOR_GROUPS |
| 359 | for image, prompt in zip(images, IMAGE_PROMPTS) |
| 360 | ] |
| 361 | |
| 362 | run_embedding_input_test( |
| 363 | vllm_runner, |
| 364 | inputs_per_case, |
| 365 | model, |
| 366 | dtype=dtype, |
| 367 | max_tokens=max_tokens, |
| 368 | num_logprobs=num_logprobs, |
| 369 | mm_limit=1, |
| 370 | tensor_parallel_size=1, |
| 371 | ) |
| 372 | |
| 373 | |
| 374 | @pytest.mark.core_model |
| 375 | @pytest.mark.parametrize("model", models) |
| 376 | @pytest.mark.parametrize("dtype", [target_dtype]) |
| 377 | @pytest.mark.parametrize("max_tokens", [128]) |
| 378 | @pytest.mark.parametrize("num_logprobs", [10]) |
| 379 | def test_qwen2_vl_multiple_image_embeddings_input( |
| 380 | vllm_runner, |
| 381 | image_assets, |
| 382 | model, |
| 383 | dtype: str, |
| 384 | max_tokens: int, |
| 385 | num_logprobs: int, |
| 386 | ) -> None: |
| 387 | images = [asset.pil_image for asset in image_assets] |
| 388 | |
| 389 | inputs_per_case: list[tuple[list[str], PromptImageInput, PromptVideoInput]] = [ |
| 390 | ( |
| 391 | [MULTIIMAGE_PROMPT for _ in size_factors], |
| 392 | [ |
| 393 | [rescale_image_size(image, factor) for image in images] |
| 394 | for factor in size_factors |
| 395 | ], |
| 396 | [], |
| 397 | ) |
| 398 | for size_factors in IMAGE_SIZE_FACTOR_GROUPS |
| 399 | ] |
| 400 | |
| 401 | run_embedding_input_test( |
| 402 | vllm_runner, |
| 403 | inputs_per_case, |
| 404 | model, |
| 405 | dtype=dtype, |
| 406 | max_tokens=max_tokens, |
| 407 | num_logprobs=num_logprobs, |
| 408 | mm_limit=2, |
| 409 | tensor_parallel_size=1, |
| 410 | ) |
| 411 | |
| 412 | |
| 413 | @pytest.mark.core_model |
| 414 | @pytest.mark.parametrize("model", models) |
| 415 | @pytest.mark.parametrize("dtype", [target_dtype]) |
| 416 | @pytest.mark.parametrize("max_tokens", [128]) |
| 417 | @pytest.mark.parametrize("num_logprobs", [10]) |
| 418 | def test_qwen2_vl_video_embeddings_input( |
| 419 | vllm_runner, |
| 420 | video_assets, |
| 421 | model, |
| 422 | dtype: str, |
| 423 | max_tokens: int, |
| 424 | num_logprobs: int, |
| 425 | ) -> None: |
| 426 | num_frames = 4 |
| 427 | sampled_vids = [ |
| 428 | sample_frames_from_video(asset.np_ndarrays, num_frames) |
| 429 | for asset in video_assets |
| 430 | ] |
| 431 | |
| 432 | inputs_per_case: list[tuple[list[str], PromptImageInput, PromptVideoInput]] = [ |
| 433 | ( |
| 434 | [prompt for _ in size_factors], |
| 435 | [], |
| 436 | [rescale_video_size(video, factor) for factor in size_factors], |
| 437 | ) |
| 438 | for size_factors in VIDEO_SIZE_FACTOR_GROUPS |
| 439 | for video, prompt in zip(sampled_vids, VIDEO_PROMPTS) |
| 440 | ] |
| 441 | |
| 442 | run_embedding_input_test( |
| 443 | vllm_runner, |
| 444 | inputs_per_case, |
| 445 | model, |
| 446 | dtype=dtype, |
| 447 | max_tokens=max_tokens, |
| 448 | num_logprobs=num_logprobs, |
| 449 | mm_limit=1, |
| 450 | tensor_parallel_size=1, |
| 451 | ) |