"""WebSocket handler for streaming text input TTS. Accepts text incrementally via WebSocket, buffers and splits at sentence boundaries, and generates audio per sentence using the existing TTS pipeline. Protocol: Client -> Server: {"type": "session.config", ...} # Session config (sent once first) {"type": "input.text", "text": "..."} # Text chunks {"type": "input.done"} # End of input Server -> Client (default, word_timestamps=false): {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "wav"} ... {"type": "audio.done", "sentence_index": 0} {"type": "session.done", "total_sentences": N} {"type": "error", "message": "..."} Server -> Client (when word_timestamps=true): {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm"} {"type": "audio.chunk", "sentence_index": 0, "chunk_id": 0, "audio_b64": "", "timestamps": null} ... {"type": "audio.chunk", "audio_b64": "", "timestamps": [{"word", "start_ms", "end_ms"}, ...]} {"type": "audio.done", "sentence_index": 0} # Audio is JSON base64 PCM (not binary). A trailing empty-audio chunk carries the # full sentence-relative alignment. timestamps: list = aligned, [] = silence, null = failed. """ import asyncio import base64 import json from contextlib import aclosing from fastapi import WebSocket, WebSocketDisconnect from pydantic import ValidationError from vllm.logger import init_logger from vllm_omni.entrypoints.openai.protocol.audio import ( OpenAICreateSpeechRequest, StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.text_splitter import ( SPLIT_CLAUSE, SPLIT_SENTENCE, SentenceSplitter, ) from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError from vllm_omni.utils.forced_aligner import align as forced_align logger = init_logger(__name__) _DEFAULT_IDLE_TIMEOUT = 30.0 # seconds _DEFAULT_CONFIG_TIMEOUT = 10.0 # seconds _PCM_SAMPLE_RATE = 24000 _BYTES_PER_SAMPLE = 2 # 16-bit mono PCM _MAX_CONFIG_MESSAGE_SIZE = 4 * 1024 * 1024 # allow large ref_audio payloads _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 class OmniStreamingSpeechHandler: """Handles WebSocket sessions for streaming text-input TTS. Each WebSocket connection is an independent session. Text arrives incrementally, is split at sentence boundaries, and audio is generated per sentence using the existing OmniOpenAIServingSpeech pipeline. Args: speech_service: The existing TTS serving instance (reused for validation and audio generation). idle_timeout: Max seconds to wait for a message before closing. config_timeout: Max seconds to wait for the initial session.config. """ def __init__( self, speech_service: OmniOpenAIServingSpeech, idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, ) -> None: self._speech_service = speech_service self._idle_timeout = idle_timeout self._config_timeout = config_timeout async def handle_session(self, websocket: WebSocket) -> None: """Main session loop for a single WebSocket connection.""" await websocket.accept() try: # 1. Wait for session.config config = await self._receive_config(websocket) if config is None: return # Error already sent, connection closing # Validate model if specified if config.model and hasattr(self._speech_service, "_check_model"): error = await self._speech_service._check_model( OpenAICreateSpeechRequest(input="ping", model=config.model) ) if error is not None: await self._send_error(websocket, str(error)) return boundary_re = SPLIT_CLAUSE if config.split_granularity == "clause" else SPLIT_SENTENCE splitter = SentenceSplitter(boundary_re=boundary_re) sentence_index = 0 # 2. Receive text chunks until input.done while True: try: raw = await asyncio.wait_for( websocket.receive_text(), timeout=self._idle_timeout, ) except asyncio.TimeoutError: await self._send_error(websocket, "Idle timeout: no message received") return if len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE: await self._send_error(websocket, "input.text message too large") continue try: msg = json.loads(raw) except json.JSONDecodeError: await self._send_error(websocket, "Invalid JSON message") continue if not isinstance(msg, dict): await self._send_error(websocket, "WebSocket messages must be JSON objects") continue msg_type = msg.get("type") if msg_type == "input.text": text = msg.get("text", "") if not isinstance(text, str): await self._send_error(websocket, "input.text requires a string value") continue sentences = splitter.add_text(text) for sentence in sentences: await self._generate_and_send(websocket, config, sentence, sentence_index) sentence_index += 1 elif msg_type == "input.done": # Flush remaining buffer remaining = splitter.flush() if remaining: await self._generate_and_send(websocket, config, remaining, sentence_index) sentence_index += 1 # Send session.done await websocket.send_json( { "type": "session.done", "total_sentences": sentence_index, } ) return else: await self._send_error( websocket, f"Unknown message type: {msg_type}", ) except WebSocketDisconnect: logger.info("Streaming speech: client disconnected") except Exception as e: logger.exception("Streaming speech session error: %s", e) try: await self._send_error(websocket, f"Internal error: {e}") except Exception: logger.debug("Failed to send error to streaming speech client", exc_info=True) async def _receive_config(self, websocket: WebSocket) -> StreamingSpeechSessionConfig | 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 if len(raw) > _MAX_CONFIG_MESSAGE_SIZE: await self._send_error(websocket, "session.config message too large") 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): await self._send_error(websocket, "session.config must be a JSON object") return None if msg.get("type") != "session.config": await self._send_error( websocket, f"Expected session.config, got: {msg.get('type')}", ) return None try: config = StreamingSpeechSessionConfig(**{k: v for k, v in msg.items() if k != "type"}) except ValidationError as e: await self._send_error(websocket, f"Invalid session config: {e}") return None return config async def _generate_and_send( self, websocket: WebSocket, config: StreamingSpeechSessionConfig, sentence_text: str, sentence_index: int, ) -> None: """Generate audio for a single sentence and send it over WebSocket.""" response_format = config.response_format or "wav" # Reject unmet word-timestamps preconditions early with a clear reason. if config.word_timestamps: if self._speech_service.forced_aligner_config is None: await self._send_error( websocket, "word_timestamps=true but the server was launched without " "--forced-aligner; either restart the server with that flag " "or set word_timestamps=false in session.config.", ) return if not (config.stream_audio and response_format == "pcm"): await self._send_error( websocket, "word_timestamps=true requires stream_audio=true and " "response_format='pcm' (the aligner consumes raw PCM).", ) return request = OpenAICreateSpeechRequest( input=sentence_text, model=config.model, voice=config.voice, task_type=config.task_type, language=config.language, instructions=config.instructions, response_format=response_format, speed=config.speed, max_new_tokens=config.max_new_tokens, initial_codec_chunk_frames=config.initial_codec_chunk_frames, non_streaming_mode=config.non_streaming_mode, ref_audio=config.ref_audio, ref_text=config.ref_text, x_vector_only_mode=config.x_vector_only_mode, speaker_embedding=config.speaker_embedding, stream=config.stream_audio, word_timestamps=config.word_timestamps, ) start_payload = { "type": "audio.start", "sentence_index": sentence_index, "sentence_text": sentence_text, "format": response_format, } if config.stream_audio and response_format == "pcm": # Nominal stream rate; each audio.chunk carries the authoritative # per-chunk sample_rate. start_payload["sample_rate"] = _PCM_SAMPLE_RATE if config.word_timestamps: start_payload["word_timestamps"] = True await websocket.send_json(start_payload) total_bytes = 0 generation_failed = False request_id = None try: if config.stream_audio: request_id, generator, _ = await self._speech_service._prepare_speech_generation(request) if config.word_timestamps: total_bytes = await self._stream_audio_with_alignments( websocket=websocket, request_id=request_id, generator=generator, sentence_text=sentence_text, sentence_index=sentence_index, language=config.language, ) else: async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream: async for chunk in stream: total_bytes += len(chunk) await websocket.send_bytes(chunk) else: audio_bytes, _ = await self._speech_service._generate_audio_bytes(request) total_bytes = len(audio_bytes) await websocket.send_bytes(audio_bytes) except WebSocketDisconnect: if request_id is not None: try: await self._speech_service.engine_client.abort(request_id) except Exception: logger.debug("Failed to abort streaming speech request %s", request_id, exc_info=True) raise except Exception as e: generation_failed = True logger.error("Generation failed for sentence %d: %s", sentence_index, e) await self._send_error(websocket, f"Generation failed for sentence {sentence_index}: {e}") finally: try: await websocket.send_json( { "type": "audio.done", "sentence_index": sentence_index, "total_bytes": total_bytes, "error": generation_failed, } ) except Exception: logger.debug("Failed to send audio.done for sentence %d", sentence_index, exc_info=True) async def _stream_audio_with_alignments( self, *, websocket: WebSocket, request_id: str, generator, sentence_text: str, sentence_index: int, language: str | None = None, ) -> int: """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. Forward each PCM chunk live (``timestamps: null``) while buffering the sentence audio, then run the forced aligner once over the whole sentence and emit a final empty-audio ``audio.chunk`` with the word timestamps. On aligner failure timestamps is ``null``; for silence it is ``[]`` (audio always flows regardless). """ aligner_config = self._speech_service.forced_aligner_config assert aligner_config is not None # gated by the precondition check sentence_audio = bytearray() total_bytes = 0 sample_rate = _PCM_SAMPLE_RATE chunk_id = 0 async def send_chunk( chunk: bytes, chunk_sample_rate: int, timestamps_payload: list[dict] | None, chunk_start_ms: int, chunk_end_ms: int, ) -> None: nonlocal chunk_id await websocket.send_json( { "type": "audio.chunk", "sentence_index": sentence_index, "chunk_id": chunk_id, "chunk_start_ms": chunk_start_ms, "chunk_end_ms": chunk_end_ms, "sample_rate": chunk_sample_rate, "audio_b64": base64.b64encode(chunk).decode("ascii"), "timestamps": timestamps_payload, } ) chunk_id += 1 async with aclosing( self._speech_service._generate_pcm_chunks( generator, request_id, include_sample_rate=True, ) ) as stream: async for chunk, chunk_sample_rate in stream: sample_rate = chunk_sample_rate chunk_start_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) sentence_audio.extend(chunk) chunk_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) total_bytes += len(chunk) # Audio first, timestamps after the whole sentence is aligned. await send_chunk(chunk, chunk_sample_rate, None, chunk_start_ms, chunk_end_ms) # Single alignment pass over the full sentence, then emit timestamps. # A load/config failure is permanent, so surface the reason once; audio # has already streamed, so the trailing frame still carries null. try: timestamps_payload = await self._align_sentence( audio=bytes(sentence_audio), text=sentence_text, sample_rate=sample_rate, config=aligner_config, language=language, ) except ForcedAlignerLoadError as exc: await self._send_error(websocket, f"forced aligner unavailable: {exc}") timestamps_payload = None sentence_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms) return total_bytes @staticmethod async def _align_sentence( *, audio: bytes, text: str, sample_rate: int, config, language: str | None = None, ) -> list[dict] | None: """Convert a sentence alignment into JSON word-timestamp dicts. Returns ``None`` on aligner failure, ``[]`` when it ran but produced no tokens. Monotonic, non-overlapping bounds are guaranteed by the decoder. ``language`` is forwarded to word segmentation. """ aligned = await forced_align( audio=audio, text=text, sample_rate=sample_rate, config=config, language=language, ) if aligned is None: return None return [{"word": ts.word, "start_ms": ts.start_ms, "end_ms": ts.end_ms} for ts in aligned] @staticmethod async def _send_error(websocket: WebSocket, message: str) -> None: """Send an error message to the client.""" try: await websocket.send_json( { "type": "error", "message": message, } ) except Exception: pass # Connection may already be closed; safe to ignore