Refactor: extract OmniStreamingVideoHandler base and QwenOmniStreamingVideoHandler (#4424)
Signed-off-by: NumberWan <wantszkin2003@gmail.com> Signed-off-by: lishunyang12 <125541396+lishunyang12@users.noreply.github.com> Co-authored-by: lishunyang12 <125541396+lishunyang12@users.noreply.github.com>
This commit is contained in:
NumberWanandlishunyang12 authored and GitHub committed
1 parent
e957d16f8d
commit
d744940b18
4 files changed
+1041
-886
No files matched your search
@@ -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 = [
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
Reference in new issue
Block a user