diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index 0787eb856..de2fcc24d 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -14,11 +14,12 @@ 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 import video_stream_base, video_stream_envs from vllm_omni.entrypoints.openai.serving_video_stream import ( - OmniStreamingVideoHandler, + QwenOmniStreamingVideoHandler, StreamingVideoSessionConfig, ) +from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -149,7 +150,7 @@ async def test_audio_in_video_sets_mm_processor_kwargs(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): captured_requests.append(request) return {"prompt": "x"} @@ -186,7 +187,7 @@ async def test_audio_in_video_disabled_omits_mm_processor_kwargs(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): captured_requests.append(request) return {"prompt": "x"} @@ -223,7 +224,7 @@ async def test_query_inline_audio_data_sets_mm_processor_kwargs(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): captured_requests.append(request) return {"prompt": "x"} @@ -309,7 +310,7 @@ async def test_async_chunk_mode_is_read_by_engine_path_at_runtime(monkeypatch): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): return {"prompt": "x"} @@ -374,7 +375,7 @@ async def test_new_query_cancels_in_flight_query(): query_cancelled = asyncio.Event() calls = 0 - class BlockingHandler(OmniStreamingVideoHandler): + class BlockingHandler(QwenOmniStreamingVideoHandler): async def _process_query(self, *args, **kwargs): nonlocal calls calls += 1 @@ -412,7 +413,7 @@ async def test_video_done_waits_for_in_flight_query(): allow_finish = asyncio.Event() query_finished = asyncio.Event() - class BlockingHandler(OmniStreamingVideoHandler): + class BlockingHandler(QwenOmniStreamingVideoHandler): async def _process_query(self, *args, **kwargs): query_started.set() await allow_finish.wait() @@ -452,11 +453,11 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch): release_decode.wait(timeout=2.0) return Image.open(io.BytesIO(raw_bytes)).convert("RGB") - class BlockingHandler(OmniStreamingVideoHandler): + class BlockingHandler(QwenOmniStreamingVideoHandler): async def _process_query(self, *args, **kwargs): query_started.set() - monkeypatch.setattr(serving_video_stream, "_decode_frame_bytes", blocked_decode) + monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", blocked_decode) ws = TimedWebSocket() handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) @@ -486,7 +487,7 @@ async def test_client_cannot_send_internal_frame_decode_failed_message(): captured_frames: list[list[str]] = [] frame = _b64(_make_jpeg()) - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _process_query( self, websocket, @@ -550,7 +551,7 @@ 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) + monkeypatch.setattr(video_stream_base.FrameSimilarityFilter, "should_retain", fail_should_retain) ws = TimedWebSocket() handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) @@ -579,7 +580,7 @@ async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _process_query_engine( self, websocket, @@ -594,7 +595,7 @@ async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch): ): captured_audio_lengths.append(len(audio_buffer)) - monkeypatch.setattr(serving_video_stream, "_MAX_AUDIO_BUFFER_BYTES", 4) + monkeypatch.setattr(video_stream_base, "_MAX_AUDIO_BUFFER_BYTES", 4) ws = TimedWebSocket() handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=5.0) @@ -618,7 +619,7 @@ async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch): def test_build_messages_keeps_recent_history_text_only(): - handler = OmniStreamingVideoHandler(chat_service=object()) + handler = QwenOmniStreamingVideoHandler(chat_service=object()) old_frame = _b64(_make_jpeg(1, 2, 3)) current_frame = _b64(_make_jpeg(4, 5, 6)) history = [ diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index e396c47d6..627a10705 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -132,7 +132,7 @@ from vllm_omni.entrypoints.openai.serving_video import ( ReferenceImage, ReferenceVideo, ) -from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler +from vllm_omni.entrypoints.openai.serving_video_stream import create_streaming_video_handler from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, get_default_sampling_params_list, @@ -1070,7 +1070,7 @@ async def omni_init_app_state( speech_service=state.openai_serving_speech, ) state.openai_streaming_video = ( - OmniStreamingVideoHandler( + create_streaming_video_handler( chat_service=state.openai_serving_chat, engine_client=engine_client, ) diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index 0c1fb52e9..af0083fcb 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""WebSocket handler for streaming video input understanding. +"""Qwen-Omni streaming video WebSocket handler. 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). +generates text + optional audio responses using the Qwen3-Omni multi-stage +pipeline (thinker -> talker -> code2wav). Protocol: Client -> Server: @@ -24,647 +24,35 @@ Protocol: {"type": "error", "message": "..."} """ -import asyncio -import base64 -import hashlib -import io -import json -import time as _time -import uuid -import wave -from collections.abc import Mapping +from __future__ import annotations + 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.entrypoints.openai.video_stream_base import ( + _BAD_FRAME, + _DEFAULT_CONFIG_TIMEOUT, + _DEFAULT_IDLE_TIMEOUT, + StreamingVideoSessionConfig, + VideoStreamTurnTrigger, +) +from vllm_omni.entrypoints.openai.video_stream_base import ( + OmniStreamingVideoHandler as OmniStreamingVideoHandlerBase, ) -from vllm_omni.outputs import OmniRequestOutput -logger = init_logger(__name__) +__all__ = [ + "QwenOmniStreamingVideoHandler", + "StreamingVideoSessionConfig", + "create_streaming_video_handler", +] -_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() +class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): + """Qwen-Omni pipeline: manual ``video.query`` trigger and image_pil prompts.""" -def _decode_frame_bytes(raw_bytes: bytes) -> Any: - return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + def should_trigger_turn(self, trigger: VideoStreamTurnTrigger) -> bool: + return False - -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( + def build_engine_prompt( self, config: StreamingVideoSessionConfig, frame_buffer: list[str], @@ -673,11 +61,6 @@ class OmniStreamingVideoHandler: 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) @@ -686,7 +69,6 @@ class OmniStreamingVideoHandler: 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: @@ -731,12 +113,6 @@ class OmniStreamingVideoHandler: 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)) @@ -745,228 +121,32 @@ class OmniStreamingVideoHandler: return messages, user_message - # ------------------------------------------------------------------ - # Audio helpers - # ------------------------------------------------------------------ + def on_turn_complete( + self, + message_history: list[dict[str, Any]], + user_message: dict[str, Any], + response_text: str, + ) -> None: + message_history.append(user_message) + message_history.append({"role": "assistant", "content": response_text}) - @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() + _build_messages = build_engine_prompt - @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]``). +def create_streaming_video_handler( + chat_service: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + engine_client: Any | None = None, +) -> OmniStreamingVideoHandlerBase: + """Create the handler for ``/v1/video/chat/stream``. - 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, Mapping): - 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 + Returns :class:`QwenOmniStreamingVideoHandler` today. Additional pipelines + can be selected here in follow-up PRs. + """ + return QwenOmniStreamingVideoHandler( + chat_service=chat_service, + idle_timeout=idle_timeout, + config_timeout=config_timeout, + engine_client=engine_client, + ) diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py new file mode 100644 index 000000000..ffd12c2b2 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -0,0 +1,974 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base WebSocket handler for streaming video input understanding. + +Shared session loop, frame/audio buffering, EVS pre-filter, prewarm, +interrupt handling, and engine ``generate()`` streaming. Pipeline-specific +behavior (trigger rules, prompt shape, history) is supplied by subclasses +via :class:`VideoStreamPipelineHooks`. + +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 collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +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") + + +@runtime_checkable +class VideoStreamPipelineHooks(Protocol): + """Pipeline-specific hooks for streaming video handlers.""" + + def should_trigger_turn(self, trigger: "VideoStreamTurnTrigger") -> bool: + """Return True to auto-start a turn after a new frame (no ``video.query``).""" + ... + + def build_engine_prompt( + 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 and the current user message.""" + ... + + def on_turn_complete( + self, + message_history: list[dict[str, Any]], + user_message: dict[str, Any], + response_text: str, + ) -> None: + """Update session state after a successful turn.""" + ... + + +@dataclass(frozen=True) +class VideoStreamTurnTrigger: + """Snapshot passed to :meth:`OmniStreamingVideoHandler.should_trigger_turn`.""" + + frame_count: int + is_generating: bool + config: "StreamingVideoSessionConfig" + + +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: + """Base handler for WebSocket streaming video sessions. + + Subclasses implement :class:`VideoStreamPipelineHooks` to customize turn + triggering, prompt construction, and history updates. + """ + + def should_trigger_turn(self, trigger: VideoStreamTurnTrigger) -> bool: + """Auto-trigger after ``video.frame`` when True (default: never).""" + return False + + def build_engine_prompt( + 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]]: + raise NotImplementedError + + def on_turn_complete( + self, + message_history: list[dict[str, Any]], + user_message: dict[str, Any], + response_text: str, + ) -> None: + raise NotImplementedError + + def create_message_history(self, config: StreamingVideoSessionConfig) -> Any: + """Per-session conversation state (default: empty OpenAI-style list).""" + return [] + + def on_frame_buffered( + self, + raw_bytes: bytes, + frame_b64: str, + message_history: Any, + config: StreamingVideoSessionConfig, + ) -> None: + """Hook after a frame is accepted into the session buffer.""" + del raw_bytes, frame_b64, message_history, config + + 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: Any = self.create_message_history(config) + 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 _start_query_turn(*, query_text: str) -> None: + """Schedule a new inference turn from the current buffers.""" + nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task + + await _cancel_active_query() + + if not frame_buffer: + await self._send_error(websocket, "No frames buffered") + return + + 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()) + + 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) + self.on_frame_buffered(raw_bytes, frame_data, message_history, config) + # 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) + + is_generating = active_request_id is not None or ( + query_task is not None and not query_task.done() + ) + if self.should_trigger_turn( + VideoStreamTurnTrigger( + frame_count=len(frame_buffer), + is_generating=is_generating, + config=config, + ) + ): + await _start_query_turn(query_text="") + + 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": + 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 + + await _start_query_turn(query_text=query_text) + + 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_engine_prompt( + 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) + self.on_turn_complete(message_history, user_message, 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}) + + # ------------------------------------------------------------------ + # 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, Mapping): + 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