diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py
index fbf56b70e..2a5c5be47 100644
--- a/tests/entrypoints/openai_api/test_decode_audio_url.py
+++ b/tests/entrypoints/openai_api/test_decode_audio_url.py
@@ -1,19 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for decode_audio_url helper and _AudioFileMediaIO."""
+"""Unit tests for decode_audio_url helper."""
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 (
- _AudioFileMediaIO,
- decode_audio_url,
-)
+from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -27,127 +24,15 @@ 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 — no mocks, real code path."""
+ """Tests for base64 data URL decoding."""
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, _permissive_connector()))
+ path = _run(decode_audio_url(url))
try:
assert os.path.isfile(path)
assert path.endswith(".mp3")
@@ -161,7 +46,7 @@ class TestDecodeAudioUrlDataUrl:
b64 = base64.b64encode(audio_bytes).decode()
url = f"data:audio/wav;base64,{b64}"
- path = _run(decode_audio_url(url, _permissive_connector()))
+ path = _run(decode_audio_url(url))
try:
assert path.endswith(".wav")
with open(path, "rb") as f:
@@ -169,103 +54,134 @@ 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="Incorrect padding"):
- _run(decode_audio_url(url, _permissive_connector()))
+
+ with pytest.raises(InvalidInputReferenceError, match="not valid base64"):
+ _run(decode_audio_url(url))
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, _permissive_connector()))
+ _run(decode_audio_url(url))
-# ---------------------------------------------------------------------------
-# decode_audio_url with file:// URIs (exercises load_file through 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)
-class TestDecodeAudioUrlFileUri:
- """Tests using file:// URIs — exercises MediaConnector + load_file."""
+class TestDecodeAudioUrlHttp:
+ """Tests for HTTP URL decoding."""
- 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)
+ 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_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))
+ async def _mock_get(self, url):
+ return fake_resp
- 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))
+ monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
+ monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None)
+ 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)
-# ---------------------------------------------------------------------------
-# decode_audio_url — invalid schemes and SSRF
-# ---------------------------------------------------------------------------
+ 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"))
class TestDecodeAudioUrlInvalid:
"""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_audio_url("/path/to/local/file.mp3", _permissive_connector()))
+ with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
+ _run(decode_audio_url("/path/to/local/file.mp3"))
def test_ftp_url_raises(self):
- 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()))
+ with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
+ _run(decode_audio_url("ftp://example.com/file.mp3"))
def test_empty_string_raises(self):
- 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)
+ with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
+ _run(decode_audio_url(""))
diff --git a/tests/entrypoints/openai_api/test_decode_video_url.py b/tests/entrypoints/openai_api/test_decode_video_url.py
deleted file mode 100644
index 64ff7e2e0..000000000
--- a/tests/entrypoints/openai_api/test_decode_video_url.py
+++ /dev/null
@@ -1,210 +0,0 @@
-# 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 49c43d7a9..31a82102a 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -13,7 +13,6 @@ 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
@@ -25,19 +24,12 @@ 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,
- _load_input_images,
- router,
-)
+from vllm_omni.entrypoints.openai.api_server import _check_max_generated_image_size, _DiffusionServingModels, 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
@@ -190,7 +182,6 @@ 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
@@ -280,10 +271,6 @@ 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)
@@ -348,10 +335,6 @@ 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)
@@ -416,10 +399,6 @@ 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)
@@ -456,24 +435,9 @@ 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
@@ -765,10 +729,6 @@ 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()
@@ -2176,107 +2136,3 @@ 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
deleted file mode 100644
index a0ce824b3..000000000
--- a/tests/entrypoints/openai_api/test_materialize_media_source.py
+++ /dev/null
@@ -1,173 +0,0 @@
-# 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 c4ac824da..3f1c5f8e5 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -68,10 +68,6 @@ 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
@@ -87,12 +83,6 @@ 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:
@@ -1869,77 +1859,3 @@ 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 2c6a3d87a..3560f26b8 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
+import base64
import dataclasses
import io
import json
@@ -19,6 +20,7 @@ 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
@@ -83,7 +85,6 @@ 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
@@ -2036,23 +2037,17 @@ 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,
- engine_client.model_config,
- normalize_rgb=normalize_edit_images_rgb,
- )
+ pil_images = await _load_input_images(input_images_list, 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], engine_client.model_config, normalize_rgb=False)
+ loaded = await _load_input_images([mask_image], normalize_rgb=False)
prompt["multi_modal_data"]["mask_image"] = loaded[0]
if reference_image is not None:
- loaded = await _load_input_images(
- [reference_image], engine_client.model_config, normalize_rgb=normalize_edit_images_rgb
- )
+ loaded = await _load_input_images([reference_image], normalize_rgb=normalize_edit_images_rgb)
prompt["multi_modal_data"]["reference_image"] = loaded[0]
# 3 Build sample params
@@ -2539,41 +2534,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:
- if isinstance(inp, str):
- # http(s), data: and file: all handled by MediaConnector
+ # 1. URL + base64
+ if isinstance(inp, str) and inp.startswith("data:image"):
try:
- # Use RGBA to preserve alpha for transparent PNG.
- result = await connector.fetch_image_async(
- inp,
- image_mode="RGBA",
- )
- images.append(result.media)
+ _, b64_data = inp.split(",", 1)
+ image_bytes = base64.b64decode(b64_data)
+ img = Image.open(io.BytesIO(image_bytes))
+ images.append(img)
except Exception as e:
- raise ValueError(f"Failed to load image: {e}") from 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
elif hasattr(inp, "file"):
- # UploadFile
try:
img_data = await inp.read()
img = Image.open(io.BytesIO(img_data))
@@ -3025,7 +3020,6 @@ 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,
)
@@ -3037,15 +3031,8 @@ 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_connector,
- )
+ audio_path = await decode_audio_url(request.audio_reference.audio_url)
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 ee5791dee..aa6479189 100644
--- a/vllm_omni/entrypoints/openai/video_api_utils.py
+++ b/vllm_omni/entrypoints/openai/video_api_utils.py
@@ -10,14 +10,12 @@ 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 (
@@ -147,12 +145,22 @@ 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, 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
+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.")
def _decode_base64_video(
@@ -176,117 +184,96 @@ def _decode_base64_video(
raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
-class _VideoFramesMediaIO(MediaIO[VideoFrames]):
- """MediaIO that decodes video bytes into VideoFrames with frame extraction."""
-
- 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(
- data,
- source=self._source,
- max_frames=self._max_frames,
- keep=self._keep,
+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,
)
- def load_base64(self, media_type: str, data: str) -> VideoFrames:
- return self.load_bytes(base64.b64decode(data))
+ 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
+ return _decode_video_bytes(
+ response.content,
+ source="video_reference.video_url",
+ max_frames=max_frames,
+ keep=keep,
+ )
- def load_file(self, filepath: Path) -> VideoFrames:
- return self.load_bytes(filepath.read_bytes())
+ raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
-class _AudioFileMediaIO(MediaIO[str]):
- """MediaIO that writes audio bytes to a temp file, returns the path."""
+async def decode_audio_url(audio_url: str) -> str:
+ """Decode an audio URL or data-URL to a temporary file path."""
+ import tempfile
- def load_bytes(self, data: bytes) -> str:
- return self._write_temp(data, suffix=".wav")
+ audio_bytes: bytes | None = None
- def load_base64(self, media_type: str, data: str) -> str:
- suffix = ".wav"
- ext = media_type.split("/")[-1] if media_type else ""
+ 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]
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)
- 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
+ tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+ tmp.write(audio_bytes)
+ tmp.close()
+ return tmp.name
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.")
@@ -300,22 +287,13 @@ async def decode_input_reference(
)
if isinstance(image_reference, UrlImageReference):
- 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)
+ return await decode_image_url(image_reference.image_url)
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 8839ffac9..b0e53cfc7 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,8 +26,6 @@ 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.
@@ -53,51 +51,29 @@ def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]:
return out
-class _TempFileMediaIO(MediaIO[str]):
- """MediaIO that writes fetched bytes to a temp file, returns the path."""
-
- 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
-
-
-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.
+ 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 not urlparse(source).scheme:
- source = Path(source).resolve().as_uri()
- connector = _get_media_connector()
- return connector.load_from_url(source, _TempFileMediaIO())
+ 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()
+ f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
+ f.write(data)
+ f.close()
+ return f.name
+ return source
def _load_video_path_pyav(