main
py 258 lines 8.06 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5 from PIL import Image
6
7 from vllm.multimodal.inputs import (
8 MultiModalBatchedField,
9 MultiModalFieldElem,
10 MultiModalKwargsItem,
11 MultiModalSharedField,
12 PlaceholderRange,
13 )
14 from vllm.multimodal.utils import (
15 argsort_mm_positions,
16 encode_image_url,
17 group_and_batch_mm_items,
18 )
19
20
21 @pytest.mark.parametrize(
22 "case",
23 [
24 # Single modality
25 ## Internally sorted
26 dict(
27 mm_positions={
28 "image": [
29 PlaceholderRange(offset=0, length=2),
30 PlaceholderRange(offset=3, length=2),
31 ]
32 },
33 expected_modality_idxs=[
34 ("image", 0),
35 ("image", 1),
36 ],
37 ),
38 ## Internally unsorted
39 dict(
40 mm_positions={
41 "image": [
42 PlaceholderRange(offset=3, length=2),
43 PlaceholderRange(offset=0, length=2),
44 ]
45 },
46 expected_modality_idxs=[
47 ("image", 1),
48 ("image", 0),
49 ],
50 ),
51 # Two modalities
52 ## Internally sorted
53 dict(
54 mm_positions={
55 "image": [
56 PlaceholderRange(offset=7, length=4),
57 PlaceholderRange(offset=11, length=5),
58 ],
59 "audio": [
60 PlaceholderRange(offset=0, length=2),
61 PlaceholderRange(offset=2, length=3),
62 ],
63 },
64 expected_modality_idxs=[
65 ("audio", 0),
66 ("audio", 1),
67 ("image", 0),
68 ("image", 1),
69 ],
70 ),
71 ## Interleaved, internally sorted
72 dict(
73 mm_positions={
74 "image": [
75 PlaceholderRange(offset=0, length=4),
76 PlaceholderRange(offset=8, length=2),
77 ],
78 "audio": [
79 PlaceholderRange(offset=5, length=2),
80 PlaceholderRange(offset=11, length=4),
81 ],
82 },
83 expected_modality_idxs=[
84 ("image", 0),
85 ("audio", 0),
86 ("image", 1),
87 ("audio", 1),
88 ],
89 ),
90 ## Interleaved, internally unsorted
91 dict(
92 mm_positions={
93 "image": [
94 PlaceholderRange(offset=8, length=2),
95 PlaceholderRange(offset=0, length=4),
96 ],
97 "audio": [
98 PlaceholderRange(offset=11, length=4),
99 PlaceholderRange(offset=5, length=2),
100 ],
101 },
102 expected_modality_idxs=[
103 ("image", 1),
104 ("audio", 1),
105 ("image", 0),
106 ("audio", 0),
107 ],
108 ),
109 # Three modalities
110 ## Internally sorted
111 dict(
112 mm_positions={
113 "image": [
114 PlaceholderRange(offset=15, length=7),
115 PlaceholderRange(offset=22, length=8),
116 ],
117 "audio": [
118 PlaceholderRange(offset=0, length=2),
119 ],
120 "video": [
121 PlaceholderRange(offset=3, length=4),
122 PlaceholderRange(offset=7, length=5),
123 PlaceholderRange(offset=12, length=6),
124 ],
125 },
126 expected_modality_idxs=[
127 ("audio", 0),
128 ("video", 0),
129 ("video", 1),
130 ("video", 2),
131 ("image", 0),
132 ("image", 1),
133 ],
134 ),
135 ## Interleaved, internally sorted
136 dict(
137 mm_positions={
138 "image": [
139 PlaceholderRange(offset=0, length=2),
140 PlaceholderRange(offset=2, length=3),
141 PlaceholderRange(offset=20, length=4),
142 ],
143 "audio": [
144 PlaceholderRange(offset=5, length=2),
145 ],
146 "video": [
147 PlaceholderRange(offset=8, length=5),
148 ],
149 },
150 expected_modality_idxs=[
151 ("image", 0),
152 ("image", 1),
153 ("audio", 0),
154 ("video", 0),
155 ("image", 2),
156 ],
157 ),
158 ## Interleaved, internally unsorted
159 dict(
160 mm_positions={
161 "image": [
162 PlaceholderRange(offset=0, length=2),
163 PlaceholderRange(offset=20, length=4),
164 PlaceholderRange(offset=2, length=3),
165 ],
166 "audio": [
167 PlaceholderRange(offset=5, length=2),
168 ],
169 "video": [
170 PlaceholderRange(offset=8, length=5),
171 ],
172 },
173 expected_modality_idxs=[
174 ("image", 0),
175 ("image", 2),
176 ("audio", 0),
177 ("video", 0),
178 ("image", 1),
179 ],
180 ),
181 ],
182 )
183 def test_argsort_mm_positions(case):
184 mm_positions = case["mm_positions"]
185 expected_modality_idxs = case["expected_modality_idxs"]
186
187 modality_idxs = argsort_mm_positions(mm_positions)
188
189 assert modality_idxs == expected_modality_idxs
190
191
192 def test_group_and_batch_mm_items_split_by_fieldset():
193 elem = MultiModalFieldElem(
194 data=torch.empty(1, dtype=torch.uint8),
195 field=MultiModalBatchedField(),
196 )
197 item1 = MultiModalKwargsItem({"x": elem, "y": elem})
198 item2 = MultiModalKwargsItem({"y": elem, "x": elem})
199 item3 = MultiModalKwargsItem({"x": elem, "y": elem, "z": elem})
200 item4 = MultiModalKwargsItem({"x": elem})
201 item5 = MultiModalKwargsItem({"x": elem, "y": elem})
202
203 res = group_and_batch_mm_items([item1, item2, item3, item4, item5])
204 assert [num_items for num_items, _ in res] == [2, 1, 1, 1]
205
206
207 def test_group_and_batch_mm_items_splits_shared_data_by_dtype():
208 elem1 = MultiModalFieldElem(
209 data=torch.zeros(1, dtype=torch.int32),
210 field=MultiModalSharedField(batch_size=1),
211 )
212 elem2 = MultiModalFieldElem(
213 data=torch.zeros(1, dtype=torch.float32),
214 field=MultiModalSharedField(batch_size=1),
215 )
216 elem3 = MultiModalFieldElem(
217 data=[torch.zeros(1, dtype=torch.int32), torch.zeros(1, dtype=torch.float32)],
218 field=MultiModalSharedField(batch_size=1),
219 )
220
221 res = group_and_batch_mm_items(
222 [
223 MultiModalKwargsItem({"x": elem1}),
224 MultiModalKwargsItem({"x": elem2}),
225 MultiModalKwargsItem({"x": elem3}),
226 ]
227 )
228
229 assert [num_items for num_items, _ in res] == [1, 1, 1]
230
231
232 def test_encode_image_url_uncommon_format_has_valid_mimetype():
233 """A format missing from mimetypes.types_map must still produce a
234 well-formed type/subtype media type, not a bare 'image'."""
235 # ".tga" is not registered in mimetypes.types_map, so it exercises the
236 # fallback path; Pillow can still encode it.
237 url = encode_image_url(Image.new("RGB", (2, 2)), format="TGA")
238 mimetype = url.removeprefix("data:").split(";", 1)[0]
239 assert mimetype == "image/tga"
240
241
242 def test_group_and_batch_mm_items_split_by_shared_data():
243 elem1 = MultiModalFieldElem(
244 data=torch.zeros(1, dtype=torch.uint8),
245 field=MultiModalSharedField(batch_size=1),
246 )
247 elem2 = MultiModalFieldElem(
248 data=torch.zeros(2, dtype=torch.uint8),
249 field=MultiModalSharedField(batch_size=1),
250 )
251 item1 = MultiModalKwargsItem({"x": elem1})
252 item2 = MultiModalKwargsItem({"x": elem1})
253 item3 = MultiModalKwargsItem({"x": elem2})
254 item4 = MultiModalKwargsItem({"x": elem1})
255 item5 = MultiModalKwargsItem({"x": elem2})
256
257 res = group_and_batch_mm_items([item1, item2, item3, item4, item5])
258 assert [num_items for num_items, _ in res] == [2, 1, 1, 1]