main
py 160 lines 5.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import importlib.util
5
6 import numpy as np
7 import pybase64 as base64
8 import pytest
9 import requests
10 import torch
11
12 from tests.utils import RemoteOpenAIServer
13 from vllm.utils.serial_utils import tensor2base64
14
15 # Prithvi requires terratorch, which is temporarily unavailable while PyPI has
16 # `lightning` quarantined (#41376). Skip just the Prithvi case; leave the
17 # Qwen3-VL case in the same file untouched.
18 _TERRATORCH_AVAILABLE = importlib.util.find_spec("terratorch") is not None
19
20
21 @pytest.mark.skipif(
22 not _TERRATORCH_AVAILABLE,
23 reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
24 )
25 @pytest.mark.parametrize(
26 "model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
27 )
28 def test_single_content(model_name: str):
29 args = [
30 "--runner",
31 "pooling",
32 # use half precision for speed and memory savings in CI environment
33 "--dtype",
34 "float16",
35 "--enforce-eager",
36 "--trust-remote-code",
37 "--max-num-seqs",
38 "32",
39 "--model-impl",
40 "terratorch",
41 "--skip-tokenizer-init",
42 "--enable-mm-embeds",
43 ]
44
45 with RemoteOpenAIServer(model_name, args) as server:
46 response = requests.post(
47 server.url_for("pooling"),
48 json={
49 "model": model_name,
50 "messages": [
51 {
52 "role": "user",
53 "content": [
54 {
55 "type": "image_embeds",
56 "image_embeds": {
57 "pixel_values": tensor2base64(
58 torch.ones((6, 512, 512), dtype=torch.float16)
59 ),
60 "location_coords": tensor2base64(
61 torch.ones((1, 2), dtype=torch.float16)
62 ),
63 },
64 },
65 ],
66 }
67 ],
68 "encoding_format": "base64",
69 },
70 )
71 response.raise_for_status()
72
73 output = response.json()["data"][0]["data"]
74
75 np_response = np.frombuffer(base64.b64decode(output), dtype=np.float32)
76 assert len(np_response) == 524288
77
78
79 @pytest.mark.parametrize("model_name", ["Qwen/Qwen3-VL-2B-Instruct"])
80 def test_multi_content(model_name: str):
81 args = [
82 "--enforce-eager",
83 "--max-num-seqs",
84 "32",
85 "--max-model-len",
86 "8192",
87 "--enable-mm-embeds",
88 ]
89
90 with RemoteOpenAIServer(model_name, args) as server:
91 client = server.get_client()
92
93 # Image only
94 chat_completion = client.chat.completions.create(
95 model=model_name,
96 messages=[
97 {
98 "role": "user",
99 "content": [
100 {
101 "type": "image_embeds",
102 "image_embeds": {
103 "image_embeds": tensor2base64(torch.zeros(220, 8192)),
104 "image_grid_thw": tensor2base64(
105 torch.tensor([1, 22, 40])
106 ),
107 },
108 },
109 {
110 "type": "image_embeds",
111 "image_embeds": {
112 "image_embeds": tensor2base64(torch.zeros(220, 8192)),
113 "image_grid_thw": tensor2base64(
114 torch.tensor([1, 22, 40])
115 ),
116 },
117 },
118 ],
119 }
120 ],
121 max_tokens=5,
122 )
123
124 assert chat_completion.id is not None
125 assert len(chat_completion.choices) == 1
126
127 # Interleaved text and image
128 chat_completion = client.chat.completions.create(
129 model=model_name,
130 messages=[
131 {
132 "role": "user",
133 "content": [
134 {
135 "type": "image_embeds",
136 "image_embeds": {
137 "image_embeds": tensor2base64(torch.zeros(220, 8192)),
138 "image_grid_thw": tensor2base64(
139 torch.tensor([1, 22, 40])
140 ),
141 },
142 },
143 {"type": "text", "text": "OCR:"},
144 {
145 "type": "image_embeds",
146 "image_embeds": {
147 "image_embeds": tensor2base64(torch.zeros(220, 8192)),
148 "image_grid_thw": tensor2base64(
149 torch.tensor([1, 22, 40])
150 ),
151 },
152 },
153 ],
154 }
155 ],
156 max_tokens=5,
157 )
158
159 assert chat_completion.id is not None
160 assert len(chat_completion.choices) == 1