diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py
index 2a5c5be47..fbf56b70e 100644
--- a/tests/entrypoints/openai_api/test_decode_audio_url.py
+++ b/tests/entrypoints/openai_api/test_decode_audio_url.py
@@ -1,16 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for decode_audio_url helper."""
+"""Unit tests for decode_audio_url helper and _AudioFileMediaIO."""
import asyncio
import base64
import os
-import httpx
import pytest
+from vllm.multimodal.media import MediaConnector
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
-from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url
+from vllm_omni.entrypoints.openai.video_api_utils import (
+ _AudioFileMediaIO,
+ decode_audio_url,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -24,15 +27,127 @@ def _cleanup(path):
os.unlink(path)
+def _permissive_connector():
+ return MediaConnector(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ )
+
+
+def _restricted_connector(allowed_domains=None):
+ return MediaConnector(
+ allowed_local_media_path="",
+ allowed_media_domains=allowed_domains or [],
+ )
+
+
+# ---------------------------------------------------------------------------
+# _AudioFileMediaIO direct tests
+# ---------------------------------------------------------------------------
+
+
+class TestAudioFileMediaIO:
+ """Direct tests for _AudioFileMediaIO methods."""
+
+ def test_load_bytes_writes_temp_file(self):
+ audio_bytes = b"\xff\xfb\x90\x00" * 50
+ io = _AudioFileMediaIO()
+ path = io.load_bytes(audio_bytes)
+ try:
+ assert os.path.isfile(path)
+ assert path.endswith(".wav")
+ with open(path, "rb") as f:
+ assert f.read() == audio_bytes
+ finally:
+ _cleanup(path)
+
+ def test_load_bytes_empty_raises(self):
+ io = _AudioFileMediaIO()
+ with pytest.raises(InvalidInputReferenceError, match="audio data is empty"):
+ io.load_bytes(b"")
+
+ def test_load_base64_mp3_suffix(self):
+ audio_bytes = b"\xff\xfb\x90\x00" * 50
+ b64 = base64.b64encode(audio_bytes).decode()
+ io = _AudioFileMediaIO()
+ path = io.load_base64("audio/mp3", b64)
+ try:
+ assert path.endswith(".mp3")
+ with open(path, "rb") as f:
+ assert f.read() == audio_bytes
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_mpeg_uses_mp3(self):
+ io = _AudioFileMediaIO()
+ path = io.load_base64("audio/mpeg", base64.b64encode(b"\x00" * 10).decode())
+ try:
+ assert path.endswith(".mp3")
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_wav_suffix(self):
+ io = _AudioFileMediaIO()
+ path = io.load_base64("audio/wav", base64.b64encode(b"\x00" * 10).decode())
+ try:
+ assert path.endswith(".wav")
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_flac_suffix(self):
+ io = _AudioFileMediaIO()
+ path = io.load_base64("audio/flac", base64.b64encode(b"\x00" * 10).decode())
+ try:
+ assert path.endswith(".flac")
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_empty_media_type_uses_wav(self):
+ io = _AudioFileMediaIO()
+ path = io.load_base64("", base64.b64encode(b"\x00" * 10).decode())
+ try:
+ assert path.endswith(".wav")
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_non_alnum_falls_back_to_wav(self):
+ io = _AudioFileMediaIO()
+ path = io.load_base64("audio/foo-bar", base64.b64encode(b"\x00" * 10).decode())
+ try:
+ assert path.endswith(".wav")
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_too_long_extension_falls_back_to_wav(self):
+ io = _AudioFileMediaIO()
+ path = io.load_base64("audio/verylongextension", base64.b64encode(b"\x00" * 10).decode())
+ try:
+ assert path.endswith(".wav")
+ finally:
+ _cleanup(path)
+
+ def test_load_file_returns_path_string(self, tmp_path):
+ audio_file = tmp_path / "test.wav"
+ audio_file.write_bytes(b"\x00" * 100)
+ io = _AudioFileMediaIO()
+ result = io.load_file(audio_file)
+ assert result == str(audio_file)
+
+
+# ---------------------------------------------------------------------------
+# decode_audio_url with data URLs (exercises real MediaConnector + MediaIO)
+# ---------------------------------------------------------------------------
+
+
class TestDecodeAudioUrlDataUrl:
- """Tests for base64 data URL decoding."""
+ """Tests for base64 data URL decoding — no mocks, real code path."""
def test_valid_mp3_data_url(self):
audio_bytes = b"\xff\xfb\x90\x00" * 100
b64 = base64.b64encode(audio_bytes).decode()
url = f"data:audio/mp3;base64,{b64}"
- path = _run(decode_audio_url(url))
+ path = _run(decode_audio_url(url, _permissive_connector()))
try:
assert os.path.isfile(path)
assert path.endswith(".mp3")
@@ -46,7 +161,7 @@ class TestDecodeAudioUrlDataUrl:
b64 = base64.b64encode(audio_bytes).decode()
url = f"data:audio/wav;base64,{b64}"
- path = _run(decode_audio_url(url))
+ path = _run(decode_audio_url(url, _permissive_connector()))
try:
assert path.endswith(".wav")
with open(path, "rb") as f:
@@ -54,134 +169,103 @@ class TestDecodeAudioUrlDataUrl:
finally:
_cleanup(path)
- def test_valid_mpeg_data_url_uses_mp3_suffix(self):
- audio_bytes = b"\xff\xfb\x90\x00" * 50
- b64 = base64.b64encode(audio_bytes).decode()
- url = f"data:audio/mpeg;base64,{b64}"
-
- path = _run(decode_audio_url(url))
- try:
- assert path.endswith(".mp3")
- finally:
- _cleanup(path)
-
- def test_valid_ogg_data_url(self):
- audio_bytes = b"OggS" + b"\x00" * 100
- b64 = base64.b64encode(audio_bytes).decode()
- url = f"data:audio/ogg;base64,{b64}"
-
- path = _run(decode_audio_url(url))
- try:
- assert path.endswith(".ogg")
- finally:
- _cleanup(path)
-
def test_invalid_base64_raises(self):
url = "data:audio/mp3;base64,!!!not-valid-base64!!!"
-
- with pytest.raises(InvalidInputReferenceError, match="not valid base64"):
- _run(decode_audio_url(url))
+ with pytest.raises(InvalidInputReferenceError, match="Incorrect padding"):
+ _run(decode_audio_url(url, _permissive_connector()))
def test_empty_audio_data_raises(self):
url = "data:audio/mp3;base64,"
-
with pytest.raises(InvalidInputReferenceError, match="audio data is empty"):
- _run(decode_audio_url(url))
+ _run(decode_audio_url(url, _permissive_connector()))
-class TestDecodeAudioUrlSuffixSanitization:
- """Tests for MIME extension sanitization (path traversal prevention)."""
-
- def test_path_traversal_in_mime_is_neutralized(self):
- audio_bytes = b"\x00" * 100
- b64 = base64.b64encode(audio_bytes).decode()
- url = f"data:audio/../..;base64,{b64}"
-
- path = _run(decode_audio_url(url))
- try:
- # ".." contains non-alnum dots, so falls back to .wav
- assert path.endswith(".wav")
- assert os.path.dirname(path) == "/tmp"
- finally:
- _cleanup(path)
-
- def test_non_alnum_extension_falls_back_to_wav(self):
- audio_bytes = b"\x00" * 100
- b64 = base64.b64encode(audio_bytes).decode()
- url = f"data:audio/foo-bar;base64,{b64}"
-
- path = _run(decode_audio_url(url))
- try:
- assert path.endswith(".wav")
- finally:
- _cleanup(path)
-
- def test_too_long_extension_falls_back_to_wav(self):
- audio_bytes = b"\x00" * 100
- b64 = base64.b64encode(audio_bytes).decode()
- url = f"data:audio/verylongextension;base64,{b64}"
-
- path = _run(decode_audio_url(url))
- try:
- assert path.endswith(".wav")
- finally:
- _cleanup(path)
-
- def test_valid_short_alnum_extension_is_used(self):
- audio_bytes = b"\x00" * 100
- b64 = base64.b64encode(audio_bytes).decode()
- url = f"data:audio/flac;base64,{b64}"
-
- path = _run(decode_audio_url(url))
- try:
- assert path.endswith(".flac")
- finally:
- _cleanup(path)
+# ---------------------------------------------------------------------------
+# decode_audio_url with file:// URIs (exercises load_file through connector)
+# ---------------------------------------------------------------------------
-class TestDecodeAudioUrlHttp:
- """Tests for HTTP URL decoding."""
+class TestDecodeAudioUrlFileUri:
+ """Tests using file:// URIs — exercises MediaConnector + load_file."""
- def test_valid_http_url(self, monkeypatch):
- audio_bytes = b"\xff\xfb\x90\x00" * 50
- fake_resp = httpx.Response(200, content=audio_bytes)
+ def test_file_uri_with_allowed_path(self, tmp_path):
+ audio_file = tmp_path / "test.wav"
+ audio_file.write_bytes(b"\xff\xfb\x90\x00" * 50)
+ file_uri = audio_file.as_uri()
+ connector = MediaConnector(
+ allowed_local_media_path=str(tmp_path),
+ allowed_media_domains=None,
+ )
+ path = _run(decode_audio_url(file_uri, connector))
+ assert path == str(audio_file)
- async def _mock_get(self, url):
- return fake_resp
+ def test_file_uri_without_allowed_path_raises(self, tmp_path):
+ audio_file = tmp_path / "test.wav"
+ audio_file.write_bytes(b"\x00" * 50)
+ file_uri = audio_file.as_uri()
+ connector = MediaConnector(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ )
+ with pytest.raises(InvalidInputReferenceError):
+ _run(decode_audio_url(file_uri, connector))
- monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
- monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None)
+ def test_file_uri_outside_allowed_path_raises(self, tmp_path):
+ audio_file = tmp_path / "test.wav"
+ audio_file.write_bytes(b"\x00" * 50)
+ file_uri = audio_file.as_uri()
+ other_dir = tmp_path / "other"
+ other_dir.mkdir()
+ connector = MediaConnector(
+ allowed_local_media_path=str(other_dir),
+ allowed_media_domains=None,
+ )
+ with pytest.raises(InvalidInputReferenceError):
+ _run(decode_audio_url(file_uri, connector))
- path = _run(decode_audio_url("https://example.com/audio.mp3"))
- try:
- assert os.path.isfile(path)
- assert path.endswith(".wav")
- with open(path, "rb") as f:
- assert f.read() == audio_bytes
- finally:
- _cleanup(path)
- def test_http_error_raises(self, monkeypatch):
- async def _mock_get(self, url):
- raise httpx.HTTPStatusError("Not Found", request=httpx.Request("GET", url), response=httpx.Response(404))
-
- monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
-
- with pytest.raises(InvalidInputReferenceError, match="failed to download audio"):
- _run(decode_audio_url("https://example.com/missing.mp3"))
+# ---------------------------------------------------------------------------
+# decode_audio_url — invalid schemes and SSRF
+# ---------------------------------------------------------------------------
class TestDecodeAudioUrlInvalid:
"""Tests for invalid URL schemes."""
def test_local_file_path_raises(self):
- with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
- _run(decode_audio_url("/path/to/local/file.mp3"))
+ with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+ _run(decode_audio_url("/path/to/local/file.mp3", _permissive_connector()))
def test_ftp_url_raises(self):
- with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
- _run(decode_audio_url("ftp://example.com/file.mp3"))
+ with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+ _run(decode_audio_url("ftp://example.com/file.mp3", _permissive_connector()))
def test_empty_string_raises(self):
- with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
- _run(decode_audio_url(""))
+ with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+ _run(decode_audio_url("", _permissive_connector()))
+
+
+class TestDecodeAudioUrlSSRF:
+ """Tests for SSRF protection via MediaConnector domain restrictions."""
+
+ def test_restricted_domain_rejects_disallowed_url(self):
+ connector = _restricted_connector(allowed_domains=["trusted.example.com"])
+ with pytest.raises(InvalidInputReferenceError, match="allowed domains"):
+ _run(
+ decode_audio_url(
+ "https://evil.internal/admin/audio.mp3",
+ connector,
+ )
+ )
+
+ def test_data_url_bypasses_domain_restriction(self):
+ audio_bytes = b"\xff\xfb\x90\x00" * 50
+ b64 = base64.b64encode(audio_bytes).decode()
+ url = f"data:audio/wav;base64,{b64}"
+
+ connector = _restricted_connector(allowed_domains=["trusted.example.com"])
+ path = _run(decode_audio_url(url, connector))
+ try:
+ assert os.path.isfile(path)
+ finally:
+ _cleanup(path)
diff --git a/tests/entrypoints/openai_api/test_decode_video_url.py b/tests/entrypoints/openai_api/test_decode_video_url.py
new file mode 100644
index 000000000..64ff7e2e0
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_decode_video_url.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for decode_video_url helper and _VideoFramesMediaIO."""
+
+import asyncio
+import base64
+import io
+
+import pytest
+from PIL import Image
+from vllm.multimodal.media import MediaConnector
+
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
+from vllm_omni.entrypoints.openai.video_api_utils import (
+ VideoFrames,
+ _VideoFramesMediaIO,
+ decode_video_url,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _run(coro):
+ return asyncio.run(coro)
+
+
+def _permissive_connector():
+ return MediaConnector(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ )
+
+
+def _restricted_connector(allowed_domains=None):
+ return MediaConnector(
+ allowed_local_media_path="",
+ allowed_media_domains=allowed_domains or [],
+ )
+
+
+def _make_mp4_bytes(width=4, height=4, num_frames=2):
+ """Create a minimal valid MP4 video using PyAV."""
+ import av
+
+ buf = io.BytesIO()
+ with av.open(buf, mode="w", format="mp4") as container:
+ stream = container.add_stream("libx264", rate=24)
+ stream.width = width
+ stream.height = height
+ stream.pix_fmt = "yuv420p"
+ stream.options = {"preset": "ultrafast"}
+ for i in range(num_frames):
+ img = Image.new("RGB", (width, height), color=(i * 40, 128, 200))
+ frame = av.VideoFrame.from_image(img)
+ for packet in stream.encode(frame):
+ container.mux(packet)
+ for packet in stream.encode():
+ container.mux(packet)
+ return buf.getvalue()
+
+
+# ---------------------------------------------------------------------------
+# _VideoFramesMediaIO direct tests
+# ---------------------------------------------------------------------------
+
+
+class TestVideoFramesMediaIO:
+ """Direct tests for _VideoFramesMediaIO methods."""
+
+ def test_load_bytes_decodes_video(self):
+ video_bytes = _make_mp4_bytes(num_frames=3)
+ media_io = _VideoFramesMediaIO()
+ result = media_io.load_bytes(video_bytes)
+ assert isinstance(result, VideoFrames)
+ assert len(result) >= 1
+ assert all(isinstance(f, Image.Image) for f in result)
+
+ def test_load_bytes_respects_max_frames(self):
+ video_bytes = _make_mp4_bytes(num_frames=5)
+ media_io = _VideoFramesMediaIO(max_frames=2)
+ result = media_io.load_bytes(video_bytes)
+ assert len(result) <= 2
+
+ def test_load_bytes_invalid_data_raises(self):
+ media_io = _VideoFramesMediaIO()
+ with pytest.raises(InvalidInputReferenceError, match="not a valid video"):
+ media_io.load_bytes(b"not a video")
+
+ def test_load_base64_decodes_video(self):
+ video_bytes = _make_mp4_bytes()
+ b64 = base64.b64encode(video_bytes).decode()
+ media_io = _VideoFramesMediaIO()
+ result = media_io.load_base64("video/mp4", b64)
+ assert isinstance(result, VideoFrames)
+ assert len(result) >= 1
+
+ def test_load_file_decodes_video(self, tmp_path):
+ video_bytes = _make_mp4_bytes()
+ video_file = tmp_path / "test.mp4"
+ video_file.write_bytes(video_bytes)
+ media_io = _VideoFramesMediaIO()
+ result = media_io.load_file(video_file)
+ assert isinstance(result, VideoFrames)
+ assert len(result) >= 1
+
+
+# ---------------------------------------------------------------------------
+# decode_video_url with data URLs (real MediaConnector + MediaIO)
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeVideoUrlDataUrl:
+ """Tests for data URL decoding — no mocks, real code path."""
+
+ def test_valid_video_data_url(self):
+ video_bytes = _make_mp4_bytes()
+ b64 = base64.b64encode(video_bytes).decode()
+ url = f"data:video/mp4;base64,{b64}"
+
+ result = _run(decode_video_url(url, _permissive_connector()))
+ assert isinstance(result, VideoFrames)
+ assert len(result) >= 1
+ assert all(isinstance(f, Image.Image) for f in result)
+
+ def test_invalid_video_data_url_raises(self):
+ b64 = base64.b64encode(b"not a video").decode()
+ url = f"data:video/mp4;base64,{b64}"
+ with pytest.raises(InvalidInputReferenceError, match="not a valid video"):
+ _run(decode_video_url(url, _permissive_connector()))
+
+
+# ---------------------------------------------------------------------------
+# decode_video_url with file:// URIs (exercises load_file through connector)
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeVideoUrlFileUri:
+ """Tests using file:// URIs — exercises MediaConnector + load_file."""
+
+ def test_file_uri_with_allowed_path(self, tmp_path):
+ video_file = tmp_path / "test.mp4"
+ video_file.write_bytes(_make_mp4_bytes())
+ file_uri = video_file.as_uri()
+ connector = MediaConnector(
+ allowed_local_media_path=str(tmp_path),
+ allowed_media_domains=None,
+ )
+ result = _run(decode_video_url(file_uri, connector))
+ assert isinstance(result, VideoFrames)
+ assert len(result) >= 1
+
+ def test_file_uri_without_allowed_path_raises(self, tmp_path):
+ video_file = tmp_path / "test.mp4"
+ video_file.write_bytes(_make_mp4_bytes())
+ file_uri = video_file.as_uri()
+ connector = MediaConnector(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ )
+ with pytest.raises(InvalidInputReferenceError):
+ _run(decode_video_url(file_uri, connector))
+
+
+# ---------------------------------------------------------------------------
+# decode_video_url — invalid schemes and SSRF
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeVideoUrlInvalid:
+ """Tests for invalid URL schemes."""
+
+ def test_local_file_path_raises(self):
+ with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+ _run(decode_video_url("/path/to/local/file.mp4", _permissive_connector()))
+
+ def test_ftp_url_raises(self):
+ with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+ _run(decode_video_url("ftp://example.com/file.mp4", _permissive_connector()))
+
+ def test_empty_string_raises(self):
+ with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+ _run(decode_video_url("", _permissive_connector()))
+
+
+class TestDecodeVideoUrlSSRF:
+ """Tests for SSRF protection via MediaConnector domain restrictions."""
+
+ def test_restricted_domain_rejects_disallowed_url(self):
+ connector = _restricted_connector(
+ allowed_domains=["trusted.example.com"],
+ )
+ with pytest.raises(InvalidInputReferenceError, match="allowed domains"):
+ _run(
+ decode_video_url(
+ "https://evil.internal/admin/video.mp4",
+ connector,
+ )
+ )
+
+ def test_data_url_bypasses_domain_restriction(self):
+ video_bytes = _make_mp4_bytes()
+ b64 = base64.b64encode(video_bytes).decode()
+ url = f"data:video/mp4;base64,{b64}"
+
+ connector = _restricted_connector(
+ allowed_domains=["trusted.example.com"],
+ )
+ result = _run(decode_video_url(url, connector))
+ assert isinstance(result, VideoFrames)
+ assert len(result) >= 1
diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 31a82102a..49c43d7a9 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -13,6 +13,7 @@ import json
from argparse import Namespace
from http import HTTPStatus
from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import FastAPI, HTTPException
@@ -24,12 +25,19 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.sampling_params import RequestOutputKind
from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.entrypoints.openai.api_server import _check_max_generated_image_size, _DiffusionServingModels, router
+from vllm_omni.entrypoints.openai.api_server import (
+ _check_max_generated_image_size,
+ _DiffusionServingModels,
+ _load_input_images,
+ router,
+)
from vllm_omni.entrypoints.openai.image_api_utils import (
encode_image_base64,
parse_size,
)
+from vllm_omni.entrypoints.openai.protocol.videos import UrlImageReference
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference
from vllm_omni.errors import GuardrailViolationError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -182,6 +190,7 @@ def mock_async_diffusion(mocker: MockerFixture):
self.captured_sampling_params_list = None
self.captured_prompt = None
self.generate_calls = 0
+ self.model_config = SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
async def generate(self, **kwargs):
self.generate_calls += 1
@@ -271,6 +280,10 @@ def async_omni_test_client():
def get_diffusion_od_config(self):
return self.od_config
+ @property
+ def model_config(self):
+ return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
app = FastAPI()
app.include_router(router)
@@ -335,6 +348,10 @@ def async_omni_rgba_test_client():
def get_diffusion_od_config(self):
return self.od_config
+ @property
+ def model_config(self):
+ return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
app = FastAPI()
app.include_router(router)
@@ -399,6 +416,10 @@ def async_omni_stage_configs_only_client():
def get_diffusion_od_config(self):
return self.od_config
+ @property
+ def model_config(self):
+ return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
app = FastAPI()
app.include_router(router)
@@ -435,9 +456,24 @@ def streaming_image_edit_client():
SamplingParams(temperature=0.1),
OmniDiffusionSamplingParams(),
]
+ stage_clients = [
+ SimpleNamespace(stage_type="llm", is_comprehension=True),
+ SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+ ]
+ stage_vllm_configs = [
+ SimpleNamespace(
+ model_config=SimpleNamespace(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ ),
+ ),
+ None,
+ ]
self.engine = SimpleNamespace(
stage_configs=stage_configs,
default_sampling_params_list=default_sampling_params_list,
+ stage_clients=stage_clients,
+ stage_vllm_configs=stage_vllm_configs,
)
self.default_sampling_params_list = default_sampling_params_list
self.captured_sampling_params_list = None
@@ -729,6 +765,10 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
def get_diffusion_od_config(self):
return self.od_config
+ @property
+ def model_config(self):
+ return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
app = FastAPI()
app.include_router(router)
engine = FakeAsyncOmniClass()
@@ -2136,3 +2176,107 @@ def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_on
assert captured_prompt["prompt"].count("
") == 2, (
f"N=2 reference images must emit 2
placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}"
)
+
+
+# ---------------------------------------------------------------------------
+# MediaConnector / SSRF protection tests
+# ---------------------------------------------------------------------------
+
+
+_DATA_IMAGE_URL = f"data:image/png;base64,{base64.b64encode(make_test_image_bytes((8, 8))).decode()}"
+
+
+def _make_upload_image():
+ upload = MagicMock()
+ upload.file = True
+ upload.read = AsyncMock(return_value=make_test_image_bytes((8, 8)))
+ return upload
+
+
+class TestLoadInputImagesMediaConnector:
+ """SSRF protection via MediaConnector for _load_input_images and
+ decode_input_reference."""
+
+ @staticmethod
+ def _config(
+ domains=None, # allow all domains
+ local_path="", # disallow all local paths
+ ):
+ return SimpleNamespace(
+ allowed_local_media_path=local_path,
+ allowed_media_domains=domains,
+ )
+
+ @pytest.mark.parametrize(
+ "domains, url",
+ [
+ # restricted + data url from allowlisted domain → accept (data: bypasses domain check)
+ (["example.com"], _DATA_IMAGE_URL),
+ # unrestricted + data url → accept
+ (None, _DATA_IMAGE_URL),
+ # restricted + file upload → accept (UploadFile bypasses MediaConnector)
+ (["example.com"], _make_upload_image()),
+ ],
+ ids=[
+ "restricted-data-url",
+ "unrestricted-data-url",
+ "restricted-file-upload",
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_accepted_inputs(self, domains, url):
+ config = self._config(domains=domains)
+
+ images = await _load_input_images([url], config)
+ assert len(images) == 1
+ assert isinstance(images[0], Image.Image)
+
+ if isinstance(url, str):
+ ref = UrlImageReference(image_url=url)
+ image = await decode_input_reference(ref, None, None, config)
+ assert isinstance(image, Image.Image)
+
+ @pytest.mark.asyncio
+ async def test_http_url_rejected_by_domain_filter(self):
+ """An http(s) URL whose domain is not in the allowlist must be rejected."""
+ config = self._config(domains=["example.com"])
+ http_url = "https://blocked.example.org/image.png"
+
+ with pytest.raises(ValueError):
+ await _load_input_images([http_url], config)
+
+ ref = UrlImageReference(image_url=http_url)
+ with pytest.raises(ValueError):
+ await decode_input_reference(ref, None, None, config)
+
+ @pytest.mark.parametrize(
+ "allow_path",
+ [True, False],
+ ids=["allowed-path", "disallowed-path"],
+ )
+ @pytest.mark.asyncio
+ async def test_file_uri_path_restriction(self, tmp_path, allow_path):
+ """file:// URIs are gated by allowed_local_media_path."""
+
+ img_path = tmp_path / "test.png"
+ img_path.write_bytes(make_test_image_bytes((8, 8)))
+ file_uri = img_path.as_uri()
+ config = self._config(local_path=str(tmp_path) if allow_path else "")
+
+ # --- _load_input_images ---
+ if allow_path:
+ images = await _load_input_images([file_uri], config)
+ assert len(images) == 1
+ assert isinstance(images[0], Image.Image)
+ else:
+ with pytest.raises(ValueError):
+ await _load_input_images([file_uri], config)
+
+ # --- decode_input_reference ---
+ ref = UrlImageReference(image_url=file_uri)
+ if allow_path:
+ image = await decode_input_reference(ref, None, None, config)
+ assert isinstance(image, Image.Image)
+ else:
+ with pytest.raises(ValueError):
+ await decode_input_reference(ref, None, None, config)
diff --git a/tests/entrypoints/openai_api/test_materialize_media_source.py b/tests/entrypoints/openai_api/test_materialize_media_source.py
new file mode 100644
index 000000000..a0ce824b3
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_materialize_media_source.py
@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for audiox materialize_media_source and _TempFileMediaIO."""
+
+import base64
+import os
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.transformers_utils.processors.audiox import (
+ _TempFileMediaIO,
+ materialize_media_source,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _cleanup(path):
+ if path and os.path.exists(path):
+ os.unlink(path)
+
+
+def _fake_forward_context(allowed_domains=None, allowed_local_media_path=""):
+ """Return a fake ForwardContext with model_config for SSRF testing."""
+ return SimpleNamespace(
+ vllm_config=SimpleNamespace(
+ model_config=SimpleNamespace(
+ allowed_local_media_path=allowed_local_media_path,
+ allowed_media_domains=allowed_domains,
+ ),
+ ),
+ )
+
+
+# ---------------------------------------------------------------------------
+# _TempFileMediaIO direct tests
+# ---------------------------------------------------------------------------
+
+
+class TestTempFileMediaIO:
+ """Direct tests for _TempFileMediaIO methods."""
+
+ def test_load_bytes_writes_temp_file(self):
+ data = b"\x00\x01\x02" * 50
+ io = _TempFileMediaIO()
+ path = io.load_bytes(data)
+ try:
+ assert os.path.isfile(path)
+ assert path.endswith(".bin")
+ with open(path, "rb") as f:
+ assert f.read() == data
+ finally:
+ _cleanup(path)
+
+ def test_load_base64_writes_temp_file(self):
+ data = b"\xff" * 100
+ b64 = base64.b64encode(data).decode()
+ io = _TempFileMediaIO()
+ path = io.load_base64("video/mp4", b64)
+ try:
+ assert os.path.isfile(path)
+ with open(path, "rb") as f:
+ assert f.read() == data
+ finally:
+ _cleanup(path)
+
+ def test_load_file_returns_path_string(self, tmp_path):
+ media_file = tmp_path / "video.mp4"
+ media_file.write_bytes(b"\x00" * 50)
+ io = _TempFileMediaIO()
+ result = io.load_file(media_file)
+ assert result == str(media_file)
+
+
+# ---------------------------------------------------------------------------
+# materialize_media_source
+# ---------------------------------------------------------------------------
+
+
+class TestMaterializeMediaSource:
+ """Tests for materialize_media_source — exercises real MediaConnector."""
+
+ def test_data_url_with_no_vllm_config(self):
+ data = b"\x00" * 50
+ b64 = base64.b64encode(data).decode()
+ url = f"data:audio/wav;base64,{b64}"
+ ctx = SimpleNamespace(vllm_config=None)
+
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=ctx,
+ ):
+ path = materialize_media_source(url)
+ try:
+ assert os.path.isfile(path)
+ finally:
+ _cleanup(path)
+
+ def test_data_url_materializes_to_temp_file(self):
+ data = b"\x00\x01\x02" * 50
+ b64 = base64.b64encode(data).decode()
+ url = f"data:audio/wav;base64,{b64}"
+
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=_fake_forward_context(),
+ ):
+ path = materialize_media_source(url)
+ try:
+ assert os.path.isfile(path)
+ with open(path, "rb") as f:
+ assert f.read() == data
+ finally:
+ _cleanup(path)
+
+ def test_local_path_with_allowed_path(self, tmp_path):
+ local_file = tmp_path / "media.bin"
+ local_file.write_bytes(b"\x00" * 50)
+ ctx = _fake_forward_context(allowed_local_media_path=str(tmp_path))
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=ctx,
+ ):
+ path = materialize_media_source(str(local_file))
+ assert path == str(local_file)
+
+ def test_local_path_without_allowed_path_raises(self, tmp_path):
+ local_file = tmp_path / "media.bin"
+ local_file.write_bytes(b"\x00" * 50)
+ ctx = _fake_forward_context(allowed_local_media_path="")
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=ctx,
+ ):
+ with pytest.raises(RuntimeError, match="allowed-local-media-path"):
+ materialize_media_source(str(local_file))
+
+ def test_http_url_rejected_by_domain_filter(self):
+ ctx = _fake_forward_context(allowed_domains=["trusted.example.com"])
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=ctx,
+ ):
+ with pytest.raises(ValueError, match="allowed domains"):
+ materialize_media_source("https://evil.internal/media.bin")
+
+ def test_file_uri_with_allowed_path(self, tmp_path):
+ media_file = tmp_path / "audio.wav"
+ media_file.write_bytes(b"\x00" * 50)
+ file_uri = media_file.as_uri()
+
+ ctx = _fake_forward_context(allowed_local_media_path=str(tmp_path))
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=ctx,
+ ):
+ path = materialize_media_source(file_uri)
+ assert path == str(media_file)
+
+ def test_file_uri_without_allowed_path_raises(self, tmp_path):
+ media_file = tmp_path / "audio.wav"
+ media_file.write_bytes(b"\x00" * 50)
+ file_uri = media_file.as_uri()
+
+ ctx = _fake_forward_context(allowed_local_media_path="")
+ with patch(
+ "vllm_omni.diffusion.forward_context.get_forward_context",
+ return_value=ctx,
+ ):
+ with pytest.raises(RuntimeError, match="allowed-local-media-path"):
+ materialize_media_source(file_uri)
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index 94b17b2b8..da97c4b94 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -67,6 +67,10 @@ class FakeAsyncOmni:
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
self.captured_prompt = None
self.captured_sampling_params_list = None
+ self.model_config = SimpleNamespace(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ )
async def generate(self, prompt, request_id, sampling_params_list):
self.captured_prompt = prompt
@@ -82,6 +86,12 @@ class BlockingVideoHandler:
self.stage_configs = None
self.started = threading.Event()
self.cancelled = threading.Event()
+ self._engine_client = SimpleNamespace(
+ model_config=SimpleNamespace(
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ ),
+ )
def set_stage_configs_if_missing(self, stage_configs):
if self.stage_configs is None:
@@ -1765,3 +1775,77 @@ def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker:
assert response.status_code == 200
assert response.content == b"fps-multiplied"
assert fps_values == [16]
+
+
+# ---------------------------------------------------------------------------
+# MediaConnector / SSRF protection (end-to-end through /v1/videos)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def restricted_test_client():
+ """Test client whose engine restricts media to trusted.example.com."""
+ engine = FakeAsyncOmni()
+ engine.model_config = SimpleNamespace(
+ allowed_local_media_path="",
+ allowed_media_domains=["trusted.example.com"],
+ )
+ app = FastAPI()
+ app.include_router(router)
+ app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
+ diffusion_engine=engine,
+ model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+ )
+ with TestClient(app) as client:
+ yield client
+
+
+class TestVideoEndpointSSRF:
+ """SSRF protection via /v1/videos endpoint with domain-restricted engine."""
+
+ def test_blocked_image_reference_url_returns_400(self, restricted_test_client):
+ response = restricted_test_client.post(
+ "/v1/videos",
+ data={
+ "prompt": "A fox.",
+ "image_reference": json.dumps({"image_url": "https://evil.internal/image.png"}),
+ },
+ )
+ assert response.status_code == 400
+ assert "allowed domains" in response.json()["detail"].lower()
+
+ def test_blocked_video_reference_url_returns_400(self, restricted_test_client):
+ response = restricted_test_client.post(
+ "/v1/videos",
+ data={
+ "prompt": "Continue.",
+ "video_reference": json.dumps({"video_url": "https://evil.internal/video.mp4"}),
+ },
+ )
+ assert response.status_code == 400
+ assert "allowed domains" in response.json()["detail"].lower()
+
+ def test_blocked_audio_reference_url_returns_400(self, restricted_test_client):
+ response = restricted_test_client.post(
+ "/v1/videos",
+ data={
+ "prompt": "Narrate.",
+ "audio_reference": json.dumps({"audio_url": "https://evil.internal/audio.mp3"}),
+ },
+ )
+ assert response.status_code == 400
+ assert "allowed domains" in response.json()["detail"].lower()
+
+ def test_data_url_bypasses_domain_restriction(self, restricted_test_client, mocker: MockerFixture):
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+ return_value=b"fake-video",
+ )
+ response = restricted_test_client.post(
+ "/v1/videos",
+ data={
+ "prompt": "A fox.",
+ "image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
+ },
+ )
+ assert response.status_code == 200
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 11186a0e7..fbfd3aafe 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
-import base64
import dataclasses
import io
import json
@@ -20,7 +19,6 @@ from numbers import Integral
from pathlib import Path
from typing import Annotated, Any, Literal, cast
-import httpx
import numpy as np
import vllm.envs as envs
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
@@ -85,6 +83,7 @@ from vllm.entrypoints.speech_to_text.translation.serving import (
OpenAIServingTranslation,
)
from vllm.logger import init_logger
+from vllm.multimodal.media import MediaConnector
from vllm.tasks import POOLING_TASKS
from vllm.tool_parsers import ToolParserManager
from vllm.utils import random_uuid
@@ -2038,17 +2037,23 @@ async def edit_images(
# Match the offline path: RGB normalize when the caller opts into
# Hunyuan-aware behavior. RGBA/P uploads otherwise diverge from offline.
normalize_edit_images_rgb = bot_task is not None or sys_type is not None
- pil_images = await _load_input_images(input_images_list, normalize_rgb=normalize_edit_images_rgb)
+ pil_images = await _load_input_images(
+ input_images_list,
+ engine_client.model_config,
+ normalize_rgb=normalize_edit_images_rgb,
+ )
prompt["multi_modal_data"] = {}
prompt["multi_modal_data"]["image"] = pil_images
if mask_image is not None:
# Mask role is different (alpha channel matters); never normalize.
- loaded = await _load_input_images([mask_image], normalize_rgb=False)
+ loaded = await _load_input_images([mask_image], engine_client.model_config, normalize_rgb=False)
prompt["multi_modal_data"]["mask_image"] = loaded[0]
if reference_image is not None:
- loaded = await _load_input_images([reference_image], normalize_rgb=normalize_edit_images_rgb)
+ loaded = await _load_input_images(
+ [reference_image], engine_client.model_config, normalize_rgb=normalize_edit_images_rgb
+ )
prompt["multi_modal_data"]["reference_image"] = loaded[0]
# 3 Build sample params
@@ -2535,41 +2540,41 @@ def _extract_images_from_result(result: Any) -> list[Any]:
async def _load_input_images(
inputs: list[str],
+ model_config: Any,
*,
normalize_rgb: bool = True,
) -> list[Image.Image]:
+ """Load images from data URIs, http(s) URLs, file:// URIs, or UploadFiles.
+
+ http(s) and ``file:`` URLs are fetched through vLLM's ``MediaConnector`` which
+ respects ``--allowed-media-domains`` and ``--allowed-local-media-path``
+ to prevent SSRF.
"""
- convert to PIL.Image.Image list
- """
+
if isinstance(inputs, str):
inputs = [inputs]
+ connector = MediaConnector(
+ allowed_local_media_path=model_config.allowed_local_media_path,
+ allowed_media_domains=model_config.allowed_media_domains,
+ )
+
images: list[Image.Image] = []
for inp in inputs:
- # 1. URL + base64
- if isinstance(inp, str) and inp.startswith("data:image"):
+ if isinstance(inp, str):
+ # http(s), data: and file: all handled by MediaConnector
try:
- _, b64_data = inp.split(",", 1)
- image_bytes = base64.b64decode(b64_data)
- img = Image.open(io.BytesIO(image_bytes))
- images.append(img)
+ # Use RGBA to preserve alpha for transparent PNG.
+ result = await connector.fetch_image_async(
+ inp,
+ image_mode="RGBA",
+ )
+ images.append(result.media)
except Exception as e:
- raise ValueError(f"Invalid base64 image: {e}")
-
- # 2. URL
- elif isinstance(inp, str) and inp.startswith("http"):
- async with httpx.AsyncClient(timeout=60) as client:
- try:
- resp = await client.get(inp)
- resp.raise_for_status()
- img = Image.open(io.BytesIO(resp.content))
- images.append(img)
- except Exception as e:
- raise ValueError(f"Failed to download image from URL {inp}: {e}")
-
- # 3. UploadFile
+ raise ValueError(f"Failed to load image: {e}") from e
elif hasattr(inp, "file"):
+ # UploadFile
try:
img_data = await inp.read()
img = Image.open(io.BytesIO(img_data))
@@ -3021,6 +3026,7 @@ async def _parse_video_form(
request.image_reference,
request.video_reference,
input_reference_bytes,
+ model_config=handler._engine_client.model_config,
max_video_frames=decode_spec.max_frames,
video_keep=decode_spec.keep,
)
@@ -3032,8 +3038,15 @@ async def _parse_video_form(
reference_audio: ReferenceAudio | None = None
if request.audio_reference is not None:
+ audio_connector = MediaConnector(
+ allowed_local_media_path=handler._engine_client.model_config.allowed_local_media_path,
+ allowed_media_domains=handler._engine_client.model_config.allowed_media_domains,
+ )
try:
- audio_path = await decode_audio_url(request.audio_reference.audio_url)
+ audio_path = await decode_audio_url(
+ request.audio_reference.audio_url,
+ audio_connector,
+ )
except InvalidInputReferenceError as exc:
raise HTTPException(400, detail=str(exc)) from exc
reference_audio = ReferenceAudio(path=audio_path)
diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py
index aa6479189..ee5791dee 100644
--- a/vllm_omni/entrypoints/openai/video_api_utils.py
+++ b/vllm_omni/entrypoints/openai/video_api_utils.py
@@ -10,12 +10,14 @@ import base64
import binascii
from collections import deque
from io import BytesIO
+from pathlib import Path
from typing import Any, Literal
-import httpx
import numpy as np
import torch
from PIL import Image, UnidentifiedImageError
+from vllm.multimodal.media import MediaConnector
+from vllm.multimodal.media.base import MediaIO
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -145,22 +147,12 @@ def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
raise InvalidInputReferenceError(f"Invalid {source}: image data is empty.")
-async def decode_image_url(image_url: str) -> Image.Image:
- if image_url.startswith("data:image"):
- return _decode_base64_image(image_url, source="image_reference.image_url")
-
- if image_url.startswith(("http://", "https://")):
- async with httpx.AsyncClient(timeout=60) as client:
- try:
- response = await client.get(image_url)
- response.raise_for_status()
- except httpx.HTTPError as exc:
- raise InvalidInputReferenceError(
- "Invalid image_reference.image_url: failed to download image."
- ) from exc
- return _decode_image_bytes(response.content, source="image_reference.image_url")
-
- raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")
+async def decode_image_url(image_url: str, connector: MediaConnector) -> Image.Image:
+ """Fetch and decode an image URL using MediaConnector for SSRF protection."""
+ try:
+ return (await connector.fetch_image_async(image_url)).media
+ except Exception as exc:
+ raise InvalidInputReferenceError(f"Invalid image_reference.image_url: {exc}") from exc
def _decode_base64_video(
@@ -184,96 +176,117 @@ def _decode_base64_video(
raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
-async def decode_video_url(
- video_url: str,
- *,
- max_frames: int | None = None,
- keep: Literal["first", "last"] = "first",
-) -> VideoFrames:
- if video_url.startswith("data:video"):
- return _decode_base64_video(
- video_url,
- source="video_reference.video_url",
- max_frames=max_frames,
- keep=keep,
- )
+class _VideoFramesMediaIO(MediaIO[VideoFrames]):
+ """MediaIO that decodes video bytes into VideoFrames with frame extraction."""
- if video_url.startswith(("http://", "https://")):
- async with httpx.AsyncClient(timeout=60) as client:
- try:
- response = await client.get(video_url)
- response.raise_for_status()
- except httpx.HTTPError as exc:
- raise InvalidInputReferenceError(
- "Invalid video_reference.video_url: failed to download video."
- ) from exc
+ def __init__(
+ self,
+ *,
+ source: str = "video_reference.video_url",
+ max_frames: int | None = None,
+ keep: Literal["first", "last"] = "first",
+ ) -> None:
+ self._source = source
+ self._max_frames = max_frames
+ self._keep = keep
+
+ def load_bytes(self, data: bytes) -> VideoFrames:
return _decode_video_bytes(
- response.content,
- source="video_reference.video_url",
- max_frames=max_frames,
- keep=keep,
+ data,
+ source=self._source,
+ max_frames=self._max_frames,
+ keep=self._keep,
)
- raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
+ def load_base64(self, media_type: str, data: str) -> VideoFrames:
+ return self.load_bytes(base64.b64decode(data))
+
+ def load_file(self, filepath: Path) -> VideoFrames:
+ return self.load_bytes(filepath.read_bytes())
-async def decode_audio_url(audio_url: str) -> str:
- """Decode an audio URL or data-URL to a temporary file path."""
- import tempfile
+class _AudioFileMediaIO(MediaIO[str]):
+ """MediaIO that writes audio bytes to a temp file, returns the path."""
- audio_bytes: bytes | None = None
+ def load_bytes(self, data: bytes) -> str:
+ return self._write_temp(data, suffix=".wav")
- if audio_url.startswith("data:audio"):
- _, b64_data = audio_url.split(",", 1)
- try:
- audio_bytes = base64.b64decode(b64_data)
- except (binascii.Error, ValueError) as exc:
- raise InvalidInputReferenceError(
- "Invalid audio_reference.audio_url: audio data is not valid base64."
- ) from exc
- elif audio_url.startswith(("http://", "https://")):
- async with httpx.AsyncClient(timeout=60) as client:
- try:
- response = await client.get(audio_url)
- response.raise_for_status()
- except httpx.HTTPError as exc:
- raise InvalidInputReferenceError(
- "Invalid audio_reference.audio_url: failed to download audio."
- ) from exc
- audio_bytes = response.content
- else:
- raise InvalidInputReferenceError("Invalid audio_reference.audio_url: must be an http(s) URL or data URL.")
-
- if not audio_bytes:
- raise InvalidInputReferenceError("Invalid audio_reference: audio data is empty.")
-
- suffix = ".wav"
- if audio_url.startswith("data:audio/"):
- mime = audio_url.split(";")[0].removeprefix("data:")
- ext = mime.split("/")[-1]
+ def load_base64(self, media_type: str, data: str) -> str:
+ suffix = ".wav"
+ ext = media_type.split("/")[-1] if media_type else ""
if ext in ("mpeg", "mp3"):
suffix = ".mp3"
elif ext == "wav":
suffix = ".wav"
elif ext.isalnum() and len(ext) <= 8:
suffix = f".{ext}"
+ return self._write_temp(base64.b64decode(data), suffix=suffix)
- tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
- tmp.write(audio_bytes)
- tmp.close()
- return tmp.name
+ def load_file(self, filepath: Path) -> str:
+ return str(filepath)
+
+ @staticmethod
+ def _write_temp(data: bytes, *, suffix: str) -> str:
+ import tempfile
+
+ if not data:
+ raise InvalidInputReferenceError("Invalid audio_reference: audio data is empty.")
+ tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+ tmp.write(data)
+ tmp.close()
+ return tmp.name
+
+
+async def decode_video_url(
+ video_url: str,
+ connector: MediaConnector,
+ *,
+ max_frames: int | None = None,
+ keep: Literal["first", "last"] = "first",
+) -> VideoFrames:
+ """Fetch and decode a video URL using MediaConnector for SSRF protection."""
+ try:
+ return await connector.load_from_url_async(
+ video_url,
+ _VideoFramesMediaIO(max_frames=max_frames, keep=keep),
+ )
+ except InvalidInputReferenceError:
+ raise
+ except Exception as exc:
+ raise InvalidInputReferenceError(f"Invalid video_reference.video_url: {exc}") from exc
+
+
+async def decode_audio_url(
+ audio_url: str,
+ connector: MediaConnector,
+) -> str:
+ """Fetch and decode an audio URL using MediaConnector for SSRF protection."""
+ try:
+ return await connector.load_from_url_async(
+ audio_url,
+ _AudioFileMediaIO(),
+ )
+ except InvalidInputReferenceError:
+ raise
+ except Exception as exc:
+ raise InvalidInputReferenceError(f"Invalid audio_reference.audio_url: {exc}") from exc
async def decode_input_reference(
image_reference: ImageReference | None,
video_reference: VideoReference | None,
input_reference_bytes: bytes | None,
+ model_config: Any,
*,
max_video_frames: int | None = None,
video_keep: Literal["first", "last"] = "first",
) -> Image.Image | VideoFrames | None:
- """Decode media input from multipart bytes, data URLs, or typed references."""
+ """Decode media input from multipart bytes, data URLs, or typed references.
+ http(s) image and video URLs are fetched through vLLM's
+ ``MediaConnector`` which respects ``--allowed-media-domains`` and
+ ``--allowed-local-media-path`` to prevent SSRF.
+ """
provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference))
if provided > 1:
raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.")
@@ -287,13 +300,22 @@ async def decode_input_reference(
)
if isinstance(image_reference, UrlImageReference):
- return await decode_image_url(image_reference.image_url)
+ connector = MediaConnector(
+ allowed_local_media_path=model_config.allowed_local_media_path,
+ allowed_media_domains=model_config.allowed_media_domains,
+ )
+ return await decode_image_url(image_reference.image_url, connector)
elif isinstance(image_reference, FileImageReference):
raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")
if isinstance(video_reference, UrlVideoReference):
+ connector = MediaConnector(
+ allowed_local_media_path=model_config.allowed_local_media_path,
+ allowed_media_domains=model_config.allowed_media_domains,
+ )
return await decode_video_url(
video_reference.video_url,
+ connector,
max_frames=max_video_frames,
keep=video_keep,
)
diff --git a/vllm_omni/transformers_utils/processors/audiox.py b/vllm_omni/transformers_utils/processors/audiox.py
index b0e53cfc7..8839ffac9 100644
--- a/vllm_omni/transformers_utils/processors/audiox.py
+++ b/vllm_omni/transformers_utils/processors/audiox.py
@@ -14,9 +14,9 @@ from __future__ import annotations
import base64
import os
import tempfile
+from pathlib import Path
from typing import Any
from urllib.parse import urlparse
-from urllib.request import urlopen
import av
import numpy as np
@@ -26,6 +26,8 @@ import torch.nn.functional as F
import torchaudio.functional as taF
from einops import rearrange
from torchvision.io import read_image
+from vllm.multimodal.media import MediaConnector
+from vllm.multimodal.media.base import MediaIO
# AudioX task taxonomy. Tasks beginning with "v" require a video input; tasks containing
# "t" carry a text prompt. tv2*/v2* share the same conditioning pathways.
@@ -51,29 +53,51 @@ def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]:
return out
-def materialize_media_source(source: str) -> str:
- """Return a local filesystem path for ``source``.
+class _TempFileMediaIO(MediaIO[str]):
+ """MediaIO that writes fetched bytes to a temp file, returns the path."""
- Accepts a local path, a ``data:;base64,...`` URI, or an ``http(s)://`` URL.
- Anything non-local is fetched into a NamedTemporaryFile and that path is returned;
- callers don't need to clean the tempfile up (the OS does on exit).
- """
- if source.startswith("data:"):
- _, _, payload = source.partition(",")
- raw = base64.b64decode(payload)
- f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
- f.write(raw)
- f.close()
- return f.name
- parsed = urlparse(source)
- if parsed.scheme in ("http", "https"):
- with urlopen(source) as resp:
- data = resp.read()
+ def load_bytes(self, data: bytes) -> str:
+ return self._write_temp(data)
+
+ def load_base64(self, media_type: str, data: str) -> str:
+ return self._write_temp(base64.b64decode(data))
+
+ def load_file(self, filepath: Path) -> str:
+ return str(filepath)
+
+ @staticmethod
+ def _write_temp(data: bytes) -> str:
f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
f.write(data)
f.close()
return f.name
- return source
+
+
+def _get_media_connector() -> MediaConnector:
+ """Build a MediaConnector from the current forward context's model config."""
+ from vllm_omni.diffusion.forward_context import get_forward_context
+
+ ctx = get_forward_context()
+ model_config = ctx.vllm_config.model_config if ctx.vllm_config else None
+ return MediaConnector(
+ allowed_local_media_path=getattr(model_config, "allowed_local_media_path", ""),
+ allowed_media_domains=getattr(model_config, "allowed_media_domains", None),
+ )
+
+
+def materialize_media_source(source: str) -> str:
+ """Return a local filesystem path for ``source``.
+
+ Accepts a local path, a ``data:;base64,...`` URI, or an ``http(s)://`` URL.
+ All inputs are routed through vLLM's ``MediaConnector`` which respects
+ ``--allowed-media-domains`` and ``--allowed-local-media-path`` to prevent
+ SSRF. Bare local paths are converted to ``file://`` URIs so that the
+ same path restriction applies.
+ """
+ if not urlparse(source).scheme:
+ source = Path(source).resolve().as_uri()
+ connector = _get_media_connector()
+ return connector.load_from_url(source, _TempFileMediaIO())
def _load_video_path_pyav(