1088 lines
44 KiB
Python
1088 lines
44 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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": "...", "frame_id": "...", "pts_ms": 0}
|
|
{"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": "video.frame.ack", ...} # when frame_id is provided
|
|
{"type": "video.frames.consumed", ...} # after first engine output
|
|
{"type": "response.start"}
|
|
{"type": "response.text.delta", "delta": "..."}
|
|
{"type": "response.text.done", "text": "..."}
|
|
{"type": "response.output_audio.delta", "data": "...", "format": "wav"}
|
|
{"type": "response.output_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 enum import Enum
|
|
from typing import Any, Final, Protocol, TypeAlias, 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.entrypoints.utils import coerce_param_message_types
|
|
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
|
|
|
|
|
|
class _FrameStatus(Enum):
|
|
BAD = "bad"
|
|
|
|
|
|
_BAD_FRAME: Final = _FrameStatus.BAD
|
|
PrewarmedFrame: TypeAlias = tuple[Any, str] | _FrameStatus
|
|
|
|
|
|
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: Mapping[str, PrewarmedFrame],
|
|
*,
|
|
frame_indices: list[int] | None = None,
|
|
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
|
"""Build messages using the supplied frame selection, or sample if omitted."""
|
|
...
|
|
|
|
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: Mapping[str, PrewarmedFrame],
|
|
*,
|
|
frame_indices: list[int] | None = None,
|
|
) -> 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
|
|
frame_metadata: list[dict[str, Any]] = []
|
|
# Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index.
|
|
frame_pil_cache: dict[str, PrewarmedFrame] = {} # 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
|
|
if msg_type == "video.frame":
|
|
msg["_receiver_received_ts_ms"] = _time.monotonic() * 1000
|
|
|
|
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
|
|
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_frame_metadata = list(frame_metadata)
|
|
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:
|
|
process_kwargs: dict[str, Any] = {}
|
|
if any(metadata.get("frame_id") for metadata in query_frame_metadata):
|
|
process_kwargs["frame_metadata"] = query_frame_metadata
|
|
await self._process_query(
|
|
websocket,
|
|
config,
|
|
query_frames,
|
|
query_audio_buffer,
|
|
message_history,
|
|
query_text,
|
|
request_id,
|
|
interrupt_event,
|
|
query_prewarmed_frames,
|
|
**process_kwargs,
|
|
)
|
|
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:
|
|
retained_indices = [
|
|
index for index, frame in enumerate(frame_buffer) if frame != frame_data
|
|
]
|
|
frame_buffer[:] = [frame_buffer[index] for index in retained_indices]
|
|
frame_metadata[:] = [frame_metadata[index] for index in retained_indices]
|
|
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):
|
|
await self._send_frame_ack(
|
|
websocket,
|
|
msg,
|
|
accepted=False,
|
|
buffered_frames=len(frame_buffer),
|
|
reason="filtered",
|
|
)
|
|
continue
|
|
except Exception:
|
|
await self._send_error(websocket, "Invalid image data")
|
|
continue
|
|
max_buf = config.max_frames
|
|
dropped_frame_id: str | None = None
|
|
if len(frame_buffer) >= max_buf:
|
|
dropped = frame_buffer.pop(0)
|
|
dropped_metadata = frame_metadata.pop(0)
|
|
dropped_frame_id = dropped_metadata.get("frame_id")
|
|
frame_pil_cache.pop(dropped, None)
|
|
frame_buffer.append(frame_data)
|
|
frame_metadata.append(
|
|
{
|
|
"frame_id": msg.get("frame_id"),
|
|
"pts_ms": msg.get("pts_ms"),
|
|
"source_pts_ms": msg.get("source_pts_ms"),
|
|
"quality_profile": msg.get("quality_profile"),
|
|
"capture_ts_ms": msg.get("capture_ts_ms"),
|
|
"receiver_received_ts_ms": msg.get("_receiver_received_ts_ms"),
|
|
}
|
|
)
|
|
self.on_frame_buffered(raw_bytes, frame_data, message_history, config)
|
|
await self._send_frame_ack(
|
|
websocket,
|
|
msg,
|
|
accepted=True,
|
|
buffered_frames=len(frame_buffer),
|
|
dropped_frame_id=dropped_frame_id,
|
|
)
|
|
# 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)
|
|
# The frame may have been evicted while decoding.
|
|
if b64 in frame_buffer:
|
|
frame_pil_cache[b64] = (pil, u)
|
|
except Exception:
|
|
if b64 not in frame_buffer:
|
|
return
|
|
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: Mapping[str, PrewarmedFrame],
|
|
frame_metadata: list[dict[str, Any]] | None = None,
|
|
) -> 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
|
|
|
|
engine_kwargs: dict[str, Any] = {}
|
|
if frame_metadata:
|
|
engine_kwargs["frame_metadata"] = frame_metadata
|
|
await self._process_query_engine(
|
|
websocket,
|
|
config,
|
|
frame_buffer,
|
|
audio_buffer,
|
|
message_history,
|
|
query_text,
|
|
request_id,
|
|
interrupt_event,
|
|
prewarmed_frames,
|
|
**engine_kwargs,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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: Mapping[str, PrewarmedFrame],
|
|
frame_metadata: list[dict[str, Any]] | None = None,
|
|
) -> None:
|
|
"""Direct engine_client.generate() path for async_chunk audio."""
|
|
engine_client = self._engine_client
|
|
assert engine_client is not None, "_process_query must validate the engine client"
|
|
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest,
|
|
)
|
|
|
|
# Select once from this turn's snapshot so prompt images and the
|
|
# consumed event refer to exactly the same frames, including decode failures.
|
|
frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames)
|
|
messages, user_message = self.build_engine_prompt(
|
|
config,
|
|
frame_buffer,
|
|
audio_buffer,
|
|
message_history,
|
|
query_text,
|
|
prewarmed_frames,
|
|
frame_indices=frame_indices,
|
|
)
|
|
|
|
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,
|
|
}
|
|
sampling_kwargs: dict[str, Any] = {}
|
|
try:
|
|
chat_request = ChatCompletionRequest(**request_kwargs)
|
|
if config.sampling_params_list:
|
|
params = self._chat_service._to_sampling_params_list(config.sampling_params_list)
|
|
# Explicit params must retain the engine's streaming output mode.
|
|
sampling_kwargs["sampling_params_list"] = coerce_param_message_types(params, is_streaming=True)
|
|
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
|
|
decoded_ready_ts_ms = _time.monotonic() * 1000
|
|
selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else []
|
|
model_selected_ts_ms = _time.monotonic() * 1000
|
|
|
|
await websocket.send_json({"type": "response.start"})
|
|
text_parts: list[str] = []
|
|
text_done_sent = False
|
|
audio_chunk_count = 0
|
|
# Count emitted tensors only to identify the first audio emission.
|
|
# Default and explicit sampling parameters use DELTA audio outputs.
|
|
audio_chunks_drained = 0
|
|
previous_text = ""
|
|
interrupted = False
|
|
frames_consumed_sent = 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 = engine_client.generate(
|
|
prompt=engine_prompt,
|
|
request_id=request_id,
|
|
output_modalities=config.modalities,
|
|
**sampling_kwargs,
|
|
)
|
|
|
|
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
|
|
|
|
if not frames_consumed_sent and frame_metadata:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "video.frames.consumed",
|
|
"request_id": request_id,
|
|
"model_selected_ts_ms": model_selected_ts_ms,
|
|
"frame_ids": [
|
|
metadata["frame_id"]
|
|
for metadata in selected_metadata
|
|
if isinstance(metadata.get("frame_id"), str)
|
|
],
|
|
"frames": [
|
|
{
|
|
"frame_id": metadata.get("frame_id"),
|
|
"pts_ms": metadata.get("pts_ms"),
|
|
"source_pts_ms": metadata.get("source_pts_ms"),
|
|
"quality_profile": metadata.get("quality_profile"),
|
|
"receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"),
|
|
"decoded_ready_ts_ms": decoded_ready_ts_ms,
|
|
}
|
|
for metadata in selected_metadata
|
|
],
|
|
"latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None,
|
|
}
|
|
)
|
|
frames_consumed_sent = True
|
|
|
|
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.output_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.extend(audio_data)
|
|
else:
|
|
audio_tail_tensors.append(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.output_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.output_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).
|
|
|
|
Explicit sampling parameters are coerced to DELTA above, and AsyncOmni
|
|
applies the same coercion to defaults. The output processor drains audio
|
|
after each result, so every tensor in the current payload is new.
|
|
`chunks_drained` only identifies the first emission for the existing
|
|
leading-artifact trim.
|
|
The legacy slow setting remains accepted as an equivalent alias.
|
|
"""
|
|
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 = result
|
|
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]:
|
|
"""Encode every fresh tensor in this DELTA payload."""
|
|
new_chunks = audio_data if isinstance(audio_data, list) else [audio_data]
|
|
if not new_chunks:
|
|
return None, 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)
|
|
if tail_np is None or len(tail_np) == 0:
|
|
return None, chunks_drained
|
|
return cls._encode_tail(
|
|
tail_np,
|
|
chunks_drained,
|
|
new_drained=chunks_drained + len(new_chunks),
|
|
is_first=(chunks_drained == 0),
|
|
)
|
|
|
|
@classmethod
|
|
def _delta_slow(
|
|
cls,
|
|
audio_data,
|
|
chunks_drained: int,
|
|
) -> tuple[str | None, int]:
|
|
"""Compatibility alias for the DELTA payload contract."""
|
|
return cls._delta_fast(audio_data, chunks_drained)
|
|
|
|
@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,
|
|
base64_encode=True,
|
|
)
|
|
mixin = AudioMixin()
|
|
resp = mixin.create_audio(audio_obj)
|
|
# base64_encode=True selects the string result of create_audio().
|
|
assert isinstance(resp.audio_data, str)
|
|
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 = result
|
|
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
|
|
|
|
@staticmethod
|
|
def _sample_frame_indices(
|
|
frame_buffer: list[str],
|
|
num_frames: int,
|
|
prewarmed_frames: Mapping[str, PrewarmedFrame],
|
|
) -> list[int]:
|
|
"""Stride-sample with the last frame, then drop known bad frames without refilling."""
|
|
n_buf = len(frame_buffer)
|
|
if n_buf <= num_frames:
|
|
indices = list(range(n_buf))
|
|
else:
|
|
stride = max(1, n_buf // num_frames)
|
|
indices = [index * stride for index in range(num_frames - 1)] + [n_buf - 1]
|
|
return [index for index in indices if prewarmed_frames.get(frame_buffer[index]) is not _BAD_FRAME]
|
|
|
|
@staticmethod
|
|
async def _send_frame_ack(
|
|
websocket: WebSocket,
|
|
message: Mapping[str, Any],
|
|
*,
|
|
accepted: bool,
|
|
buffered_frames: int,
|
|
reason: str | None = None,
|
|
dropped_frame_id: str | None = None,
|
|
) -> None:
|
|
frame_id = message.get("frame_id")
|
|
if not isinstance(frame_id, str) or not frame_id:
|
|
return
|
|
ack: dict[str, Any] = {
|
|
"type": "video.frame.ack",
|
|
"frame_id": frame_id,
|
|
"pts_ms": message.get("pts_ms"),
|
|
"capture_ts_ms": message.get("capture_ts_ms"),
|
|
"accepted": accepted,
|
|
"buffered_frames": buffered_frames,
|
|
"server_receive_ts_ms": message.get("_receiver_received_ts_ms", _time.monotonic() * 1000),
|
|
}
|
|
if reason is not None:
|
|
ack["reason"] = reason
|
|
if dropped_frame_id is not None:
|
|
ack["dropped_frame_id"] = dropped_frame_id
|
|
await websocket.send_json(ack)
|