A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 19:45 UTC
[Bugfix][Frontend] Return 400 for corrupt/truncated image inputs instead of 500 (#38253)
Signed-off-by: aliialsaeedii <ali.al-saeedi@nscale.com>
aliialsaeedii committed
Mar 30, 2026 at 12:26 UTC
7e76af14fabfc72144c627331b87d581df361158
2 files changed
+82
-6
tests/multimodal/media/test_image.py
+74
@@ -131,3 +131,77 @@ def test_image_media_io_rgba_background_color_validation():
131
ImageMediaIO(rgba_background_color=(0, 0, 0)) # Should not raise
132
ImageMediaIO(rgba_background_color=[255, 255, 255]) # Should not raise
133
ImageMediaIO(rgba_background_color=(128, 128, 128)) # Should not raise
134
+
135
+
136
+def test_image_media_io_load_bytes(tmp_path):
137
+ """Test load_bytes with valid and invalid image data."""
138
+ # Save a valid RGB image to use as source bytes
139
+ valid_image = Image.new("RGB", (8, 8), (100, 150, 200))
140
+ valid_path = tmp_path / "valid.png"
141
+ valid_image.save(valid_path)
142
+
143
+ valid_data = valid_path.read_bytes()
144
+
145
+ # Test 1: Valid image bytes load successfully and are fully decoded
146
+ image_io = ImageMediaIO()
147
+ result = image_io.load_bytes(valid_data)
148
+
149
+ # Check the returned media is a properly loaded image
150
+ assert isinstance(result.media, Image.Image)
151
+ assert result.media.size == (8, 8)
152
+ assert result.media.getpixel((0, 0)) == (100, 150, 200)
153
+
154
+ # Test 2: Garbage bytes raise ValueError
155
+ with pytest.raises(ValueError, match="Failed to load image"):
156
+ image_io.load_bytes(b"not an image")
157
+
158
+ # Test 3: Truncated PNG header raises ValueError
159
+ with pytest.raises(ValueError, match="Failed to load image"):
160
+ image_io.load_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 10)
161
+
162
+ # Test 4: Real PNG truncated mid-stream raises ValueError
163
+ with pytest.raises(ValueError, match="Failed to load image"):
164
+ image_io.load_bytes(valid_data[: len(valid_data) // 2])
165
+
166
+ # Test 5: Empty bytes raise ValueError
167
+ with pytest.raises(ValueError, match="Failed to load image"):
168
+ image_io.load_bytes(b"")
169
+
170
+
171
+def test_image_media_io_load_file(tmp_path):
172
+ """Test load_file with valid and invalid image files."""
173
+ # Save a valid RGB image to disk
174
+ valid_image = Image.new("RGB", (4, 4), (10, 20, 30))
175
+ valid_path = tmp_path / "valid.png"
176
+ valid_image.save(valid_path)
177
+
178
+ # Test 1: Valid image file loads successfully and is fully decoded
179
+ image_io = ImageMediaIO()
180
+ result = image_io.load_file(valid_path)
181
+
182
+ # Check the returned media is a properly loaded image
183
+ assert isinstance(result.media, Image.Image)
184
+ assert result.media.size == (4, 4)
185
+ assert result.media.getpixel((0, 0)) == (10, 20, 30)
186
+
187
+ # Test 2: File with garbage content raises ValueError
188
+ bad_file = tmp_path / "bad.png"
189
+ bad_file.write_bytes(b"this is not an image")
190
+
191
+ with pytest.raises(ValueError, match="Failed to load image"):
192
+ image_io.load_file(bad_file)
193
+
194
+ # Test 3: File with truncated PNG header raises ValueError
195
+ truncated_file = tmp_path / "truncated.png"
196
+ truncated_file.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 10)
197
+
198
+ with pytest.raises(ValueError, match="Failed to load image"):
199
+ image_io.load_file(truncated_file)
200
+
201
+ # Test 4: Real PNG file truncated mid-stream raises ValueError
202
+ valid_data = valid_path.read_bytes()
203
+ truncated_real_file = tmp_path / "truncated_real.png"
204
+ truncated_real_file.write_bytes(valid_data[: len(valid_data) // 2])
205
+
206
+ with pytest.raises(ValueError, match="Failed to load image"):
207
+ image_io.load_file(truncated_real_file)
vllm/multimodal/media/image.py
+8
-6
@@ -68,17 +68,19 @@ class ImageMediaIO(MediaIO[Image.Image]):
68
return convert_image_mode(image, self.image_mode)
69
70
def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]:
71
- image = Image.open(BytesIO(data))
72
- return MediaWithBytes(self._convert_image_mode(image), data)
71
+ try:
72
+ image = Image.open(BytesIO(data))
73
+ image.load()
74
+ image = self._convert_image_mode(image)
75
+ except (OSError, Image.UnidentifiedImageError) as e:
76
+ raise ValueError(f"Failed to load image: {e}") from e
77
+ return MediaWithBytes(image, data)
78
79
def load_base64(self, media_type: str, data: str) -> MediaWithBytes[Image.Image]:
80
return self.load_bytes(pybase64.b64decode(data, validate=True))
81
82
def load_file(self, filepath: Path) -> MediaWithBytes[Image.Image]:
78
- with open(filepath, "rb") as f:
79
- data = f.read()
80
- image = Image.open(BytesIO(data))
81
- return MediaWithBytes(self._convert_image_mode(image), data)
83
+ return self.load_bytes(filepath.read_bytes())
84
85
def encode_base64(
86
self,