@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 166 lines 5.55 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import numpy as np
4 import pytest
5 import torch
6 from PIL import Image
7
8 from vllm.multimodal.parse import (
9 AudioProcessorItems,
10 ImageProcessorItems,
11 MultiModalDataParser,
12 VideoProcessorItems,
13 )
14
15 H, W = 480, 640
16
17
18 class AudioMetadataParser(MultiModalDataParser):
19 embedding_fields = {
20 "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
21 }
22
23
24 @pytest.mark.parametrize("allow_missing", [False, True])
25 def test_audio_metadata_requires_ec_consumer(allow_missing):
26 parser = AudioMetadataParser(allow_missing_mm_embeddings=allow_missing)
27 data = {"audio_num_tokens": torch.tensor([[3], [5]])}
28 if not allow_missing:
29 with pytest.raises(ValueError, match="audio_embeds"):
30 parser.parse_mm_data({"audio": data})
31 else:
32 items = parser.parse_mm_data({"audio": data})["audio"]
33 assert len(items) == 2
34 assert items.get(1)["audio_num_tokens"].item() == 5
35 assert items.get_processor_data() == {}
36
37
38 @pytest.mark.parametrize("counts", [[0], [-1], [1.5], [True], [[1, 2]]])
39 def test_audio_metadata_rejects_invalid_token_counts(counts):
40 parser = AudioMetadataParser(allow_missing_mm_embeddings=True)
41 with pytest.raises(ValueError, match="positive integer"):
42 parser.parse_mm_data({"audio": {"audio_num_tokens": torch.tensor(counts)}})
43
44
45 def test_audio_metadata_checks_supplied_embedding_lengths():
46 parser = AudioMetadataParser()
47 data = {
48 "audio_num_tokens": torch.tensor([3, 5]),
49 "audio_embeds": [torch.zeros(3, 8), torch.zeros(5, 8)],
50 }
51 assert len(parser.parse_mm_data({"audio": data})["audio"]) == 2
52 data["audio_num_tokens"] = torch.tensor([3, 4])
53 with pytest.raises(ValueError, match="does not match"):
54 parser.parse_mm_data({"audio": data})
55
56
57 @pytest.mark.parametrize(
58 "image",
59 [
60 Image.new("RGB", (W, H)),
61 # HWC, e.g. from np.array(PIL.Image)
62 np.zeros((H, W, 3), dtype=np.uint8),
63 torch.zeros((H, W, 3), dtype=torch.uint8),
64 # CHW, standard PyTorch / numpy convention
65 np.zeros((3, H, W), dtype=np.uint8),
66 torch.zeros((3, H, W), dtype=torch.uint8),
67 ],
68 )
69 def test_image_size_hwc_chw(image):
70 """Image sizes must be channel-layout agnostic.
71
72 `get_image_size` determines the multimodal placeholder count; reading an
73 HWC array (the layout `np.array(PIL.Image)` produces) as CHW yields a
74 bogus size and a placeholder/embedding count mismatch at inference time.
75 """
76 items = ImageProcessorItems([image])
77
78 assert items.get_image_size(0) == (W, H)
79
80
81 @pytest.mark.parametrize(
82 "frame",
83 [
84 Image.new("RGB", (W, H)),
85 np.zeros((H, W, 3), dtype=np.uint8),
86 torch.zeros((H, W, 3), dtype=torch.uint8),
87 np.zeros((3, H, W), dtype=np.uint8),
88 torch.zeros((3, H, W), dtype=torch.uint8),
89 ],
90 )
91 def test_frame_size_hwc_chw(frame):
92 """`get_frame_size` must stay consistent with `get_image_size`."""
93 items = VideoProcessorItems([[frame]])
94
95 assert items.get_frame_size(0) == (W, H)
96
97
98 def test_video_with_metadata_tensor_passthrough():
99 """Tensor frames pass through unchanged regardless of device: HF video
100 processors accept tensors, and device-resident frames (e.g. NVDEC-decoded)
101 must not be copied back to host."""
102 frames = torch.zeros((4, H, W, 3), dtype=torch.uint8)
103 video, metadata = MultiModalDataParser()._get_video_with_metadata(frames)
104
105 assert video is frames
106 assert metadata is None
107
108
109 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
110 def test_video_with_metadata_keeps_device_tensor():
111 """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors,
112 so a device-side HF processor can consume them without a D2H copy."""
113 frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda")
114 video, metadata = MultiModalDataParser()._get_video_with_metadata(frames)
115
116 assert video is frames
117 assert metadata is None
118
119
120 @pytest.mark.parametrize(
121 "frames",
122 [
123 [np.zeros((H, W, 3), dtype=np.uint8) for _ in range(2)],
124 [torch.zeros((H, W, 3), dtype=torch.uint8) for _ in range(2)],
125 ],
126 )
127 def test_parse_video_frame_list_as_single_video(frames):
128 """A list of decoded frames must represent one video item."""
129 items = MultiModalDataParser().parse_mm_data({"video": frames})["video"]
130
131 assert items.get_count() == 1
132 video = items.get(0)
133 assert isinstance(video, np.ndarray)
134 np.testing.assert_array_equal(video, np.stack([np.asarray(f) for f in frames]))
135
136
137 @pytest.mark.parametrize(
138 "modality,processor_cls",
139 [
140 ("audio", AudioProcessorItems),
141 ("image", ImageProcessorItems),
142 ("video", VideoProcessorItems),
143 ],
144 )
145 def test_parse_mm_data_accepts_none_cached_item(modality, processor_cls):
146 mm_items = MultiModalDataParser().parse_mm_data({modality: [None]})
147 items = mm_items[modality]
148 assert isinstance(items, processor_cls)
149 assert len(items) == 1
150 assert items.get(0) is None
151
152
153 def test_cached_audio_items_preserve_positions_during_resampling():
154 waveform = np.arange(16, dtype=np.float32)
155 parser = MultiModalDataParser(
156 target_sr=16000, target_channels=1, audio_resample_method="scipy"
157 )
158 items = parser.parse_mm_data(
159 {"audio": [None, (waveform, 8000), None, (waveform, 16000)]}
160 )["audio"]
161
162 assert len(items) == 4
163 assert items.get(0) is None
164 assert items.get(2) is None
165 assert len(items.get(1)) == 32
166 np.testing.assert_array_equal(items.get(3), waveform)