From 7b876324c24d129d6bfa1b5778d8169b2e42bbde Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 24 Apr 2026 22:54:35 +0800 Subject: [PATCH] [Feature] Streaming video input with EVS frame filtering (RFC #2201 Phase 2-4) (#2342) Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: shimei.jmh Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- docs/.nav.yml | 1 + docs/serving/video_stream_api.md | 93 ++ .../qwen3_omni/streaming_video_client.py | 208 ++++ .../entrypoints/openai_api/conftest_video.py | 31 + .../openai_api/test_serving_video_stream.py | 650 ++++++++++++ .../openai_api/test_video_frame_filter.py | 89 ++ .../openai_api/test_video_stream_handler.py | 593 +++++++++++ .../openai_api/test_video_stream_session.py | 289 ++++++ vllm_omni/entrypoints/openai/api_server.py | 27 + .../openai/serving_video_stream.py | 971 ++++++++++++++++++ .../entrypoints/openai/video_frame_filter.py | 116 +++ .../openai/video_stream_context.py | 28 + .../entrypoints/openai/video_stream_envs.py | 62 ++ .../openai/video_stream_session.py | 439 ++++++++ 14 files changed, 3597 insertions(+) create mode 100644 docs/serving/video_stream_api.md create mode 100644 examples/online_serving/qwen3_omni/streaming_video_client.py create mode 100644 tests/entrypoints/openai_api/conftest_video.py create mode 100644 tests/entrypoints/openai_api/test_serving_video_stream.py create mode 100644 tests/entrypoints/openai_api/test_video_frame_filter.py create mode 100644 tests/entrypoints/openai_api/test_video_stream_handler.py create mode 100644 tests/entrypoints/openai_api/test_video_stream_session.py create mode 100644 vllm_omni/entrypoints/openai/serving_video_stream.py create mode 100644 vllm_omni/entrypoints/openai/video_frame_filter.py create mode 100644 vllm_omni/entrypoints/openai/video_stream_context.py create mode 100644 vllm_omni/entrypoints/openai/video_stream_envs.py create mode 100644 vllm_omni/entrypoints/openai/video_stream_session.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 8b0f9a023..b44e8e6b5 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -10,6 +10,7 @@ nav: - Image Generation: serving/image_generation_api.md - Image Edit: serving/image_edit_api.md - Text to Speech: serving/speech_api.md + - Streaming Video Input: serving/video_stream_api.md - Examples: - examples/README.md - Offline Inference: diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md new file mode 100644 index 000000000..88f74affc --- /dev/null +++ b/docs/serving/video_stream_api.md @@ -0,0 +1,93 @@ +# Streaming Video Input API + +vLLM-Omni provides a WebSocket API for streaming video frames and optional audio chunks into Qwen3-Omni, then asking questions over the buffered session context. + +Each server instance runs a single model specified at startup with `vllm serve --omni`. + +## Quick Start + +### Start the Server + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ + --deploy-config vllm_omni/deploy/qwen3_omni.yaml \ + --omni \ + --port 8000 \ + --trust-remote-code +``` + +### Run the Example Client + +```bash +python examples/online_serving/qwen3_omni/streaming_video_client.py \ + --url ws://localhost:8000/v1/video/chat/stream \ + --video /path/to/video.mp4 \ + --query "Describe what is happening in the video." +``` + +## API Reference + +### Endpoint + +```text +WebSocket /v1/video/chat/stream +``` + +### Protocol + +| Direction | Type | Required fields | Description | +|-----------|------|-----------------|-------------| +| Client -> Server | `session.config` | none | First message. Configures output modalities, frame sampling, EVS, and prompts. | +| Client -> Server | `video.frame` | `data` | Base64 JPEG/PNG frame. | +| Client -> Server | `audio.chunk` | `data` | Base64 PCM16 16 kHz mono audio bytes. | +| Client -> Server | `video.query` | `text` | Ask a question over the buffered frames and audio. | +| Client -> Server | `video.done` | none | End the WebSocket session. | +| Server -> Client | `response.start` | none | Query generation started. | +| Server -> Client | `response.text.delta` | `delta` | Incremental text output. | +| Server -> Client | `response.text.done` | `text` | Final text output for the query. | +| Server -> Client | `response.audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. | +| Server -> Client | `response.audio.done` | none | Audio output finished. | +| Server -> Client | `session.done` | none | Session closed. | +| Server -> Client | `error` | `message` | Recoverable protocol or generation error. | + +### `session.config` Fields + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. | +| `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. | +| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. | +| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. | +| `system_prompt` | string or null | null | Optional custom system prompt. | +| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. | +| `sampling_params_list` | list or null | null | Optional per-stage sampling parameter overrides. | +| `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. | +| `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. | + +### Legacy Aliases + +The server accepts these legacy field names and rewrites them before validation. New clients should send the canonical names above. + +| Legacy field | Canonical field | +|--------------|-----------------| +| `num_sample_frames` | `num_frames` | +| `evs_enabled` | `enable_frame_filter` | +| `evs_threshold` | `frame_filter_threshold` | + +### Environment Variables + +| Variable | Values | Default | Description | +|----------|--------|---------|-------------| +| `VLLM_VIDEO_ASYNC_CHUNK` | `on`, `off` | `on` | Wire-level streaming switch. `off` buffers server-side deltas and emits coalesced outputs at the end of a query. | +| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Audio delta extraction strategy. `fast` emits only newly produced chunks; `slow` recomputes from accumulated audio and exists for A/B verification. | + +## EVS Semantics + +EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames. + +## Known Limitations + +- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers. +- Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts. +- If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session. +- The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments. diff --git a/examples/online_serving/qwen3_omni/streaming_video_client.py b/examples/online_serving/qwen3_omni/streaming_video_client.py new file mode 100644 index 000000000..58f26d245 --- /dev/null +++ b/examples/online_serving/qwen3_omni/streaming_video_client.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Example WebSocket client for the /v1/video/chat/stream endpoint. + +Sends video frames from a local file (or generates synthetic ones), submits a +query, and prints the streamed text response. + +Requirements: + pip install websockets pillow + +Usage: + # With a video file (requires opencv-python): + python streaming_video_client.py --video my_clip.mp4 \\ + --query "What is happening in this video?" + + # Synthetic frames (no extra deps): + python streaming_video_client.py \\ + --query "Describe what you see." \\ + --synthetic-frames 10 + + # With audio (Phase 3): + python streaming_video_client.py --video my_clip.mp4 \\ + --audio my_audio.pcm \\ + --query "What is the person saying and doing?" +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import io +import json +import sys + +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + sys.exit(1) + +from PIL import Image + + +def _generate_synthetic_frame(index: int, width: int = 320, height: int = 240) -> bytes: + """Generate a simple synthetic JPEG frame with a colour gradient.""" + r = (index * 37) % 256 + g = (index * 73) % 256 + b = (index * 113) % 256 + img = Image.new("RGB", (width, height), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=80) + return buf.getvalue() + + +def _load_video_frames(path: str, max_frames: int = 64, fps: int = 2) -> list[bytes]: + """Extract frames from a video file using OpenCV.""" + try: + import cv2 + except ImportError: + print("opencv-python is required to read video files: pip install opencv-python") + sys.exit(1) + + cap = cv2.VideoCapture(path) + if not cap.isOpened(): + print(f"Cannot open video: {path}") + sys.exit(1) + + video_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + frame_interval = max(1, int(video_fps / fps)) + + frames: list[bytes] = [] + idx = 0 + while len(frames) < max_frames: + ret, frame = cap.read() + if not ret: + break + if idx % frame_interval == 0: + _, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) + frames.append(buf.tobytes()) + idx += 1 + + cap.release() + print(f"Loaded {len(frames)} frames from {path} (interval={frame_interval})") + return frames + + +async def run(args: argparse.Namespace) -> None: + uri = f"ws://{args.host}:{args.port}/v1/video/chat/stream" + + # Prepare frames + if args.video: + frames = _load_video_frames(args.video, max_frames=args.max_frames, fps=args.fps) + else: + frames = [_generate_synthetic_frame(i) for i in range(args.synthetic_frames)] + print(f"Generated {len(frames)} synthetic frames") + + # Prepare audio (optional, Phase 3) + audio_data: bytes | None = None + if args.audio: + with open(args.audio, "rb") as f: + audio_data = f.read() + print(f"Loaded audio: {len(audio_data)} bytes") + + async with websockets.connect(uri, max_size=16 * 1024 * 1024) as ws: + # 1. Send session.config + config = { + "type": "session.config", + "model": args.model, + "modalities": ["text", "audio"] if audio_data else ["text"], + "max_frames": args.max_frames, + "num_frames": args.num_sample_frames, + "enable_frame_filter": args.evs, + "frame_filter_threshold": args.evs_threshold, + "use_audio_in_video": bool(audio_data), + } + await ws.send(json.dumps(config)) + print(f"Sent session.config: model={args.model} evs={args.evs}") + + # 2. Send frames + for i, frame in enumerate(frames): + msg = { + "type": "video.frame", + "data": base64.b64encode(frame).decode(), + } + await ws.send(json.dumps(msg)) + if (i + 1) % 10 == 0: + print(f" Sent {i + 1}/{len(frames)} frames") + print(f"Sent all {len(frames)} frames") + + # 3. Send audio chunks (Phase 3) + if audio_data: + chunk_size = 16000 * 2 # 1 second of 16 kHz 16-bit PCM + for offset in range(0, len(audio_data), chunk_size): + chunk = audio_data[offset : offset + chunk_size] + msg = { + "type": "audio.chunk", + "data": base64.b64encode(chunk).decode(), + } + await ws.send(json.dumps(msg)) + print(f"Sent audio in {(len(audio_data) + chunk_size - 1) // chunk_size} chunks") + + # 4. Send query, then immediately send video.done so the server + # knows the session is complete (avoids deadlock where client + # waits for session.done while server waits for video.done). + await ws.send(json.dumps({"type": "video.query", "text": args.query})) + print(f"\nQuery: {args.query}") + print("Response: ", end="", flush=True) + + # Signal end of session right after the query. The server will + # process the query first (it's already queued), then handle + # video.done and reply with session.done. + await ws.send(json.dumps({"type": "video.done"})) + + # 5. Receive response until session.done + recv_timeout = 120 # seconds — avoid infinite hang if server stalls + while True: + raw = await asyncio.wait_for(ws.recv(), timeout=recv_timeout) + data = json.loads(raw) + msg_type = data.get("type") + + if msg_type == "response.text.delta": + print(data.get("delta", ""), end="", flush=True) + elif msg_type == "response.text.done": + print() # newline + elif msg_type == "response.evs_stats": + retained = data.get("retained_count", 0) + dropped = data.get("dropped_count", 0) + rate = data.get("drop_rate", 0) + print(f"\nEVS stats: retained={retained} dropped={dropped} drop_rate={rate:.1%}") + elif msg_type == "session.done": + print("Session complete.") + break + elif msg_type == "error": + print(f"\nError: {data.get('message')}") + break + elif msg_type == "response.start": + pass # expected + else: + print(f"\n[unknown message] {data}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Streaming video chat client") + parser.add_argument("--host", default="localhost") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--model", default="Qwen/Qwen3-Omni-MoE") + parser.add_argument("--video", help="Path to video file (requires opencv-python)") + parser.add_argument("--audio", help="Path to raw PCM 16kHz audio file (Phase 3)") + parser.add_argument("--query", default="What do you see in this video?") + parser.add_argument( + "--synthetic-frames", type=int, default=10, help="Number of synthetic frames if --video is not set" + ) + parser.add_argument("--max-frames", type=int, default=64) + parser.add_argument("--num-sample-frames", type=int, default=16) + parser.add_argument("--fps", type=int, default=2, help="Frame extraction rate from video") + parser.add_argument( + "--no-evs", dest="evs", action="store_false", help="Disable EVS frame filtering (enabled by default)" + ) + parser.set_defaults(evs=True) + parser.add_argument("--evs-threshold", type=float, default=0.95) + args = parser.parse_args() + asyncio.run(run(args)) + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/openai_api/conftest_video.py b/tests/entrypoints/openai_api/conftest_video.py new file mode 100644 index 000000000..3d5e0d855 --- /dev/null +++ b/tests/entrypoints/openai_api/conftest_video.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared test helpers for video streaming tests.""" + +from __future__ import annotations + +import io + +import pytest + +np = pytest.importorskip("numpy", reason="numpy required for video stream tests") +PIL = pytest.importorskip("PIL", reason="Pillow required for video stream tests") +from PIL import Image # noqa: E402 + + +def make_jpeg(r: int = 128, g: int = 128, b: int = 128, size: int = 64) -> bytes: + """Create a solid-colour JPEG image.""" + img = Image.new("RGB", (size, size), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def make_gradient_jpeg(seed: int, size: int = 64) -> bytes: + """Create a random-gradient JPEG that varies based on *seed*.""" + rng = np.random.RandomState(seed) + arr = rng.randint(0, 256, (size, size, 3), dtype=np.uint8) + img = Image.fromarray(arr, "RGB") + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py new file mode 100644 index 000000000..0787eb856 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -0,0 +1,650 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the serving-layer streaming video WebSocket handler.""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import json +import threading +from typing import Any + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.openai import serving_video_stream, video_stream_envs +from vllm_omni.entrypoints.openai.serving_video_stream import ( + OmniStreamingVideoHandler, + StreamingVideoSessionConfig, +) +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes: + img = Image.new("RGB", (64, 64), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def _b64(data: bytes) -> str: + return base64.b64encode(data).decode() + + +def _text_result(text: str) -> OmniRequestOutput: + class Output: + pass + + class RequestOutput: + pass + + output = Output() + output.text = text + request_output = RequestOutput() + request_output.outputs = [output] + return OmniRequestOutput(final_output_type="text", request_output=request_output) + + +def _audio_result(audio_data: Any) -> OmniRequestOutput: + class Output: + pass + + class RequestOutput: + pass + + output = Output() + output.multimodal_output = {"audio": audio_data} + request_output = RequestOutput() + request_output.outputs = [output] + return OmniRequestOutput(final_output_type="audio", request_output=request_output) + + +class MockWebSocket: + def __init__(self, messages: list[str] | None = None): + self._messages = list(messages or []) + self._idx = 0 + self.accepted = False + self.sent: list[dict[str, Any]] = [] + + async def accept(self): + self.accepted = True + + async def receive_text(self) -> str: + if self._idx >= len(self._messages): + await asyncio.sleep(999) + msg = self._messages[self._idx] + self._idx += 1 + return msg + + async def send_json(self, data: dict[str, Any]): + self.sent.append(data) + + +class TimedWebSocket: + def __init__(self): + self._q: asyncio.Queue[str] = asyncio.Queue() + self.accepted = False + self.sent: list[dict[str, Any]] = [] + + async def accept(self): + self.accepted = True + + async def receive_text(self) -> str: + return await self._q.get() + + async def send_json(self, data: dict[str, Any]): + self.sent.append(data) + + def put(self, msg: dict[str, Any]): + self._q.put_nowait(json.dumps(msg)) + + def sent_types(self) -> list[str]: + return [m.get("type", "") for m in self.sent] + + +def test_api_server_registers_video_stream_route(): + from vllm_omni.entrypoints.openai.api_server import router + + assert any(getattr(route, "path", None) == "/v1/video/chat/stream" for route in router.routes) + + +@pytest.mark.asyncio +async def test_receive_config_accepts_client_legacy_aliases(): + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "num_sample_frames": 7, + "evs_enabled": False, + "evs_threshold": 0.87, + } + ) + ] + ) + handler = OmniStreamingVideoHandler(chat_service=object()) + + config = await handler._receive_config(ws) + + assert config is not None + assert config.num_frames == 7 + assert config.enable_frame_filter is False + assert config.frame_filter_threshold == 0.87 + + +@pytest.mark.asyncio +async def test_audio_in_video_sets_mm_processor_kwargs(): + captured_requests = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + captured_requests.append(request) + return {"prompt": "x"} + + ws = MockWebSocket() + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=True) + + await handler._process_query_engine( + ws, + config, + [_b64(_make_jpeg())], + bytearray(b"\x00\x00"), + [], + "what is happening?", + "req-1", + asyncio.Event(), + {}, + ) + + assert captured_requests + assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True} + + +@pytest.mark.asyncio +async def test_audio_in_video_disabled_omits_mm_processor_kwargs(): + captured_requests = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + captured_requests.append(request) + return {"prompt": "x"} + + ws = MockWebSocket() + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=False) + + await handler._process_query_engine( + ws, + config, + [_b64(_make_jpeg())], + bytearray(b"\x00\x00"), + [], + "what is happening?", + "req-1", + asyncio.Event(), + {}, + ) + + assert captured_requests + assert captured_requests[0].mm_processor_kwargs is None + + +@pytest.mark.asyncio +async def test_query_inline_audio_data_sets_mm_processor_kwargs(): + captured_requests = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + captured_requests.append(request) + return {"prompt": "x"} + + ws = MockWebSocket( + [ + json.dumps({"type": "session.config", "model": "test"}), + json.dumps({"type": "video.frame", "data": _b64(_make_jpeg())}), + json.dumps( + { + "type": "video.query", + "text": "describe", + "audio_data": _b64(b"\x00\x00"), + } + ), + json.dumps({"type": "video.done"}), + ] + ) + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=2.0) + + await handler.handle_session(ws) + + assert captured_requests + assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True} + assert "session.done" in [m.get("type") for m in ws.sent] + + +def test_audio_delta_mode_is_read_by_serving_code_at_runtime(monkeypatch): + handler = OmniStreamingVideoHandler(chat_service=object()) + result = _audio_result([object()]) + + monkeypatch.setattr( + OmniStreamingVideoHandler, + "_delta_fast", + classmethod(lambda cls, audio_data, chunks_drained: ("fast-path", chunks_drained)), + ) + monkeypatch.setattr( + OmniStreamingVideoHandler, + "_delta_slow", + classmethod(lambda cls, audio_data, chunks_drained: ("slow-path", chunks_drained)), + ) + + monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "fast") + assert handler._extract_audio_delta_b64(result, 0)[0] == "fast-path" + + monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "slow") + assert handler._extract_audio_delta_b64(result, 0)[0] == "slow-path" + + +def test_video_stream_envs_strip_and_warn_once_per_invalid_value(monkeypatch): + warnings = [] + + video_stream_envs._warned_invalid_envs.clear() + try: + monkeypatch.setattr( + video_stream_envs.logger, + "warning", + lambda message, *args, **_kwargs: warnings.append((message, args)), + ) + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", " off ") + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "off" + assert not warnings + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "bad") + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on" + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on" + assert len(warnings) == 1 + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "still_bad") + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on" + assert len(warnings) == 2 + finally: + video_stream_envs._warned_invalid_envs.clear() + + +@pytest.mark.asyncio +async def test_async_chunk_mode_is_read_by_engine_path_at_runtime(monkeypatch): + class TextEngine: + def generate(self, **_kwargs): + async def _gen(): + yield _text_result("hello") + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt": "x"} + + handler = CapturingHandler(chat_service=object(), engine_client=TextEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text"]) + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "on") + ws_on = MockWebSocket() + await handler._process_query_engine( + ws_on, + config, + [_b64(_make_jpeg())], + bytearray(), + [], + "describe", + "req-on", + asyncio.Event(), + {}, + ) + assert {"type": "response.text.delta", "delta": "hello"} in ws_on.sent + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "off") + ws_off = MockWebSocket() + await handler._process_query_engine( + ws_off, + config, + [_b64(_make_jpeg())], + bytearray(), + [], + "describe", + "req-off", + asyncio.Event(), + {}, + ) + assert {"type": "response.text.done", "text": "hello"} in ws_off.sent + assert not any(m.get("type") == "response.text.delta" for m in ws_off.sent) + + +@pytest.mark.asyncio +async def test_query_without_engine_client_sends_error(): + ws = MockWebSocket() + handler = OmniStreamingVideoHandler(chat_service=object(), engine_client=None) + + await handler._process_query( + ws, + StreamingVideoSessionConfig(model="test"), + [], + bytearray(), + [], + "describe", + "req-1", + asyncio.Event(), + {}, + ) + + assert {"type": "error", "message": "Streaming video requires an engine client"} in ws.sent + + +@pytest.mark.asyncio +async def test_new_query_cancels_in_flight_query(): + query_started = asyncio.Event() + query_cancelled = asyncio.Event() + calls = 0 + + class BlockingHandler(OmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + nonlocal calls + calls += 1 + if calls > 1: + return + query_started.set() + try: + await asyncio.sleep(999) + except asyncio.CancelledError: + query_cancelled.set() + raise + + ws = TimedWebSocket() + handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + ws.put({"type": "video.query", "text": "interrupt"}) + await asyncio.wait_for(query_cancelled.wait(), timeout=2.0) + ws.put({"type": "video.done"}) + + await asyncio.wait_for(task, timeout=2.0) + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_video_done_waits_for_in_flight_query(): + query_started = asyncio.Event() + allow_finish = asyncio.Event() + query_finished = asyncio.Event() + + class BlockingHandler(OmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + query_started.set() + await allow_finish.wait() + query_finished.set() + + ws = TimedWebSocket() + handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + ws.put({"type": "video.done"}) + await asyncio.sleep(0.05) + assert not task.done() + assert not query_finished.is_set() + + allow_finish.set() + await asyncio.wait_for(task, timeout=2.0) + + assert query_finished.is_set() + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_frame_prewarm_does_not_block_following_query(monkeypatch): + decode_started = threading.Event() + release_decode = threading.Event() + query_started = asyncio.Event() + + def blocked_decode(raw_bytes: bytes): + decode_started.set() + release_decode.wait(timeout=2.0) + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + class BlockingHandler(OmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + query_started.set() + + monkeypatch.setattr(serving_video_stream, "_decode_frame_bytes", blocked_decode) + + ws = TimedWebSocket() + handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + + for _ in range(100): + if decode_started.is_set(): + break + await asyncio.sleep(0.01) + assert decode_started.is_set() + + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + release_decode.set() + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_client_cannot_send_internal_frame_decode_failed_message(): + captured_frames: list[list[str]] = [] + frame = _b64(_make_jpeg()) + + class CapturingHandler(OmniStreamingVideoHandler): + async def _process_query( + self, + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ): + captured_frames.append(list(frame_buffer)) + + ws = TimedWebSocket() + handler = CapturingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": frame}) + await asyncio.sleep(0) + ws.put({"type": "_internal.frame_decode_failed", "b64": frame}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "Unknown type: _internal.frame_decode_failed"} in ws.sent + assert captured_frames == [[frame]] + + +@pytest.mark.asyncio +async def test_failed_frame_prewarm_removes_frame_before_query(): + ws = TimedWebSocket() + handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test", "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(b"not-a-jpeg")}) + + for _ in range(100): + if any(m.get("message") == "Frame decode failed" for m in ws.sent): + break + await asyncio.sleep(0.01) + + assert {"type": "error", "message": "Frame decode failed"} in ws.sent + + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "No frames buffered"} in ws.sent + + +@pytest.mark.asyncio +async def test_frame_filter_error_sends_invalid_image(monkeypatch): + def fail_should_retain(self, frame_jpeg): + raise ValueError("decode failed") + + monkeypatch.setattr(serving_video_stream.FrameSimilarityFilter, "should_retain", fail_should_retain) + + ws = TimedWebSocket() + handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "Invalid image data"} in ws.sent + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch): + captured_audio_lengths: list[int] = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _process_query_engine( + self, + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ): + captured_audio_lengths.append(len(audio_buffer)) + + monkeypatch.setattr(serving_video_stream, "_MAX_AUDIO_BUFFER_BYTES", 4) + + ws = TimedWebSocket() + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "audio.chunk", "data": _b64(b"1234")}) + await asyncio.sleep(0) + ws.put({"type": "audio.chunk", "data": _b64(b"5")}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "Audio buffer overflow"} in ws.sent + assert captured_audio_lengths == [0] + + +def test_build_messages_keeps_recent_history_text_only(): + handler = OmniStreamingVideoHandler(chat_service=object()) + old_frame = _b64(_make_jpeg(1, 2, 3)) + current_frame = _b64(_make_jpeg(4, 5, 6)) + history = [ + {"role": "user", "content": [{"type": "text", "text": "old question"}]}, + {"role": "assistant", "content": "old answer"}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{old_frame}"}}, + {"type": "input_audio", "input_audio": {"data": "ignored", "format": "wav"}}, + {"type": "text", "text": "recent question"}, + ], + }, + {"role": "assistant", "content": "recent answer"}, + ] + + messages, user_message = handler._build_messages( + StreamingVideoSessionConfig(model="test", num_frames=1), + [current_frame], + bytearray(), + history, + "current question", + {}, + ) + + assert messages[0] == {"role": "user", "content": "recent question"} + assert messages[1] == {"role": "assistant", "content": "recent answer"} + assert messages[2] == user_message + assert user_message["content"][-1] == {"type": "text", "text": "current question"} diff --git a/tests/entrypoints/openai_api/test_video_frame_filter.py b/tests/entrypoints/openai_api/test_video_frame_filter.py new file mode 100644 index 000000000..a734aec75 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_frame_filter.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for FrameSimilarityFilter (Phase 2 EVS).""" + +from __future__ import annotations + +import pytest + +from tests.entrypoints.openai_api.conftest_video import ( + make_gradient_jpeg, + make_jpeg, +) +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter + + +class TestFrameSimilarityFilter: + def test_first_frame_always_retained(self): + f = FrameSimilarityFilter(threshold=0.99) + assert f.should_retain(make_jpeg()) is True + + def test_identical_frames_dropped(self): + f = FrameSimilarityFilter(threshold=0.90) + frame = make_jpeg(100, 100, 100) + assert f.should_retain(frame) is True + assert f.should_retain(frame) is False + + def test_very_different_frames_retained(self): + f = FrameSimilarityFilter(threshold=0.95) + assert f.should_retain(make_jpeg(255, 255, 255)) is True + assert f.should_retain(make_jpeg(0, 0, 0)) is True + + def test_low_threshold_keeps_slightly_different(self): + f = FrameSimilarityFilter(threshold=0.50) + assert f.should_retain(make_jpeg(100, 100, 100)) is True + assert f.should_retain(make_jpeg(110, 110, 110)) is True + + def test_random_frames_all_retained(self): + f = FrameSimilarityFilter(threshold=0.95) + for i in range(5): + assert f.should_retain(make_gradient_jpeg(seed=i)) is True + + def test_reset_clears_state(self): + f = FrameSimilarityFilter(threshold=0.90) + frame = make_jpeg(128, 128, 128) + assert f.should_retain(frame) is True + assert f.should_retain(frame) is False + f.reset() + assert f.should_retain(frame) is True + + def test_stats_counting(self): + f = FrameSimilarityFilter(threshold=0.90) + frame = make_jpeg(50, 50, 50) + f.should_retain(frame) # retained + f.should_retain(frame) # dropped + f.should_retain(frame) # dropped + stats = f.stats + assert stats["retained_count"] == 1 + assert stats["dropped_count"] == 2 + assert stats["total_count"] == 3 + assert abs(stats["drop_rate"] - 2.0 / 3.0) < 1e-6 + + def test_stats_empty(self): + stats = FrameSimilarityFilter().stats + assert stats["total_count"] == 0 + assert stats["drop_rate"] == 0.0 + + def test_stats_reset(self): + f = FrameSimilarityFilter(threshold=0.90) + f.should_retain(make_jpeg()) + f.reset() + assert f.stats["total_count"] == 0 + + def test_invalid_threshold(self): + with pytest.raises(ValueError, match="threshold"): + FrameSimilarityFilter(threshold=1.5) + with pytest.raises(ValueError, match="threshold"): + FrameSimilarityFilter(threshold=-0.1) + + def test_invalid_thumbnail_size(self): + with pytest.raises(ValueError, match="thumbnail_size"): + FrameSimilarityFilter(thumbnail_size=0) + + def test_different_image_sizes_same_colour(self): + """Filter should handle frames of varying resolutions.""" + f = FrameSimilarityFilter(threshold=0.90) + small = make_jpeg(100, 100, 100, size=32) + large = make_jpeg(100, 100, 100, size=256) + assert f.should_retain(small) is True + assert f.should_retain(large) is False diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py new file mode 100644 index 000000000..bee5e0de3 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_stream_handler.py @@ -0,0 +1,593 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runtime behavior tests for VideoStreamHandler. + +Every test creates a mock WebSocket, drives handle_session through a +specific code path, and asserts on the JSON messages actually sent back. +No inspect.getsource() tricks — these test real async execution. +""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import json +from unittest.mock import AsyncMock + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.openai.video_stream_session import ( + VideoStreamHandler, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes: + img = Image.new("RGB", (64, 64), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def _b64(data: bytes) -> str: + return base64.b64encode(data).decode() + + +def _config_msg(**overrides) -> str: + msg = {"type": "session.config", "model": "test", "evs_enabled": False} + msg.update(overrides) + return json.dumps(msg) + + +def _frame_msg(jpeg: bytes | None = None) -> str: + if jpeg is None: + jpeg = _make_jpeg() + return json.dumps({"type": "video.frame", "data": _b64(jpeg)}) + + +def _query_msg(text: str = "What do you see?") -> str: + return json.dumps({"type": "video.query", "text": text}) + + +def _done_msg() -> str: + return json.dumps({"type": "video.done"}) + + +def _audio_msg(pcm: bytes = b"\x00" * 320) -> str: + return json.dumps({"type": "audio.chunk", "data": _b64(pcm)}) + + +def _make_sse_line(content: str) -> str: + payload = {"choices": [{"delta": {"content": content}}]} + return f"data: {json.dumps(payload)}\n\n" + + +class MockWebSocket: + """Async-compatible mock WebSocket that feeds messages from a list.""" + + def __init__(self, messages: list[str]): + self._messages = list(messages) + self._idx = 0 + self._sent: list[dict | bytes] = [] + self._accepted = False + + async def accept(self): + self._accepted = True + + async def receive_text(self) -> str: + if self._idx >= len(self._messages): + # Simulate connection close by hanging forever (will be timed out) + await asyncio.sleep(999) + msg = self._messages[self._idx] + self._idx += 1 + return msg + + async def send_json(self, data: dict): + self._sent.append(data) + + async def send_bytes(self, data: bytes): + self._sent.append(data) + + @property + def sent_messages(self) -> list[dict]: + return [m for m in self._sent if isinstance(m, dict)] + + def sent_types(self) -> list[str]: + return [m.get("type", "") for m in self.sent_messages] + + +class MockChatHandler: + """Mock OmniOpenAIServingChat for testing _handle_query paths.""" + + def __init__(self, response=None): + self._response = response + + async def create_chat_completion(self, request, raw_request=None): + if self._response is not None: + return self._response + + # Default: return a simple streaming generator + async def _gen(): + yield _make_sse_line("Hello") + yield _make_sse_line("World") + yield "data: [DONE]\n\n" + + return _gen() + + +# --------------------------------------------------------------------------- +# _receive_config path tests +# --------------------------------------------------------------------------- + + +class TestReceiveConfig: + @pytest.mark.asyncio + async def test_config_timeout(self): + """No message within config_timeout → error sent, session ends.""" + ws = MockWebSocket([]) # no messages — will hang + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=0.05) + await handler.handle_session(ws) + + assert ws._accepted + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert len(errors) == 1 + assert "Timeout" in errors[0]["message"] + + @pytest.mark.asyncio + async def test_config_invalid_json(self): + """Non-JSON config message → error, session ends.""" + ws = MockWebSocket(["not json"]) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Invalid JSON" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_config_wrong_type(self): + """Message with wrong type → error, session ends.""" + ws = MockWebSocket([json.dumps({"type": "wrong"})]) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Expected session.config" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_config_invalid_field_type(self): + """Config with wrong field type → error, session ends.""" + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "max_frames": "potato", + } + ) + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("max_frames" in e["message"] for e in errors) + + +# --------------------------------------------------------------------------- +# Normal session flow +# --------------------------------------------------------------------------- + + +class TestNormalFlow: + @pytest.mark.asyncio + async def test_frame_query_done(self): + """Happy path: config → frame → query → done → session.done.""" + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + _query_msg("Describe."), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + types = ws.sent_types() + assert "response.start" in types + assert "response.text.delta" in types + assert "response.text.done" in types + assert "session.done" in types + + # Check text content was streamed + deltas = [m["delta"] for m in ws.sent_messages if m.get("type") == "response.text.delta"] + assert "Hello" in deltas + assert "World" in deltas + + # response.text.done has full text + done_msg = next(m for m in ws.sent_messages if m.get("type") == "response.text.done") + assert done_msg["text"] == "HelloWorld" + + @pytest.mark.asyncio + async def test_multiple_frames_before_query(self): + """Multiple frames accumulate, query uses sampled frames.""" + frames = [_frame_msg(_make_jpeg(r=i * 50)) for i in range(5)] + ws = MockWebSocket( + [ + _config_msg(num_sample_frames=3), + *frames, + _query_msg(), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_evs_stats_in_done(self): + """EVS stats sent before session.done when evs_enabled=True.""" + frame = _frame_msg(_make_jpeg(100, 100, 100)) + ws = MockWebSocket( + [ + _config_msg(evs_enabled=True, evs_threshold=0.90), + frame, + frame, # second should be dropped + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + evs_msgs = [m for m in ws.sent_messages if m.get("type") == "response.evs_stats"] + assert len(evs_msgs) == 1 + assert evs_msgs[0]["retained_count"] == 1 + assert evs_msgs[0]["dropped_count"] == 1 + + +# --------------------------------------------------------------------------- +# _handle_query error paths +# --------------------------------------------------------------------------- + + +class TestHandleQueryErrors: + @pytest.mark.asyncio + async def test_query_with_no_frames(self): + """Query before any frames → error, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + _query_msg("Hello?"), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("No frames" in e["message"] for e in errors) + # Session still ends cleanly + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_query_returns_error_response(self): + """create_chat_completion returns ErrorResponse → error forwarded.""" + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + + err = ErrorResponse( + message="model overloaded", + type="server_error", + code=503, + ) + + chat = AsyncMock() + chat.create_chat_completion = AsyncMock(return_value=err) + + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + _query_msg(), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=chat, idle_timeout=2.0) + await handler.handle_session(ws) + + # Should see response.start, then error, then response.text.done + assert "response.start" in ws.sent_types() + assert "response.text.done" in ws.sent_types() + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("model overloaded" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_query_generator_raises(self): + """Exception during streaming → error sent, session continues.""" + + async def _exploding_gen(): + yield _make_sse_line("partial") + raise RuntimeError("CUDA OOM") + + chat = MockChatHandler() + chat._response = None # override below + + class BoomChat: + async def create_chat_completion(self, req, raw_request=None): + return _exploding_gen() + + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + _query_msg(), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=BoomChat(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Query processing failed" in e["message"] for e in errors) + # Crucially: no raw exception message leaked + assert not any("CUDA OOM" in e["message"] for e in errors) + # Protocol: response.text.done always sent after response.start + assert "response.text.done" in ws.sent_types() + # Session still ends cleanly + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_empty_query_text_rejected(self): + """video.query with empty text → error, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + json.dumps({"type": "video.query", "text": ""}), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("non-empty" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + +# --------------------------------------------------------------------------- +# Timeout paths +# --------------------------------------------------------------------------- + + +class TestTimeouts: + @pytest.mark.asyncio + async def test_idle_timeout(self): + """No message after config within idle_timeout → error, session ends.""" + ws = MockWebSocket([_config_msg()]) # config then nothing + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=0.05) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Idle timeout" in e["message"] for e in errors) + + +# --------------------------------------------------------------------------- +# Reader error paths +# --------------------------------------------------------------------------- + + +class TestReaderErrors: + @pytest.mark.asyncio + async def test_invalid_json_mid_session(self): + """Invalid JSON mid-session → error sent, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + "{{bad json", + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Invalid JSON" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_non_dict_message(self): + """Non-object JSON → error sent, continues.""" + ws = MockWebSocket( + [ + _config_msg(), + json.dumps([1, 2, 3]), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("JSON objects" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_unknown_message_type(self): + """Unknown type → error, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + json.dumps({"type": "teleport.now"}), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Unknown message type" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_invalid_base64_frame(self): + """Invalid base64 in video.frame → error, continues.""" + ws = MockWebSocket( + [ + _config_msg(), + json.dumps({"type": "video.frame", "data": "!!!not-base64!!!"}), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Invalid base64" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + +# --------------------------------------------------------------------------- +# Concurrency: frames arrive during query processing +# --------------------------------------------------------------------------- + + +class TestConcurrency: + @pytest.mark.asyncio + async def test_frames_during_query(self): + """Frames sent while a query is processing must be buffered.""" + query_started = asyncio.Event() + query_can_finish = asyncio.Event() + captured_frame_count: list[int] = [] + + class SlowChat: + async def create_chat_completion(self, req, raw_request=None): + query_started.set() + await query_can_finish.wait() + + async def _gen(): + yield _make_sse_line("answer") + yield "data: [DONE]\n\n" + + return _gen() + + class InspectingSlowChat(SlowChat): + """On the second query, record how many frames are in the request.""" + + _call_count = 0 + + async def create_chat_completion(self, req, raw_request=None): + self._call_count += 1 + if self._call_count == 2: + # Count image_url blocks in the request to verify buffering + for msg in req.messages: + if isinstance(msg.get("content", msg), list): + content = msg.get("content", msg) + else: + content = getattr(msg, "content", []) + if isinstance(content, list): + n = sum( + 1 + for p in content + if (isinstance(p, dict) and p.get("type") == "image_url") + or (hasattr(p, "type") and p.type == "image_url") + ) + captured_frame_count.append(n) + return await super().create_chat_completion(req, raw_request) + + # Custom WebSocket that can inject frames after query starts + class TimedWebSocket(MockWebSocket): + def __init__(self): + self._q: asyncio.Queue[str] = asyncio.Queue() + self._sent: list[dict | bytes] = [] + self._accepted = False + + async def accept(self): + self._accepted = True + + async def receive_text(self) -> str: + return await self._q.get() + + def put(self, msg: str): + self._q.put_nowait(msg) + + ws = TimedWebSocket() + chat = InspectingSlowChat() + handler = VideoStreamHandler(chat_handler=chat, idle_timeout=5.0) + + session_task = asyncio.create_task(handler.handle_session(ws)) + + # 1. Config + ws.put(_config_msg()) + await asyncio.sleep(0.01) + + # 2. Initial frame + query + ws.put(_frame_msg(_make_jpeg(10, 10, 10))) + await asyncio.sleep(0.01) + ws.put(_query_msg("first")) + + # 3. Wait for query to start processing + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + # 4. Send more frames WHILE query is running + ws.put(_frame_msg(_make_jpeg(200, 200, 200))) + ws.put(_frame_msg(_make_jpeg(50, 50, 50))) + await asyncio.sleep(0.05) + + # 5. Let query finish + query_can_finish.set() + await asyncio.sleep(0.1) + + # 6. Second query — should see all 3 frames (1 initial + 2 during) + query_started.clear() + query_can_finish.clear() + ws.put(_query_msg("second")) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + query_can_finish.set() + await asyncio.sleep(0.1) + + # 7. Done + ws.put(_done_msg()) + + await asyncio.wait_for(session_task, timeout=5.0) + + assert "session.done" in ws.sent_types() + # Verify: the 2 frames sent during query 1 were buffered and + # included in query 2 (total 3 frames). + assert captured_frame_count == [3] + + +# --------------------------------------------------------------------------- +# SSE parsing (kept from original, these are valid unit tests) +# --------------------------------------------------------------------------- + + +class TestParseSSEDeltas: + def test_single_delta(self): + assert VideoStreamHandler._parse_sse_deltas(_make_sse_line("hello")) == ["hello"] + + def test_multiple_deltas(self): + chunk = _make_sse_line("a") + _make_sse_line("b") + assert VideoStreamHandler._parse_sse_deltas(chunk) == ["a", "b"] + + def test_done_skipped(self): + chunk = _make_sse_line("x") + "data: [DONE]\n\n" + assert VideoStreamHandler._parse_sse_deltas(chunk) == ["x"] + + def test_empty_content_skipped(self): + p = json.dumps({"choices": [{"delta": {"content": ""}}]}) + assert VideoStreamHandler._parse_sse_deltas(f"data: {p}\n\n") == [] + + def test_malformed_json_skipped(self): + chunk = "data: {bad}\n\n" + _make_sse_line("ok") + assert VideoStreamHandler._parse_sse_deltas(chunk) == ["ok"] + + def test_empty_string(self): + assert VideoStreamHandler._parse_sse_deltas("") == [] + + def test_unicode(self): + assert VideoStreamHandler._parse_sse_deltas(_make_sse_line("你好🌍")) == ["你好🌍"] diff --git a/tests/entrypoints/openai_api/test_video_stream_session.py b/tests/entrypoints/openai_api/test_video_stream_session.py new file mode 100644 index 000000000..b0a12b1e4 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_stream_session.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for VideoStreamSession (Phase 2 + 3).""" + +from __future__ import annotations + +import pytest + +from tests.entrypoints.openai_api.conftest_video import ( + make_gradient_jpeg, + make_jpeg, +) +from vllm_omni.entrypoints.openai.video_stream_session import ( + VideoStreamConfig, + VideoStreamSession, +) + +# --------------------------------------------------------------------------- +# VideoStreamConfig +# --------------------------------------------------------------------------- + + +class TestVideoStreamConfig: + def test_defaults(self): + cfg = VideoStreamConfig() + assert cfg.max_frames == 64 + assert cfg.num_sample_frames == 16 + assert cfg.evs_enabled is True + assert cfg.evs_threshold == 0.95 + + def test_from_dict(self): + cfg = VideoStreamConfig.from_dict( + { + "model": "test-model", + "max_frames": 32, + "evs_threshold": 0.90, + "unknown_field": "ignored", + } + ) + assert cfg.model == "test-model" + assert cfg.max_frames == 32 + assert cfg.evs_threshold == 0.90 + + def test_from_dict_empty(self): + cfg = VideoStreamConfig.from_dict({}) + assert cfg.model == "" + assert cfg.max_frames == 64 + + def test_from_dict_invalid_type(self): + with pytest.raises(TypeError, match="max_frames.*expected int.*got str"): + VideoStreamConfig.from_dict({"max_frames": "potato"}) + + def test_from_dict_invalid_bool(self): + with pytest.raises(TypeError, match="evs_enabled.*expected bool"): + VideoStreamConfig.from_dict({"evs_enabled": "yes"}) + + def test_from_dict_invalid_modalities(self): + with pytest.raises(TypeError, match="modalities.*expected list"): + VideoStreamConfig.from_dict({"modalities": 42}) + + def test_from_dict_evs_threshold_int_accepted(self): + """JSON doesn't distinguish int/float — int 1 is a valid threshold.""" + cfg = VideoStreamConfig.from_dict({"evs_threshold": 1}) + assert cfg.evs_threshold == 1 + + +# --------------------------------------------------------------------------- +# Frame buffer & sliding window (uses deque) +# --------------------------------------------------------------------------- + + +class TestFrameBuffer: + def test_add_frame_basic(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=10) + session = VideoStreamSession(cfg) + assert session.add_frame(make_jpeg()) is True + assert session.frame_count == 1 + + def test_add_frame_too_large(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=10) + session = VideoStreamSession(cfg) + huge = b"\xff" * (10 * 1024 * 1024 + 1) # just over 10 MB + with pytest.raises(ValueError, match="Frame too large"): + session.add_frame(huge) + + def test_sliding_window(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=3) + session = VideoStreamSession(cfg) + for i in range(5): + session.add_frame(make_jpeg(r=i * 50)) + assert session.frame_count == 3 + + def test_sliding_window_keeps_newest(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=2) + session = VideoStreamSession(cfg) + f1 = make_jpeg(10, 10, 10) + f2 = make_jpeg(20, 20, 20) + f3 = make_jpeg(30, 30, 30) + session.add_frame(f1) + session.add_frame(f2) + session.add_frame(f3) + sampled = session.sample_frames() + assert len(sampled) == 2 + assert sampled[0] == f2 + assert sampled[1] == f3 + + +# --------------------------------------------------------------------------- +# EVS integration +# --------------------------------------------------------------------------- + + +class TestEVSIntegration: + def test_evs_drops_identical_frames(self): + cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.90) + session = VideoStreamSession(cfg) + frame = make_jpeg(100, 100, 100) + assert session.add_frame(frame) is True + assert session.add_frame(frame) is False + assert session.frame_count == 1 + + def test_evs_keeps_different_frames(self): + cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.95) + session = VideoStreamSession(cfg) + for i in range(5): + assert session.add_frame(make_gradient_jpeg(seed=i)) is True + assert session.frame_count == 5 + + def test_evs_disabled(self): + cfg = VideoStreamConfig(evs_enabled=False) + session = VideoStreamSession(cfg) + frame = make_jpeg() + assert session.add_frame(frame) is True + assert session.add_frame(frame) is True + assert session.frame_count == 2 + + def test_evs_stats(self): + cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.90) + session = VideoStreamSession(cfg) + frame = make_jpeg() + session.add_frame(frame) + session.add_frame(frame) + stats = session.evs_stats + assert stats is not None + assert stats["retained_count"] == 1 + assert stats["dropped_count"] == 1 + + def test_evs_stats_none_when_disabled(self): + cfg = VideoStreamConfig(evs_enabled=False) + session = VideoStreamSession(cfg) + assert session.evs_stats is None + + +# --------------------------------------------------------------------------- +# Uniform sampling +# --------------------------------------------------------------------------- + + +class TestSampling: + def test_sample_exact(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=3, max_frames=10) + session = VideoStreamSession(cfg) + frames = [make_gradient_jpeg(i) for i in range(3)] + for f in frames: + session.add_frame(f) + sampled = session.sample_frames() + assert len(sampled) == 3 + assert sampled == frames + + def test_sample_fewer_than_requested(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=10, max_frames=64) + session = VideoStreamSession(cfg) + for i in range(3): + session.add_frame(make_gradient_jpeg(i)) + assert len(session.sample_frames()) == 3 + + def test_sample_uniform(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=4, max_frames=64) + session = VideoStreamSession(cfg) + for i in range(10): + session.add_frame(make_gradient_jpeg(i)) + sampled = session.sample_frames() + assert len(sampled) == 4 + expected = [session._frames[i] for i in [0, 3, 6, 9]] + assert sampled == expected + + def test_sample_empty(self): + cfg = VideoStreamConfig(evs_enabled=False) + session = VideoStreamSession(cfg) + assert session.sample_frames() == [] + + def test_sample_single_frame_from_multi_frame_buffer(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=1, max_frames=64) + session = VideoStreamSession(cfg) + first = make_gradient_jpeg(0) + second = make_gradient_jpeg(1) + session.add_frame(first) + session.add_frame(second) + + assert session.sample_frames() == [second] + + +# --------------------------------------------------------------------------- +# Audio buffer (Phase 3) +# --------------------------------------------------------------------------- + + +class TestAudioBuffer: + def test_add_audio_chunk(self): + session = VideoStreamSession(VideoStreamConfig()) + assert session.has_audio is False + session.add_audio_chunk(b"\x00" * 100) + assert session.has_audio is True + + def test_clear_audio(self): + session = VideoStreamSession(VideoStreamConfig()) + session.add_audio_chunk(b"\x00" * 100) + session.clear_audio() + assert session.has_audio is False + + +# --------------------------------------------------------------------------- +# build_chat_request +# --------------------------------------------------------------------------- + + +class TestBuildChatRequest: + def test_video_only_request(self): + cfg = VideoStreamConfig(model="test-model", evs_enabled=False, num_sample_frames=4) + session = VideoStreamSession(cfg) + for i in range(4): + session.add_frame(make_gradient_jpeg(i)) + + request = session.build_chat_request("Describe this scene.") + assert request.model == "test-model" + assert request.stream is True + + content = request.messages[0]["content"] + assert len(content) == 5 # 4 image_url + 1 text + image_parts = [p for p in content if p["type"] == "image_url"] + text_parts = [p for p in content if p["type"] == "text"] + assert len(image_parts) == 4 + assert len(text_parts) == 1 + assert text_parts[0]["text"] == "Describe this scene." + + mm_kw = getattr(request, "mm_processor_kwargs", None) + assert mm_kw is None or not mm_kw.get("use_audio_in_video", False) + + def test_video_plus_audio_request(self): + cfg = VideoStreamConfig(model="test-model", evs_enabled=False, num_sample_frames=2) + session = VideoStreamSession(cfg) + session.add_frame(make_gradient_jpeg(0)) + session.add_frame(make_gradient_jpeg(1)) + session.add_audio_chunk(b"\x00" * 3200) + + request = session.build_chat_request("What is being said?") + + content = request.messages[0]["content"] + assert len(content) == 4 # 2 image_url + 1 audio_url + 1 text + audio_parts = [p for p in content if p["type"] == "audio_url"] + assert len(audio_parts) == 1 + # RFC 3551: audio/L16 for linear 16-bit PCM + assert audio_parts[0]["audio_url"]["url"].startswith("data:audio/L16;rate=16000;base64,") + + mm_kw = getattr(request, "mm_processor_kwargs", None) or {} + assert mm_kw.get("use_audio_in_video") is True + + def test_image_url_is_valid_base64(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=1) + session = VideoStreamSession(cfg) + session.add_frame(make_jpeg(200, 100, 50)) + request = session.build_chat_request("test") + content = request.messages[0]["content"] + img_url = content[0]["image_url"]["url"] + assert img_url.startswith("data:image/jpeg;base64,") + import base64 + + b64_data = img_url.split(",", 1)[1] + decoded = base64.b64decode(b64_data) + assert len(decoded) > 0 + + def test_clear_audio_after_query(self): + session = VideoStreamSession(VideoStreamConfig(evs_enabled=False)) + session.add_frame(make_jpeg()) + session.add_audio_chunk(b"\x00" * 100) + session.build_chat_request("test") + session.clear_audio() + assert session.has_audio is False + assert session.frame_count == 1 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 78073ae21..6345f0006 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -117,6 +117,7 @@ from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage +from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request @@ -606,6 +607,7 @@ async def omni_init_app_state( stage_configs=diffusion_stage_configs, ) state.openai_streaming_speech = None + state.openai_streaming_video = None state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) state.server_load_metrics = 0 @@ -907,6 +909,14 @@ async def omni_init_app_state( state.openai_streaming_speech = OmniStreamingSpeechHandler( speech_service=state.openai_serving_speech, ) + state.openai_streaming_video = ( + OmniStreamingVideoHandler( + chat_service=state.openai_serving_chat, + engine_client=engine_client, + ) + if state.openai_serving_chat is not None + else None + ) state.openai_serving_realtime = OpenAIServingRealtime( engine_client=engine_client, models=state.openai_serving_models, @@ -1292,6 +1302,23 @@ async def streaming_speech(websocket: WebSocket): await handler.handle_session(websocket) +@router.websocket("/v1/video/chat/stream") +async def streaming_video_chat(websocket: WebSocket): + """WebSocket endpoint for streaming video input chat.""" + handler = getattr(websocket.app.state, "openai_streaming_video", None) + if handler is None: + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "message": "Streaming video chat is not available", + } + ) + await websocket.close() + return + await handler.handle_session(websocket) + + @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py new file mode 100644 index 000000000..a76b241c5 --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -0,0 +1,971 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""WebSocket handler for streaming video input understanding. + +Accepts video frames incrementally via WebSocket, buffers them, and +generates text + optional audio responses using the existing Qwen3-Omni +multi-stage pipeline (thinker -> talker -> code2wav). + +Protocol: + Client -> Server: + {"type": "session.config", ...} # Session config (sent once) + {"type": "video.frame", "data": "..."} # base64 JPEG/PNG frame + {"type": "audio.chunk", "data": "..."} # base64 PCM16 16kHz mono + {"type": "video.query", "text": "..."} # Submit query about buffered frames + {"type": "video.done"} # End of session + + Server -> Client: + {"type": "response.start"} + {"type": "response.text.delta", "delta": "..."} + {"type": "response.text.done", "text": "..."} + {"type": "response.audio.delta", "data": "...", "format": "wav"} + {"type": "response.audio.done"} + {"type": "session.done"} + {"type": "error", "message": "..."} +""" + +import asyncio +import base64 +import hashlib +import io +import json +import time as _time +import uuid +import wave +from typing import Any + +import torch +from fastapi import WebSocket, WebSocketDisconnect +from PIL import Image +from pydantic import BaseModel, Field, ValidationError +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai import video_stream_envs +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter +from vllm_omni.entrypoints.openai.video_stream_context import ( + text_only_message, +) +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + +_DEFAULT_IDLE_TIMEOUT = 60.0 +_DEFAULT_CONFIG_TIMEOUT = 10.0 +_MAX_FRAME_SIZE = 10 * 1024 * 1024 # 10MB per frame +_MAX_BUFFER_FRAMES = 64 +_MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 +_MAX_MSG_QUEUE = 200 +_CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length +_BAD_FRAME = object() + + +def _decode_frame_bytes(raw_bytes: bytes) -> Any: + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + +class StreamingVideoSessionConfig(BaseModel): + """Configuration sent as the first WebSocket message.""" + + model: str | None = None + modalities: list[str] = Field( + default_factory=lambda: ["text", "audio"], + description="Output modalities: 'text', 'audio', or both.", + ) + num_frames: int = Field( + default=4, + ge=1, + le=128, + description="Max frames to sample from buffer for the model.", + ) + max_frames: int = Field( + default=50, + ge=1, + le=256, + description="Max frames to keep in the buffer.", + ) + system_prompt: str | None = Field( + default=None, + description="Custom system prompt.", + ) + use_audio_in_video: bool = Field( + default=True, + description="Interleave audio chunks with video frames when audio input is present.", + ) + sampling_params_list: list[dict[str, Any]] | None = Field( + default=None, + description="Per-stage sampling params [thinker, talker, code2wav].", + ) + enable_frame_filter: bool = Field( + default=True, + description="EVS pixel-similarity pre-filter to drop near-duplicate frames.", + ) + frame_filter_threshold: float = Field( + default=0.95, + ge=0.0, + le=1.0, + description="EVS similarity threshold (higher = keep more frames).", + ) + + +class OmniStreamingVideoHandler: + """Handles WebSocket sessions for streaming video input. + + Supports: + - Concurrent frame reception during query processing (reader/processor split) + - PCM audio input (``audio.chunk``) + - Async-chunk incremental audio output via ``engine_client.generate()`` + - Multi-turn conversation history + - Soft interrupt (new query cancels current generation) + """ + + def __init__( + self, + chat_service: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + engine_client: Any | None = None, + ) -> None: + self._chat_service = chat_service + self._idle_timeout = idle_timeout + self._config_timeout = config_timeout + self._engine_client = engine_client + + async def handle_session(self, websocket: WebSocket) -> None: + """Main session loop for a single WebSocket connection.""" + await websocket.accept() + + try: + config = await self._receive_config(websocket) + if config is None: + return + + frame_buffer: list[str] = [] # base64-encoded JPEG frames + # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. + frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME + frame_filter = ( + FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None + ) + audio_buffer = bytearray() # raw PCM16 16kHz mono + message_history: list[dict[str, Any]] = [] + active_request_id: str | None = None + prev_request_id: str | None = None # abort target iff prev was interrupted + prev_was_interrupted: bool = False + interrupt_event = asyncio.Event() + prewarm_tasks: set[asyncio.Task[Any]] = set() + query_task: asyncio.Task[Any] | None = None + + msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=_MAX_MSG_QUEUE) + + async def _reader() -> None: + """Receive WebSocket messages and enqueue them.""" + try: + while True: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._idle_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Idle timeout") + await msg_queue.put(None) + return + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON") + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "Messages must be JSON objects") + continue + + msg_type = str(msg.get("type", "")) + if msg_type.startswith("_internal."): + await self._send_error(websocket, f"Unknown type: {msg_type}") + continue + + await msg_queue.put(msg) + if msg.get("type") == "video.done": + return + except WebSocketDisconnect: + await msg_queue.put(None) + except Exception: + await msg_queue.put(None) + raise + + async def _cancel_active_query(*, abort_now: bool = False) -> None: + """Signal soft interrupt for the active query.""" + nonlocal active_request_id, prev_was_interrupted, query_task + if active_request_id is not None: + interrupt_event.set() + prev_was_interrupted = True + logger.info("Interrupt signaled for %s", active_request_id) + if abort_now and self._engine_client: + try: + await self._engine_client.abort(active_request_id) + except Exception: + logger.debug("Abort failed for %s", active_request_id, exc_info=True) + if query_task is not None and not query_task.done(): + query_task.cancel() + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + + async def _processor() -> None: + """Process enqueued messages.""" + nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task + + while True: + msg = await msg_queue.get() + if msg is None: + await _cancel_active_query(abort_now=True) + return + + msg_type = msg.get("type") + + if msg_type == "_internal.frame_decode_failed": + frame_data = msg.get("b64", "") + removed = frame_data in frame_buffer + if removed: + frame_buffer[:] = [f for f in frame_buffer if f != frame_data] + if frame_pil_cache.get(frame_data) is _BAD_FRAME: + frame_pil_cache.pop(frame_data, None) + if removed: + await self._send_error(websocket, "Frame decode failed") + + elif msg_type == "video.frame": + frame_data = msg.get("data", "") + if not frame_data: + continue + if len(frame_data) > _MAX_FRAME_SIZE: + await self._send_error(websocket, "Frame too large") + continue + try: + raw_bytes = base64.b64decode(frame_data, validate=True) + except Exception: + await self._send_error(websocket, "Invalid image data") + continue + if frame_filter is not None: + try: + if not frame_filter.should_retain(raw_bytes): + continue + except Exception: + await self._send_error(websocket, "Invalid image data") + continue + max_buf = config.max_frames + if len(frame_buffer) >= max_buf: + dropped = frame_buffer.pop(0) + frame_pil_cache.pop(dropped, None) + frame_buffer.append(frame_data) + # Prewarm: decode PIL off the event loop so query-time chat_template + # can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames. + if frame_data not in frame_pil_cache: + mm_uuid = hashlib.md5(raw_bytes, usedforsecurity=False).hexdigest() + + async def _prewarm(b64: str, b: bytes, u: str) -> None: + try: + pil = await asyncio.to_thread(_decode_frame_bytes, b) + frame_pil_cache[b64] = (pil, u) + except Exception: + frame_pil_cache[b64] = _BAD_FRAME + logger.warning("prewarm decode failed for frame (len=%d)", len(b)) + try: + msg_queue.put_nowait({"type": "_internal.frame_decode_failed", "b64": b64}) + except asyncio.QueueFull: + logger.warning( + "frame decode failure event dropped because message queue is full" + ) + + task = asyncio.create_task(_prewarm(frame_data, raw_bytes, mm_uuid)) + prewarm_tasks.add(task) + task.add_done_callback(prewarm_tasks.discard) + + elif msg_type == "audio.chunk": + data_b64 = msg.get("data", "") + try: + pcm_bytes = base64.b64decode(data_b64) + except Exception: + continue + if len(audio_buffer) + len(pcm_bytes) > _MAX_AUDIO_BUFFER_BYTES: + await self._send_error(websocket, "Audio buffer overflow") + audio_buffer.clear() + continue + audio_buffer.extend(pcm_bytes) + + elif msg_type == "video.query": + await _cancel_active_query() + + query_text = msg.get("text", "") + audio_data_b64 = msg.get("audio_data") + if audio_data_b64: + try: + decoded = base64.b64decode(audio_data_b64) + if len(audio_buffer) + len(decoded) <= _MAX_AUDIO_BUFFER_BYTES: + audio_buffer.extend(decoded) + else: + await self._send_error(websocket, "Audio buffer overflow") + audio_buffer.clear() + except Exception: + pass + + if not frame_buffer: + await self._send_error(websocket, "No frames buffered") + continue + + # Abort only if the previous turn was interrupted mid-flight. + # A naturally-finished request is already released by the scheduler; + # aborting it again can race with stage-1/2 tear-down and has been + # observed to crash flash_attn with a mixed prefill+decode batch + # (scheduler_metadata shape mismatch) under longer sessions. + if prev_was_interrupted and prev_request_id and self._engine_client: + try: + await self._engine_client.abort(prev_request_id) + except Exception: + pass + await asyncio.sleep(0.1) + prev_was_interrupted = False + + request_id = f"video-{uuid.uuid4().hex[:12]}" + active_request_id = request_id + interrupt_event.clear() + query_frames = list(frame_buffer) + query_audio_buffer = bytearray(audio_buffer) + audio_buffer.clear() + query_prewarmed_frames = dict(frame_pil_cache) + + async def _run_query() -> None: + nonlocal active_request_id, prev_request_id + try: + await self._process_query( + websocket, + config, + query_frames, + query_audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + query_prewarmed_frames, + ) + finally: + if active_request_id == request_id: + prev_request_id = request_id + active_request_id = None + + query_task = asyncio.create_task(_run_query()) + + elif msg_type == "video.done": + if query_task is not None and not query_task.done(): + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + await websocket.send_json({"type": "session.done"}) + return + + elif msg_type == "ping": + try: + await websocket.send_json({"type": "pong"}) + except Exception: + pass + + else: + await self._send_error(websocket, f"Unknown type: {msg_type}") + + reader_task = asyncio.create_task(_reader()) + try: + await _processor() + finally: + reader_task.cancel() + try: + await reader_task + except (asyncio.CancelledError, Exception): + pass + for t in list(prewarm_tasks): + t.cancel() + if prewarm_tasks: + await asyncio.gather(*prewarm_tasks, return_exceptions=True) + if query_task is not None and not query_task.done(): + await _cancel_active_query(abort_now=True) + + except WebSocketDisconnect: + logger.info("Streaming video: client disconnected") + except Exception as e: + logger.exception("Streaming video session error: %s", e) + try: + await self._send_error(websocket, f"Internal error: {e}") + except Exception: + pass + + async def _receive_config(self, websocket: WebSocket) -> StreamingVideoSessionConfig | None: + """Wait for and validate the session.config message.""" + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._config_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.config") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.config") + return None + + if not isinstance(msg, dict) or msg.get("type") != "session.config": + await self._send_error( + websocket, + f"Expected session.config, got: {msg.get('type') if isinstance(msg, dict) else type(msg).__name__}", + ) + return None + + config_data = {k: v for k, v in msg.items() if k != "type"} + alias_map = { + "num_sample_frames": "num_frames", + "evs_enabled": "enable_frame_filter", + "evs_threshold": "frame_filter_threshold", + } + for old_key, new_key in alias_map.items(): + if old_key in config_data and new_key not in config_data: + config_data[new_key] = config_data[old_key] + + try: + config = StreamingVideoSessionConfig(**config_data) + except ValidationError as e: + await self._send_error(websocket, f"Invalid session config: {e}") + return None + + return config + + async def _process_query( + self, + websocket: WebSocket, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + request_id: str, + interrupt_event: asyncio.Event, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> None: + """Build prompt, run inference, stream text + audio response.""" + + if self._engine_client is None: + await self._send_error(websocket, "Streaming video requires an engine client") + return + + await self._process_query_engine( + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ) + + # ------------------------------------------------------------------ + # Engine-client path (async_chunk audio streaming) + # ------------------------------------------------------------------ + + async def _process_query_engine( + self, + websocket: WebSocket, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + request_id: str, + interrupt_event: asyncio.Event, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> None: + """Direct engine_client.generate() path for async_chunk audio.""" + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + + messages, user_message = self._build_messages( + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + prewarmed_frames, + ) + + request_kwargs: dict[str, Any] = { + "model": config.model or "default", + "messages": messages, + "stream": True, + "modalities": config.modalities, + "add_generation_prompt": True, + "continue_final_message": False, + "add_special_tokens": False, + } + if config.use_audio_in_video and len(audio_buffer) > 0: + request_kwargs["mm_processor_kwargs"] = { + "use_audio_in_video": True, + } + if config.sampling_params_list: + request_kwargs["sampling_params_list"] = config.sampling_params_list + + try: + chat_request = ChatCompletionRequest(**request_kwargs) + except Exception as e: + await self._send_error(websocket, f"Failed to build request: {e}") + return + + try: + engine_prompt = await self._preprocess_to_engine_prompt(chat_request) + except Exception as e: + await self._send_error(websocket, f"Preprocess failed: {e}") + return + + await websocket.send_json({"type": "response.start"}) + text_parts: list[str] = [] + text_done_sent = False + audio_chunk_count = 0 + # Number of per-step tensors in OmniRequestOutput.audio_data already + # drained. Used by the fast path to skip already-emitted history. + audio_chunks_drained = 0 + previous_text = "" + interrupted = False + t_start = _time.monotonic() + t_first_text = None + t_first_audio = None + + # Wire-level async-chunk switch. "off" means + # buffer all deltas server-side and flush once at the end; the engine + # pipeline still overlaps internally. + async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK + streaming = async_chunk_mode == "on" + audio_tail_tensors: list[Any] = [] + + try: + result_gen = self._engine_client.generate( + prompt=engine_prompt, + request_id=request_id, + output_modalities=config.modalities, + ) + + async for output in result_gen: + # Soft interrupt: drain without sending + if interrupt_event.is_set(): + if not interrupted: + logger.info("Generation interrupted — draining") + interrupted = True + continue + + if not isinstance(output, OmniRequestOutput): + continue + + out_type = getattr(output, "final_output_type", "text") + + if out_type == "audio": + if streaming and not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if t_first_audio is None: + t_first_audio = _time.monotonic() + audio_chunk_count += 1 + if streaming: + b64, audio_chunks_drained = self._extract_audio_delta_b64( + output, + audio_chunks_drained, + ) + if b64: + await websocket.send_json( + { + "type": "response.audio.delta", + "data": b64, + "format": "wav", + } + ) + else: + audio_data = self._get_audio_data(output) + if audio_data is not None: + if isinstance(audio_data, list): + audio_tail_tensors = list(audio_data) + else: + audio_tail_tensors = [audio_data] + else: + delta_text, previous_text = self._extract_text_delta( + output, + previous_text, + ) + if delta_text: + if t_first_text is None: + t_first_text = _time.monotonic() + text_parts.append(delta_text) + if streaming: + await websocket.send_json({"type": "response.text.delta", "delta": delta_text}) + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if not streaming and audio_tail_tensors: + try: + coalesced = ( + audio_tail_tensors[0] if len(audio_tail_tensors) == 1 else torch.cat(audio_tail_tensors, dim=-1) + ) + tail_np = self._tensor_to_1d_np(coalesced) + b64, _ = self._encode_tail( + tail_np, + 0, + new_drained=len(audio_tail_tensors), + is_first=True, + ) + if b64: + await websocket.send_json( + { + "type": "response.audio.delta", + "data": b64, + "format": "wav", + } + ) + except Exception: + logger.exception("Failed to coalesce off-path audio") + + if audio_chunk_count > 0: + await websocket.send_json({"type": "response.audio.done"}) + + response_text = "".join(text_parts) + message_history.append(user_message) + message_history.append({"role": "assistant", "content": response_text}) + + t_end = _time.monotonic() + logger.info( + "[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d", + async_chunk_mode, + t_end - t_start, + (t_first_text - t_start) if t_first_text else -1, + (t_first_audio - t_start) if t_first_audio else -1, + audio_chunk_count, + ) + + except Exception: + logger.exception("Engine query failed") + await self._send_error(websocket, "Query processing failed") + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + + # ------------------------------------------------------------------ + # Message building + # ------------------------------------------------------------------ + + def _build_messages( + self, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Build OpenAI-style messages list and the current user message. + + Returns (messages, user_message). + """ + # Stride sampling (index 0 anchor, last slot = newest). Covers full buffer + stable mm_hash. + n_buf = len(frame_buffer) + if n_buf <= config.num_frames: + frames = list(frame_buffer) + else: + stride = max(1, n_buf // config.num_frames) + idx = [i * stride for i in range(config.num_frames - 1)] + [n_buf - 1] + frames = [frame_buffer[i] for i in idx] + + # Prefer prewarmed PIL + uuid so mm_cache can dedupe by hash. + prewarmed = prewarmed_frames or {} + user_content: list[dict] = [] + for frame_b64 in frames: + cached = prewarmed.get(frame_b64) + if cached is _BAD_FRAME: + continue + if cached is not None: + pil, pil_uuid = cached + user_content.append( + { + "type": "image_pil", + "image_pil": pil, + "uuid": pil_uuid, + } + ) + else: + user_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"}, + } + ) + + if len(audio_buffer) > 0: + wav_b64 = self._pcm_to_wav_b64(bytes(audio_buffer)) + user_content.append( + { + "type": "input_audio", + "input_audio": { + "data": wav_b64, + "format": "wav", + }, + } + ) + + if query_text: + user_content.append({"type": "text", "text": query_text}) + + user_message: dict[str, Any] = {"role": "user", "content": user_content} + + messages: list[dict[str, Any]] = [] + if config.system_prompt: + messages.append({"role": "system", "content": config.system_prompt}) + + # Add text-only history (strip images/audio from old turns). + # Keep only the last turn (2 messages) to keep prompt short + # enough for single-step mm_encoder scheduling. When prompt + # exceeds ~50 tokens, the V1 scheduler splits mm_encoder and + # text prefill, causing incomplete thinker embeddings and + # garbled audio. + recent_history = message_history[-2:] if len(message_history) > 2 else message_history + for hist_msg in recent_history: + messages.append(self._text_only_message(hist_msg)) + + messages.append(user_message) + + return messages, user_message + + # ------------------------------------------------------------------ + # Audio helpers + # ------------------------------------------------------------------ + + @staticmethod + def _pcm_to_wav_b64(pcm_data: bytes, sample_rate: int = 16000) -> str: + """Wrap raw PCM16 mono in a WAV container and return base64.""" + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(pcm_data) + return base64.b64encode(buf.getvalue()).decode() + + @classmethod + def _extract_audio_delta_b64( + cls, + result: OmniRequestOutput, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Return (base64 WAV of new samples, updated chunks_drained). + + `chunks_drained` is the number of per-step tensors in + ``audio_data`` that have already been emitted. Each engine step appends + one tensor, so new samples are ``audio_data[chunks_drained:]`` — no + matter how many steps accumulated between reads (handles backpressure + cleanly, unlike a simple ``audio_data[-1]``). + + Two paths, selected at runtime by ``VLLM_VIDEO_AUDIO_DELTA_MODE``: + * fast — only D2H the new tail. Per-call cost ∝ new chunks. + * slow — full cat + D2H each call. Per-call cost ∝ total history. + Retained for A/B; remove once downstream callers confirm. + """ + audio_data = cls._get_audio_data(result) + if audio_data is None: + return None, chunks_drained + + if video_stream_envs.VLLM_VIDEO_AUDIO_DELTA_MODE == "slow": + return cls._delta_slow(audio_data, chunks_drained) + return cls._delta_fast(audio_data, chunks_drained) + + @staticmethod + def _get_audio_data(result: OmniRequestOutput): + """Navigate OmniRequestOutput → multimodal_output['audio']. None on miss.""" + request_output = getattr(result, "request_output", None) + if request_output is None: + return None + outputs = getattr(request_output, "outputs", None) + if not isinstance(outputs, list) or not outputs: + return None + mm_output = getattr(outputs[0], "multimodal_output", None) + if not isinstance(mm_output, dict): + return None + return mm_output.get("audio") + + @classmethod + def _delta_fast( + cls, + audio_data, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Emit only tensors appended since the last call.""" + # Single tensor: output_processor hands us one tensor before it becomes a + # list (see output_processor.py:89). Treat it as chunk #0. + if not isinstance(audio_data, list): + if chunks_drained >= 1: + return None, chunks_drained + tail_np = cls._tensor_to_1d_np(audio_data) + return cls._encode_tail(tail_np, chunks_drained, new_drained=1, is_first=True) + + n = len(audio_data) + if n <= chunks_drained: + return None, chunks_drained + + new_chunks = audio_data[chunks_drained:] + tail = new_chunks[0] if len(new_chunks) == 1 else torch.cat(new_chunks, dim=-1) + tail_np = cls._tensor_to_1d_np(tail) + return cls._encode_tail(tail_np, chunks_drained, new_drained=n, is_first=(chunks_drained == 0)) + + @classmethod + def _delta_slow( + cls, + audio_data, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Pre-fix behaviour: concat everything each call and slice on CPU.""" + if isinstance(audio_data, list): + if not audio_data: + return None, chunks_drained + audio_tensor = torch.cat(audio_data, dim=-1) + new_drained = len(audio_data) + else: + audio_tensor = audio_data + new_drained = 1 + + full_np = cls._tensor_to_1d_np(audio_tensor) + if full_np is None: + return None, chunks_drained + # chunks_drained doesn't map directly to sample offset without tracking + # per-chunk lengths, so we re-derive: replay the tail that corresponds + # to chunks appended since last call by slicing off the part produced + # by the already-drained prefix. For slow path this is intentionally + # wasteful — the point is to reproduce the pre-fix hot loop. + if chunks_drained == 0: + tail_np = full_np + else: + # Recover prefix length by re-concatenating the already-drained + # prefix tensors (cost intentionally identical to the baseline + # implementation this was lifted from). + if isinstance(audio_data, list) and chunks_drained < len(audio_data): + prefix_len = sum(int(t.shape[-1]) for t in audio_data[:chunks_drained]) + tail_np = full_np[prefix_len:] + else: + tail_np = full_np[0:0] + return cls._encode_tail(tail_np, chunks_drained, new_drained=new_drained, is_first=(chunks_drained == 0)) + + @classmethod + def _encode_tail( + cls, + tail_np, + old_drained: int, + *, + new_drained: int, + is_first: bool, + ) -> tuple[str | None, int]: + """Strip the CausalConv leading artifact on first emit, then b64-encode.""" + if tail_np is None or len(tail_np) == 0: + return None, new_drained + if is_first and len(tail_np) > _CODEC_FRAME_SAMPLES * 2: + tail_np = tail_np[_CODEC_FRAME_SAMPLES:] + if len(tail_np) == 0: + return None, new_drained + try: + return cls._encode_audio_wav_b64(tail_np), new_drained + except Exception: + logger.exception("Failed to encode audio delta WAV") + return None, old_drained + + @staticmethod + def _tensor_to_1d_np(t): + """Tensor → flat float32 numpy on CPU. None on failure.""" + if t is None or not hasattr(t, "float"): + return None + arr = t.float().detach().cpu().numpy() + if arr.ndim > 1: + arr = arr.flatten() + return arr + + @staticmethod + def _encode_audio_wav_b64(audio_np) -> str: + """Encode numpy float32 audio to base64 WAV (24kHz).""" + from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin + from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio + + audio_obj = CreateAudio( + audio_tensor=audio_np, + sample_rate=24000, + response_format="wav", + speed=1.0, + stream_format="audio", + base64_encode=True, + ) + mixin = AudioMixin() + resp = mixin.create_audio(audio_obj) + return resp.audio_data + + @staticmethod + def _extract_text_delta( + result: OmniRequestOutput, + previous_text: str, + ) -> tuple[str, str]: + """Extract incremental text delta from OmniRequestOutput.""" + if result.final_output_type != "text": + return "", previous_text + + request_output = getattr(result, "request_output", None) + if request_output is None: + return "", previous_text + + outputs = getattr(request_output, "outputs", None) + if not isinstance(outputs, list) or not outputs: + return "", previous_text + + text = getattr(outputs[0], "text", None) + if not isinstance(text, str) or not text: + return "", previous_text + + if text.startswith(previous_text): + return text[len(previous_text) :], text + return text, text + + # ------------------------------------------------------------------ + # Preprocessing + # ------------------------------------------------------------------ + + async def _preprocess_to_engine_prompt(self, request) -> Any: + """Use the chat handler's preprocessing to build an engine prompt.""" + handler = self._chat_service + renderer = handler.renderer + + _conversation, engine_prompts = await handler._preprocess_chat( + request, + request.messages, + default_template=getattr(request, "chat_template", None) or handler.chat_template, + default_template_content_format=handler.chat_template_content_format, + renderer=renderer, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + add_special_tokens=request.add_special_tokens, + ) + return engine_prompts[0] + + # ------------------------------------------------------------------ + # Utilities + # ------------------------------------------------------------------ + + _text_only_message = staticmethod(text_only_message) + + async def _send_error(self, websocket: WebSocket, message: str) -> None: + """Send an error message to the client.""" + try: + await websocket.send_json({"type": "error", "message": message}) + except Exception: + pass diff --git a/vllm_omni/entrypoints/openai/video_frame_filter.py b/vllm_omni/entrypoints/openai/video_frame_filter.py new file mode 100644 index 000000000..4ea427835 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_frame_filter.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EVS (Efficient Video Sampling) frame pre-filter for streaming video input. + +Lightweight pixel-level similarity filter that runs before frames reach the +vision encoder. For static or slow-moving scenes this can reduce the number +of frames by 2-5x, proportionally cutting encoder compute and KV-cache usage. + +Usage: + filter = FrameSimilarityFilter(threshold=0.95) + for jpeg_bytes in incoming_frames: + if filter.should_retain(jpeg_bytes): + buffer.append(jpeg_bytes) +""" + +from __future__ import annotations + +import io + +import numpy as np +from PIL import Image + +_DEFAULT_THRESHOLD = 0.95 +_DEFAULT_THUMBNAIL_SIZE = 64 + + +class FrameSimilarityFilter: + """Drop near-duplicate frames based on pixel-level similarity. + + Each incoming JPEG frame is down-scaled to a small thumbnail and compared + against the last *retained* frame. If the normalised similarity exceeds + ``threshold`` the frame is considered redundant and dropped. + + Args: + threshold: Similarity score in [0, 1] above which a frame is dropped. + Higher values keep more frames (less aggressive filtering). + Default 0.95 works well for typical webcam / surveillance feeds. + thumbnail_size: Edge length of the square thumbnail used for + comparison. Larger values are more accurate but slower. + """ + + def __init__( + self, + threshold: float = _DEFAULT_THRESHOLD, + thumbnail_size: int = _DEFAULT_THUMBNAIL_SIZE, + ) -> None: + if not 0.0 <= threshold <= 1.0: + raise ValueError(f"threshold must be in [0, 1], got {threshold}") + if thumbnail_size < 1: + raise ValueError(f"thumbnail_size must be >= 1, got {thumbnail_size}") + + self._threshold = threshold + self._thumbnail_size = thumbnail_size + self._last_retained: np.ndarray | None = None + self._retained_count = 0 + self._dropped_count = 0 + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def should_retain(self, frame_jpeg: bytes) -> bool: + """Return ``True`` if *frame_jpeg* is sufficiently different from the + last retained frame and should be kept in the buffer.""" + current = self._decode_and_resize(frame_jpeg) + + if self._last_retained is None: + self._last_retained = current + self._retained_count += 1 + return True + + similarity = self._compute_similarity(self._last_retained, current) + if similarity >= self._threshold: + self._dropped_count += 1 + return False + + self._last_retained = current + self._retained_count += 1 + return True + + def reset(self) -> None: + """Clear internal state so the next frame is always retained.""" + self._last_retained = None + self._retained_count = 0 + self._dropped_count = 0 + + @property + def stats(self) -> dict[str, int | float]: + """Return filtering statistics.""" + total = self._retained_count + self._dropped_count + return { + "retained_count": self._retained_count, + "dropped_count": self._dropped_count, + "total_count": total, + "drop_rate": self._dropped_count / total if total > 0 else 0.0, + } + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + @staticmethod + def _compute_similarity(a: np.ndarray, b: np.ndarray) -> float: + """Normalised pixel similarity in [0, 1]. 1 = identical.""" + mse = float(np.mean((a.astype(np.float32) - b.astype(np.float32)) ** 2)) + return 1.0 - mse / (255.0 * 255.0) + + def _decode_and_resize(self, jpeg_bytes: bytes) -> np.ndarray: + """Decode JPEG and resize to a small square thumbnail for fast + comparison.""" + img = Image.open(io.BytesIO(jpeg_bytes)) + img = img.resize( + (self._thumbnail_size, self._thumbnail_size), + Image.Resampling.BILINEAR, + ).convert("RGB") + return np.asarray(img, dtype=np.uint8) diff --git a/vllm_omni/entrypoints/openai/video_stream_context.py b/vllm_omni/entrypoints/openai/video_stream_context.py new file mode 100644 index 000000000..e72dfcd41 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_context.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small helpers shared by streaming video handlers.""" + +from __future__ import annotations + +from typing import Any + + +def text_only_message(message: dict[str, Any]) -> dict[str, Any]: + """Return a history message with multimodal content stripped out.""" + role = message.get("role", "user") + content = message.get("content", "") + + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": ""} + + text_parts: list[str] = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + text = part.get("text") + if isinstance(text, str): + text_parts.append(text) + + return {"role": role, "content": "".join(text_parts)} diff --git a/vllm_omni/entrypoints/openai/video_stream_envs.py b/vllm_omni/entrypoints/openai/video_stream_envs.py new file mode 100644 index 000000000..f8571e7f9 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_envs.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Environment variables for the streaming video OpenAI entrypoint.""" + +import logging +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + VLLM_VIDEO_AUDIO_DELTA_MODE: Literal["fast", "slow"] = "fast" + VLLM_VIDEO_ASYNC_CHUNK: Literal["on", "off"] = "on" + +logger = logging.getLogger(__name__) +_warned_invalid_envs: set[tuple[str, str]] = set() +_VIDEO_AUDIO_DELTA_MODE = "VLLM_VIDEO_AUDIO_DELTA_MODE" +_VIDEO_ASYNC_CHUNK = "VLLM_VIDEO_ASYNC_CHUNK" + + +def _choice_env( + name: str, + default: str, + allowed: tuple[str, ...], +) -> str: + value = os.getenv(name, default).strip().lower() + if value in allowed: + return value + warning_key = (name, value) + if warning_key not in _warned_invalid_envs: + logger.warning("%s=%s not recognized; falling back to %r", name, value, default) + _warned_invalid_envs.add(warning_key) + return default + + +environment_variables: dict[str, Callable[[], str]] = { + _VIDEO_AUDIO_DELTA_MODE: lambda: _choice_env( + _VIDEO_AUDIO_DELTA_MODE, + "fast", + ("fast", "slow"), + ), + _VIDEO_ASYNC_CHUNK: lambda: _choice_env( + _VIDEO_ASYNC_CHUNK, + "on", + ("on", "off"), + ), +} + + +def __getattr__(name: str) -> str: + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return list(environment_variables.keys()) + + +def is_set(name: str) -> bool: + if name in environment_variables: + return name in os.environ + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/entrypoints/openai/video_stream_session.py b/vllm_omni/entrypoints/openai/video_stream_session.py new file mode 100644 index 000000000..ead3b1ec5 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_session.py @@ -0,0 +1,439 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Video streaming session manager and WebSocket handler. + +Provides ``VideoStreamConfig``, ``VideoStreamSession`` (frame/audio buffer +with EVS filtering), and ``VideoStreamHandler`` (WebSocket session loop). +Builds standard ``ChatCompletionRequest`` so ``OmniOpenAIServingChat`` is +reused with zero changes. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +from collections import deque +from typing import Any + +from fastapi import WebSocket, WebSocketDisconnect +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter + +logger = init_logger(__name__) + +_DEFAULT_IDLE_TIMEOUT = 60.0 +_DEFAULT_CONFIG_TIMEOUT = 10.0 +_MAX_FRAME_BYTES = 10 * 1024 * 1024 # 10 MB per frame + + +class VideoStreamConfig: + """Per-session configuration sent by the client in ``session.config``.""" + + __slots__ = ( + "model", + "modalities", + "max_frames", + "num_sample_frames", + "evs_enabled", + "evs_threshold", + ) + + def __init__( + self, + model: str = "", + modalities: list[str] | None = None, + max_frames: int = 64, + num_sample_frames: int = 16, + evs_enabled: bool = True, + evs_threshold: float = 0.95, + ) -> None: + self.model = model + self.modalities = modalities if modalities is not None else ["text"] + self.max_frames = max_frames + self.num_sample_frames = num_sample_frames + self.evs_enabled = evs_enabled + self.evs_threshold = evs_threshold + + # JSON doesn't distinguish int/float, so float fields accept both. + _FIELD_TYPES: dict[str, tuple[type, ...]] = { + "model": (str,), + "modalities": (list,), + "max_frames": (int,), + "num_sample_frames": (int,), + "evs_enabled": (bool,), + "evs_threshold": (int, float), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> VideoStreamConfig: + known = set(cls.__slots__) + filtered = {k: v for k, v in data.items() if k in known} + + for field, accepted in cls._FIELD_TYPES.items(): + if field in filtered and not isinstance(filtered[field], accepted): + names = "/".join(t.__name__ for t in accepted) + raise TypeError(f"Invalid type for '{field}': expected {names}, got {type(filtered[field]).__name__}") + + return cls(**filtered) + + +class VideoStreamSession: + """Manages frame and audio buffers for a single streaming session. + + Frames are optionally filtered through ``FrameSimilarityFilter`` (EVS) + before being stored in a fixed-size ring buffer. ``build_chat_request`` + merges buffers into a ``ChatCompletionRequest`` with ``image_url`` (+ + ``audio_url``) content blocks. + """ + + def __init__(self, config: VideoStreamConfig) -> None: + self._config = config + self._frame_filter: FrameSimilarityFilter | None = ( + FrameSimilarityFilter(threshold=config.evs_threshold) if config.evs_enabled else None + ) + self._frames: deque[bytes] = deque(maxlen=config.max_frames) + self._audio_chunks: list[bytes] = [] + + def add_frame(self, jpeg_bytes: bytes) -> bool: + """Add a JPEG frame after EVS filtering. Returns ``True`` if kept.""" + if len(jpeg_bytes) > _MAX_FRAME_BYTES: + raise ValueError(f"Frame too large: {len(jpeg_bytes)} bytes (limit {_MAX_FRAME_BYTES})") + + if self._frame_filter and not self._frame_filter.should_retain(jpeg_bytes): + return False + + self._frames.append(jpeg_bytes) + return True + + def sample_frames(self) -> list[bytes]: + """Return up to ``num_sample_frames`` uniformly sampled frames.""" + n = len(self._frames) + k = min(n, self._config.num_sample_frames) + if k == 0: + return [] + if k == 1: + return [self._frames[-1]] + if k >= n: + return list(self._frames) + indices = [int(i * (n - 1) / (k - 1)) for i in range(k)] + return [self._frames[i] for i in indices] + + @property + def frame_count(self) -> int: + return len(self._frames) + + def add_audio_chunk(self, pcm_bytes: bytes) -> None: + """Append raw PCM 16 kHz audio bytes.""" + self._audio_chunks.append(pcm_bytes) + + def clear_audio(self) -> None: + """Clear audio buffer after a query is submitted.""" + self._audio_chunks.clear() + + @property + def has_audio(self) -> bool: + return bool(self._audio_chunks) + + def build_chat_request(self, query_text: str) -> ChatCompletionRequest: + """Build a ``ChatCompletionRequest`` from the current buffers.""" + sampled = self.sample_frames() + content_parts: list[dict[str, Any]] = [] + + for frame in sampled: + frame_b64 = base64.b64encode(frame).decode() + content_parts.append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{frame_b64}", + }, + } + ) + + if self.has_audio: + combined_pcm = b"".join(self._audio_chunks) + audio_b64 = base64.b64encode(combined_pcm).decode() + content_parts.append( + { + "type": "audio_url", + "audio_url": { + "url": f"data:audio/L16;rate=16000;base64,{audio_b64}", + }, + } + ) + + content_parts.append({"type": "text", "text": query_text}) + + messages: list[dict[str, Any]] = [ + {"role": "user", "content": content_parts}, + ] + + request_dict: dict[str, Any] = { + "model": self._config.model, + "messages": messages, + "stream": True, + } + + if self.has_audio: + request_dict["mm_processor_kwargs"] = { + "use_audio_in_video": True, + } + + return ChatCompletionRequest(**request_dict) + + @property + def evs_stats(self) -> dict[str, int | float] | None: + return self._frame_filter.stats if self._frame_filter else None + + +class VideoStreamHandler: + """Drives a ``VideoStreamSession`` over a FastAPI ``WebSocket``. + + Instantiate once at server startup and call ``handle_session`` for each + incoming connection. Uses two concurrent tasks (_reader + _processor) + so frames keep arriving while a query is being processed. + """ + + def __init__( + self, + chat_handler: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + ) -> None: + self._chat_handler = chat_handler + self._idle_timeout = idle_timeout + self._config_timeout = config_timeout + + async def handle_session(self, websocket: WebSocket) -> None: + """Main session loop for one WebSocket connection.""" + await websocket.accept() + + try: + config = await self._receive_config(websocket) + if config is None: + return + + session = VideoStreamSession(config) + logger.info( + "Video stream session started: model=%s modalities=%s max_frames=%d evs=%s", + config.model, + config.modalities, + config.max_frames, + config.evs_enabled, + ) + + msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() + + async def _reader() -> None: + try: + while True: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._idle_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Idle timeout: no message received") + await msg_queue.put(None) + return + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON") + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "Messages must be JSON objects") + continue + + await msg_queue.put(msg) + + if msg.get("type") == "video.done": + return + except WebSocketDisconnect: + await msg_queue.put(None) + except Exception: + await msg_queue.put(None) + raise + + async def _processor() -> None: + while True: + msg = await msg_queue.get() + if msg is None: + return + + msg_type = msg.get("type") + + if msg_type == "video.frame": + data_b64 = msg.get("data", "") + try: + jpeg_bytes = base64.b64decode(data_b64) + except Exception: + await self._send_error(websocket, "Invalid base64 in video.frame") + continue + try: + session.add_frame(jpeg_bytes) + except ValueError as exc: + await self._send_error(websocket, str(exc)) + continue + + elif msg_type == "audio.chunk": + data_b64 = msg.get("data", "") + try: + pcm_bytes = base64.b64decode(data_b64) + except Exception: + await self._send_error(websocket, "Invalid base64 in audio.chunk") + continue + session.add_audio_chunk(pcm_bytes) + + elif msg_type == "video.query": + query_text = msg.get("text", "") + if not query_text: + await self._send_error( + websocket, + "video.query requires a non-empty 'text' field", + ) + continue + await self._handle_query(websocket, session, query_text) + + elif msg_type == "video.done": + evs = session.evs_stats + if evs is not None: + await websocket.send_json({"type": "response.evs_stats", **evs}) + await websocket.send_json({"type": "session.done"}) + return + + else: + await self._send_error(websocket, f"Unknown message type: {msg_type}") + + reader_task = asyncio.create_task(_reader()) + try: + await _processor() + finally: + reader_task.cancel() + try: + await reader_task + except asyncio.CancelledError: + pass + except Exception: + logger.debug( + "Reader task raised an exception", + exc_info=True, + ) + + except WebSocketDisconnect: + logger.info("Video stream: client disconnected") + except Exception: + logger.exception("Video stream session error") + try: + await self._send_error(websocket, "Internal server error") + except Exception: + pass + + async def _receive_config(self, websocket: WebSocket) -> VideoStreamConfig | None: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._config_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.config") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.config") + return None + + if not isinstance(msg, dict) or msg.get("type") != "session.config": + await self._send_error( + websocket, + f"Expected session.config, got: {msg.get('type') if isinstance(msg, dict) else type(msg).__name__}", + ) + return None + + try: + return VideoStreamConfig.from_dict(msg) + except TypeError as exc: + await self._send_error(websocket, str(exc)) + return None + + async def _handle_query( + self, + websocket: WebSocket, + session: VideoStreamSession, + query_text: str, + ) -> None: + """Build a ChatCompletionRequest, run it, and stream the response.""" + if session.frame_count == 0: + await self._send_error( + websocket, + "No frames available. Send video.frame before video.query.", + ) + return + + request = session.build_chat_request(query_text) + await websocket.send_json({"type": "response.start"}) + + # After response.start, the protocol contract requires us to always + # send response.text.done — even on error / exception paths. + text_parts: list[str] = [] + try: + # raw_request=None: serving_chat.py guards with `if raw_request:` + generator = await self._chat_handler.create_chat_completion(request, raw_request=None) + + if isinstance(generator, ErrorResponse): + error_msg = generator.error.message if generator.error else "Unknown error" + await self._send_error(websocket, error_msg) + else: + async for chunk in generator: + if isinstance(chunk, str): + for delta in self._parse_sse_deltas(chunk): + text_parts.append(delta) + await websocket.send_json( + { + "type": "response.text.delta", + "delta": delta, + } + ) + + except Exception: + logger.exception("Query failed") + await self._send_error(websocket, "Query processing failed") + + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + + session.clear_audio() + + @staticmethod + def _parse_sse_deltas(chunk: str) -> list[str]: + """Extract content deltas from SSE-formatted chunk.""" + deltas: list[str] = [] + for line in chunk.split("\n"): + line = line.strip() + if not line or not line.startswith("data: ") or line == "data: [DONE]": + continue + try: + data = json.loads(line[6:]) + delta = data.get("choices", [{}])[0].get("delta", {}).get("content", "") + if delta: + deltas.append(delta) + except (json.JSONDecodeError, IndexError, AttributeError): + pass + return deltas + + @staticmethod + async def _send_error(websocket: WebSocket, message: str) -> None: + try: + await websocket.send_json({"type": "error", "message": message}) + except Exception: + pass