| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | from transformers import CLIPModel |
| 7 | |
| 8 | from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner |
| 9 | from ...utils import check_embeddings_close |
| 10 | |
| 11 | HF_TEXT_PROMPTS = [ |
| 12 | "a photo of a stop sign", |
| 13 | "a photo of a cherry blossom", |
| 14 | ] |
| 15 | |
| 16 | HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( |
| 17 | { |
| 18 | "stop_sign": "", |
| 19 | "cherry_blossom": "", |
| 20 | } |
| 21 | ) |
| 22 | |
| 23 | MODELS = ["openai/clip-vit-base-patch32"] |
| 24 | |
| 25 | |
| 26 | def _run_test( |
| 27 | hf_runner: type[HfRunner], |
| 28 | vllm_runner: type[VllmRunner], |
| 29 | input_cases: list[tuple[list[str], PromptImageInput]], |
| 30 | model: str, |
| 31 | *, |
| 32 | dtype: str, |
| 33 | ) -> None: |
| 34 | # NOTE: take care of the order. run vLLM first, and then run HF. |
| 35 | # vLLM needs a fresh new process without cuda initialization. |
| 36 | # if we run HF first, the cuda initialization will be done and it |
| 37 | # will hurt multiprocessing backend with fork method (the default method). |
| 38 | with vllm_runner( |
| 39 | model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77 |
| 40 | ) as vllm_model: |
| 41 | vllm_outputs_per_case = [ |
| 42 | vllm_model.embed(input_texts, images=input_images) |
| 43 | for input_texts, input_images in input_cases |
| 44 | ] |
| 45 | |
| 46 | texts = [HF_TEXT_PROMPTS[0]] |
| 47 | images = [input_cases[1][1][0]] |
| 48 | with pytest.raises(ValueError, match="not both"): |
| 49 | vllm_model.embed(texts, images=images) |
| 50 | |
| 51 | # Should still be able to run subsequent requests |
| 52 | vllm_model.embed(texts) |
| 53 | vllm_model.embed([""], images=images) |
| 54 | |
| 55 | # Mixed image+text batch must not skip the text encoder (#53091). |
| 56 | mixed_outputs = vllm_outputs_per_case[2] |
| 57 | check_embeddings_close( |
| 58 | embeddings_0_lst=[vllm_outputs_per_case[0][0]], |
| 59 | embeddings_1_lst=[mixed_outputs[0]], |
| 60 | name_0="text_only", |
| 61 | name_1="mixed_text", |
| 62 | ) |
| 63 | check_embeddings_close( |
| 64 | embeddings_0_lst=[vllm_outputs_per_case[1][0]], |
| 65 | embeddings_1_lst=[mixed_outputs[1]], |
| 66 | name_0="image_only", |
| 67 | name_1="mixed_image", |
| 68 | ) |
| 69 | empty_text = vllm_model.embed([""]) |
| 70 | empty_sim = torch.nn.functional.cosine_similarity( |
| 71 | torch.tensor(mixed_outputs[0]), |
| 72 | torch.tensor(empty_text[0]), |
| 73 | dim=0, |
| 74 | ) |
| 75 | assert empty_sim < 0.99, ( |
| 76 | "Mixed-batch text embedding collapsed to the empty-string vector " |
| 77 | f"(cosine={empty_sim:.4f})" |
| 78 | ) |
| 79 | |
| 80 | with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model: |
| 81 | hf_outputs_per_case = [] |
| 82 | for input_texts, input_images in input_cases: |
| 83 | all_inputs = hf_model.get_inputs(input_texts, images=input_images) |
| 84 | |
| 85 | hf_outputs = [] |
| 86 | for inputs in all_inputs: |
| 87 | inputs = hf_model.wrap_device(inputs) |
| 88 | |
| 89 | if "pixel_values" in inputs: |
| 90 | pooled_output = hf_model.model.get_image_features( |
| 91 | pixel_values=inputs.pixel_values, |
| 92 | ) |
| 93 | else: |
| 94 | pooled_output = hf_model.model.get_text_features( |
| 95 | input_ids=inputs.input_ids, |
| 96 | attention_mask=inputs.attention_mask, |
| 97 | ) |
| 98 | |
| 99 | if not isinstance(pooled_output, torch.Tensor): |
| 100 | pooled_output = pooled_output.pooler_output |
| 101 | pooled_output = pooled_output.squeeze(0) |
| 102 | hf_outputs.append(pooled_output.tolist()) |
| 103 | |
| 104 | hf_outputs_per_case.append(hf_outputs) |
| 105 | |
| 106 | for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): |
| 107 | check_embeddings_close( |
| 108 | embeddings_0_lst=hf_outputs, |
| 109 | embeddings_1_lst=vllm_outputs, |
| 110 | name_0="hf", |
| 111 | name_1="vllm", |
| 112 | ) |
| 113 | |
| 114 | |
| 115 | @pytest.mark.parametrize("model", MODELS) |
| 116 | @pytest.mark.parametrize("dtype", ["float"]) |
| 117 | def test_models( |
| 118 | hf_runner, |
| 119 | vllm_runner, |
| 120 | image_assets, |
| 121 | model: str, |
| 122 | dtype: str, |
| 123 | ) -> None: |
| 124 | text_images = [None] * len(HF_TEXT_PROMPTS) |
| 125 | images = [asset.pil_image for asset in image_assets] |
| 126 | input_cases = [ |
| 127 | (HF_TEXT_PROMPTS, text_images), |
| 128 | (HF_IMAGE_PROMPTS, images), |
| 129 | ([HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]], [None, images[0]]), |
| 130 | ] |
| 131 | |
| 132 | _run_test( |
| 133 | hf_runner, |
| 134 | vllm_runner, |
| 135 | input_cases, |
| 136 | model, |
| 137 | dtype=dtype, |
| 138 | ) |