[Bugfix][Frontend] Enforce image pixel limits for video input references (#6963)

Signed-off-by: BANANASJIM <bananasjim1@gmail.com>

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -23,9 +23,9 @@ from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _png_bytes() -> bytes:
+def _png_bytes(size: tuple[int, int] = (2, 1)) -> bytes:
     buffer = BytesIO()
-    Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG")
+    Image.new("RGB", size, color=(12, 34, 56)).save(buffer, format="PNG")
     return buffer.getvalue()
 
 
@@ -122,6 +122,52 @@ async def test_decode_image_url_keeps_data_urls_local(monkeypatch):
     assert image.mode == "RGB"
 
 
+def test_decode_image_bytes_rejects_image_over_pixel_limit_before_conversion(monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100)
+
+    def _unexpected_convert(*args, **kwargs):
+        pytest.fail("over-limit images must be rejected before conversion")
+
+    monkeypatch.setattr(Image.Image, "convert", _unexpected_convert)
+
+    with pytest.raises(
+        InvalidInputReferenceError,
+        match=r"Image dimensions 20x20 \(400 pixels\) exceed the maximum of 100 pixels",
+    ):
+        video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference")
+
+
+def test_decode_image_bytes_maps_pillow_pixel_limit_error(monkeypatch):
+    monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100)
+
+    with pytest.raises(InvalidInputReferenceError, match="decoder pixel limit") as exc_info:
+        video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference")
+
+    assert isinstance(exc_info.value.__cause__, Image.DecompressionBombError)
+
+
+@pytest.mark.parametrize("max_pixels", [0, 400])
+def test_decode_image_bytes_allows_disabled_or_inclusive_pixel_limit(monkeypatch, max_pixels):
+    monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", max_pixels)
+
+    image = video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference")
+
+    assert image.size == (20, 20)
+    assert image.mode == "RGB"
+
+
+@pytest.mark.asyncio
+async def test_decode_input_reference_preserves_image_pixel_limit_error(monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100)
+
+    with pytest.raises(InvalidInputReferenceError, match="VLLM_MAX_IMAGE_PIXELS"):
+        await video_api_utils.decode_input_reference(
+            image_reference=None,
+            video_reference=None,
+            input_reference_bytes=_png_bytes((20, 20)),
+        )
+
+
 def _install_fake_video_mux(monkeypatch, mux_calls):
     def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None):
         mux_calls.append(

tests/entrypoints/openai_api/test_video_server.py

@@ -1630,7 +1630,8 @@ def test_generic_video_model_rejects_mixed_image_and_video_references(test_clien
     assert "does not support mixed image and video" in response.json()["detail"].lower()
 
 
-def test_h3_multipart_rejects_bmp_image_reference(test_client):
+def test_h3_multipart_rejects_bmp_image_reference(test_client, monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100)
     image = Image.new("RGB", (64, 64), color="blue")
     image_buffer = io.BytesIO()
     image.save(image_buffer, format="BMP")
@@ -1646,6 +1647,37 @@ def test_h3_multipart_rejects_bmp_image_reference(test_client):
     assert "must use jpg" in response.json()["detail"].lower()
 
 
+@pytest.mark.parametrize("field", ["input_reference", "input_references"])
+def test_h3_multipart_rejects_image_over_pixel_limit(field, test_client, monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100)
+    test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline"
+
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "reject oversized image", "extra_params": '{"task":"ref2va"}'},
+        files=[(field, ("reference.png", _make_test_image_bytes((20, 20)), "image/png"))],
+    )
+
+    assert response.status_code == 400
+    assert "VLLM_MAX_IMAGE_PIXELS" in response.json()["detail"]
+
+
+@pytest.mark.parametrize("field", ["input_reference", "input_references"])
+def test_h3_multipart_maps_pillow_pixel_limit_error(field, test_client, monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 0)
+    monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100)
+    test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline"
+
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "reject decoder bomb", "extra_params": '{"task":"ref2va"}'},
+        files=[(field, ("reference.png", _make_test_image_bytes((20, 20)), "image/png"))],
+    )
+
+    assert response.status_code == 400
+    assert "decoder pixel limit" in response.json()["detail"]
+
+
 @pytest.mark.asyncio
 async def test_h3_upload_limit_checks_declared_size_before_read():
     class OversizedUpload: