[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:

vllm_omni/entrypoints/openai/api_server.py

@@ -161,6 +161,8 @@ from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import (
     VideoFrames,
+    _decode_image_bytes,
+    _validate_image_pixel_limit,
     decode_audio_url,
     decode_input_reference,
 )
@@ -3438,6 +3440,15 @@ def _validate_minimax_h3_image_payload(
     try:
         with Image.open(io.BytesIO(payload)) as image:
             image_format = str(image.format or "").lower()
+            if image_format in MINIMAX_H3_REFERENCE_IMAGE_FORMATS:
+                _validate_image_pixel_limit(image)
+    except InvalidInputReferenceError as exc:
+        raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc
+    except Image.DecompressionBombError as exc:
+        raise HTTPException(
+            HTTPStatus.BAD_REQUEST.value,
+            detail=f"Invalid uploaded image reference: {filename}; image exceeds the decoder pixel limit.",
+        ) from exc
     except (OSError, ValueError) as exc:
         if allow_non_image:
             return
@@ -3478,10 +3489,14 @@ async def _persist_uploaded_media_references(
             if kind == "image":
                 try:
                     _validate_minimax_h3_image_payload(payload, filename=upload.filename)
-                    with Image.open(io.BytesIO(payload)) as image:
-                        images.append(image.convert("RGB"))
-                except (OSError, ValueError) as exc:
-                    raise HTTPException(400, detail=f"Invalid uploaded image reference: {upload.filename}") from exc
+                    images.append(
+                        _decode_image_bytes(
+                            payload,
+                            source=f"uploaded image reference: {upload.filename}",
+                        )
+                    )
+                except InvalidInputReferenceError as exc:
+                    raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc
                 continue
             suffix = Path(upload.filename or "").suffix.lower()
             if kind == "video" and suffix and suffix not in MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES:

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -71,6 +71,21 @@ class VideoFrames(list[Image.Image]):
         self.source_path = source_path
 
 
+class _ImagePixelLimitError(InvalidInputReferenceError):
+    """An image exceeded a configured or decoder-enforced pixel limit."""
+
+
+def _validate_image_pixel_limit(image: Image.Image) -> None:
+    width, height = image.size
+    max_pixels = envs.VLLM_MAX_IMAGE_PIXELS
+    if max_pixels > 0 and width * height > max_pixels:
+        raise _ImagePixelLimitError(
+            f"Image dimensions {width}x{height} ({width * height} pixels) exceed "
+            f"the maximum of {max_pixels} pixels. Set "
+            f"VLLM_MAX_IMAGE_PIXELS to increase this limit."
+        )
+
+
 def positive_float(value: Any) -> float | None:
     if value is None:
         return None
@@ -87,7 +102,13 @@ def positive_float(value: Any) -> float | None:
 
 def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
     try:
-        return Image.open(BytesIO(image_bytes)).convert("RGB")
+        with Image.open(BytesIO(image_bytes)) as image:
+            _validate_image_pixel_limit(image)
+            return image.convert("RGB")
+    except _ImagePixelLimitError:
+        raise
+    except Image.DecompressionBombError as exc:
+        raise _ImagePixelLimitError(f"Invalid {source}: image exceeds the decoder pixel limit.") from exc
     except (UnidentifiedImageError, OSError, ValueError) as exc:
         raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc
 
@@ -158,6 +179,8 @@ def _decode_media_bytes(
 ) -> Image.Image | VideoFrames:
     try:
         return _decode_image_bytes(media_bytes, source=source)
+    except _ImagePixelLimitError:
+        raise
     except InvalidInputReferenceError:
         try:
             return _decode_video_bytes(