# 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)