| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | import torch.nn.functional as F |
| 7 | import transformers.utils |
| 8 | from PIL import Image |
| 9 | |
| 10 | from vllm.assets.base import get_vllm_public_assets |
| 11 | from vllm.assets.image import VLM_IMAGES_DIR |
| 12 | from vllm.config import ModelConfig |
| 13 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner |
| 17 | from ....utils import large_gpu_test |
| 18 | from ...utils import check_embeddings_close |
| 19 | |
| 20 | # BC for method that was deleted in Transformers v5. |
| 21 | # Only needed for generating the HF reference. |
| 22 | transformers.utils.is_flash_attn_greater_or_equal_2_10 = ( |
| 23 | lambda: transformers.utils.is_flash_attn_greater_or_equal("2.1.0") |
| 24 | ) |
| 25 | |
| 26 | HF_TEXT_PROMPTS = [ |
| 27 | # T -> X |
| 28 | "Find me an everyday image that matches the given caption: The label of the object is stop sign", # noqa: E501 |
| 29 | # T -> X |
| 30 | "Retrieve an image of this caption: cherry blossom", |
| 31 | ] |
| 32 | |
| 33 | HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( |
| 34 | { |
| 35 | # T + I -> X |
| 36 | "stop_sign": "<|image_1|> Select the portion of the image that isolates the object of the given label: The label of the object is stop sign", # noqa: E501 |
| 37 | # I -> X |
| 38 | "cherry_blossom": "<|image_1|> Represent the given image for classification", # noqa: E501 |
| 39 | } |
| 40 | ) |
| 41 | |
| 42 | MODELS = ["TIGER-Lab/VLM2Vec-Full"] |
| 43 | |
| 44 | SPECIAL_TOKEN_IMAGE_PROMPT = ( |
| 45 | "\n<s><|user|>\n <|image_1|>\n\t <s>" |
| 46 | "Represent the given image for classification<|end|>" |
| 47 | "\n<|assistant|>\n" |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def _get_cherry_blossom_image() -> Image.Image: |
| 52 | return Image.open( |
| 53 | get_vllm_public_assets(filename="cherry_blossom.jpg", s3_prefix=VLM_IMAGES_DIR) |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | @torch.inference_mode() |
| 58 | def _run_test( |
| 59 | hf_runner: type[HfRunner], |
| 60 | vllm_runner: type[VllmRunner], |
| 61 | input_texts: list[str], |
| 62 | input_images: PromptImageInput, |
| 63 | model: str, |
| 64 | *, |
| 65 | dtype: str, |
| 66 | ) -> None: |
| 67 | # NOTE: take care of the order. run vLLM first, and then run HF. |
| 68 | # vLLM needs a fresh new process without cuda initialization. |
| 69 | # if we run HF first, the cuda initialization will be done and it |
| 70 | # will hurt multiprocessing backend with fork method (the default method). |
| 71 | with vllm_runner( |
| 72 | model, runner="pooling", dtype=dtype, enforce_eager=True |
| 73 | ) as vllm_model: |
| 74 | vllm_outputs = vllm_model.embed(input_texts, images=input_images) |
| 75 | |
| 76 | # use eager mode for hf runner, since phi3_v didn't work with flash_attn |
| 77 | hf_model_kwargs = {"_attn_implementation": "eager"} |
| 78 | with hf_runner(model, dtype=dtype, model_kwargs=hf_model_kwargs) as hf_model: |
| 79 | all_inputs = hf_model.get_inputs(input_texts, images=input_images) |
| 80 | |
| 81 | all_outputs = [] |
| 82 | for inputs in all_inputs: |
| 83 | # Based on: https://github.com/TIGER-AI-Lab/VLM2Vec/blob/db3b951bccabba220c1f53ab46a734e50dd2fc08/src/model.py |
| 84 | outputs = hf_model.model( |
| 85 | **hf_model.wrap_device(inputs), |
| 86 | return_dict=True, |
| 87 | output_hidden_states=True, |
| 88 | ) |
| 89 | last_hidden_state = outputs.hidden_states[-1][0] |
| 90 | reps = last_hidden_state[inputs.attention_mask[0].sum() - 1] |
| 91 | pooled_output = F.normalize(reps, p=2, dim=-1) |
| 92 | |
| 93 | all_outputs.append(pooled_output.tolist()) |
| 94 | |
| 95 | hf_outputs = all_outputs |
| 96 | |
| 97 | check_embeddings_close( |
| 98 | embeddings_0_lst=hf_outputs, |
| 99 | embeddings_1_lst=vllm_outputs, |
| 100 | name_0="hf", |
| 101 | name_1="vllm", |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | @pytest.mark.core_model |
| 106 | @pytest.mark.skipif( |
| 107 | current_platform.is_xpu(), |
| 108 | reason=( |
| 109 | "vllm-xpu-kernels paged_decode_default.conf is missing the " |
| 110 | "8,96,64,false,true,false tuple needed by VLM2Vec-Full at " |
| 111 | "block_size=64. Re-enable once vllm-xpu-kernels ships the updated " |
| 112 | "config. See github.com/vllm-project/vllm-xpu-kernels/pull/587." |
| 113 | ), |
| 114 | ) |
| 115 | @pytest.mark.parametrize("model", MODELS) |
| 116 | @pytest.mark.parametrize("dtype", ["half"]) |
| 117 | def test_models_text( |
| 118 | hf_runner, |
| 119 | vllm_runner, |
| 120 | image_assets, |
| 121 | model: str, |
| 122 | dtype: str, |
| 123 | ) -> None: |
| 124 | input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] |
| 125 | input_texts = [text for text, _ in input_texts_images] |
| 126 | input_images = [image for _, image in input_texts_images] |
| 127 | |
| 128 | _run_test( |
| 129 | hf_runner, |
| 130 | vllm_runner, |
| 131 | input_texts, |
| 132 | input_images, # type: ignore |
| 133 | model, |
| 134 | dtype=dtype, |
| 135 | ) |
| 136 | |
| 137 | |
| 138 | @large_gpu_test(min_gb=48) |
| 139 | @pytest.mark.core_model |
| 140 | @pytest.mark.parametrize("model", MODELS) |
| 141 | @pytest.mark.parametrize("dtype", ["half"]) |
| 142 | def test_models_image( |
| 143 | hf_runner, |
| 144 | vllm_runner, |
| 145 | image_assets, |
| 146 | model: str, |
| 147 | dtype: str, |
| 148 | ) -> None: |
| 149 | input_texts_images = [ |
| 150 | (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) |
| 151 | ] |
| 152 | input_texts = [text for text, _ in input_texts_images] |
| 153 | input_images = [image for _, image in input_texts_images] |
| 154 | |
| 155 | _run_test( |
| 156 | hf_runner, |
| 157 | vllm_runner, |
| 158 | input_texts, |
| 159 | input_images, |
| 160 | model, |
| 161 | dtype=dtype, |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | @pytest.mark.core_model |
| 166 | @pytest.mark.parametrize("model", MODELS) |
| 167 | @pytest.mark.parametrize("dtype", ["half"]) |
| 168 | def test_models_image_special_tokens_processing( |
| 169 | model: str, |
| 170 | dtype: str, |
| 171 | ) -> None: |
| 172 | model_config = ModelConfig( |
| 173 | model, |
| 174 | runner="pooling", |
| 175 | trust_remote_code=True, |
| 176 | dtype=dtype, |
| 177 | max_model_len=1024, |
| 178 | ) |
| 179 | processor = MULTIMODAL_REGISTRY.create_processor(model_config) |
| 180 | image = _get_cherry_blossom_image() |
| 181 | |
| 182 | processed_inputs = processor( |
| 183 | SPECIAL_TOKEN_IMAGE_PROMPT, |
| 184 | mm_items=processor.info.parse_mm_data({"image": image}), |
| 185 | hf_processor_mm_kwargs={}, |
| 186 | ) |
| 187 | |
| 188 | hf_processor = processor.info.get_hf_processor() |
| 189 | hf_inputs = hf_processor( |
| 190 | SPECIAL_TOKEN_IMAGE_PROMPT, |
| 191 | images=image, |
| 192 | return_tensors="pt", |
| 193 | ) |
| 194 | |
| 195 | image_token_id = hf_processor.get_special_image_token_id() |
| 196 | hf_prompt_token_ids = [ |
| 197 | image_token_id if token_id < 0 else token_id |
| 198 | for token_id in hf_inputs["input_ids"][0].tolist() |
| 199 | ] |
| 200 | |
| 201 | prompt_token_ids = processed_inputs["prompt_token_ids"] |
| 202 | |
| 203 | assert prompt_token_ids == hf_prompt_token_ids |
| 204 | assert prompt_token_ids.count(image_token_id) == hf_prompt_token_ids.count( |
| 205 | image_token_id |
| 206 | ) |
| 207 | assert prompt_token_ids.count(image_token_id) > 0 |