Qwen3-tts/docker/patch_qwen3_tts_runtime.py

#!/usr/bin/env python3
"""Runtime patches for the vllm-omni Qwen3-TTS container.

- Propagate WebSocket session.config.seed into every generated speech segment.
- Log WebSocket config and segment generation on the real 8091 endpoint.
"""
from pathlib import Path

STREAM = Path('/usr/local/lib/python3.12/dist-packages/vllm_omni/entrypoints/openai/serving_speech_stream.py')
AUDIO = Path('/usr/local/lib/python3.12/dist-packages/vllm_omni/entrypoints/openai/protocol/audio.py')


def replace_once(text: str, old: str, new: str, label: str) -> str:
    if new in text:
        return text
    if old not in text:
        raise RuntimeError(f'Patch anchor not found: {label}')
    return text.replace(old, new, 1)


def patch_seed_field() -> None:
    """Add an optional `seed` field to StreamingSpeechSessionConfig.

    The streaming WebSocket config model in protocol/audio.py ships without a
    `seed` attribute, but the stream server reads `config.seed` (see seed
    propagation patch below). Without this field every session.config fails with
    AttributeError: 'StreamingSpeechSessionConfig' object has no attribute 'seed'.
    """
    text = AUDIO.read_text()
    # Anchor on this class's unique docstring + first field: the task_type line
    # is not unique (it also appears in SpeechBatchItem / BatchSpeechRequest).
    new = replace_once(
        text,
        '    """Configuration sent as the first WebSocket message for streaming TTS."""\n\n'
        '    model: str | None = None\n',
        '    """Configuration sent as the first WebSocket message for streaming TTS."""\n\n'
        '    seed: int | None = None\n'
        '    model: str | None = None\n',
        'StreamingSpeechSessionConfig.seed field',
    )
    if new != text:
        AUDIO.write_text(new)
        print('[runtime-patch] StreamingSpeechSessionConfig.seed field added', flush=True)
    else:
        print('[runtime-patch] StreamingSpeechSessionConfig.seed field already present', flush=True)


_PIPELINE_CODE = '''

async def _pipelined_handle_session(self, websocket):
    """Pipeline: N+1 synthesis overlaps with N audio transfer to client.

    A synthesizer task submits each sentence to vllm as soon as the previous
    one is submitted (not after it is fully sent), so vllm starts work on N+1
    the moment GPU capacity is free from N.  A separate sender task streams
    audio chunks to the WebSocket in order.  Lookahead depth is controlled by
    QWEN3_TTS_PIPELINE_LOOKAHEAD (default 1).
    """
    await websocket.accept()
    try:
        config = await self._receive_config(websocket)
        if config is None:
            return
        logger.info(
            "Streaming speech config: model=%s voice=%s task_type=%s language=%s"
            " seed=%s split=%s stream_audio=%s instructions=%r speaker_embedding_dim=%s",
            config.model, config.voice, config.task_type, config.language,
            config.seed, config.split_granularity, config.stream_audio, config.instructions,
            len(config.speaker_embedding) if config.speaker_embedding is not None else None,
        )
        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)
        response_format = config.response_format or "wav"

        sentence_q = asyncio.Queue()
        lookahead = int(_os.environ.get("QWEN3_TTS_PIPELINE_LOOKAHEAD", "1"))
        ready_q = asyncio.Queue(maxsize=max(1, lookahead))
        pending_rids = set()  # all request IDs submitted but not yet completed

        def _build_req(text):
            return OpenAICreateSpeechRequest(
                input=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,
                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, seed=config.seed,
            )

        async def _abort_all():
            for rid in list(pending_rids):
                try:
                    await self._speech_service.engine_client.abort(rid)
                except Exception:
                    pass
            pending_rids.clear()

        async def _synthesizer():
            while True:
                item = await sentence_q.get()
                if item is None:
                    await ready_q.put(None)
                    return
                idx, text = item
                logger.info(
                    "Streaming speech segment: index=%s chars=%s voice=%s"
                    " task_type=%s language=%s seed=%s speaker_embedding_dim=%s",
                    idx, len(text), config.voice, config.task_type, config.language,
                    config.seed,
                    len(config.speaker_embedding) if config.speaker_embedding is not None else None,
                )
                try:
                    req = _build_req(text)
                    if config.stream_audio:
                        rid, gen, _ = await self._speech_service._prepare_speech_generation(req)
                        pending_rids.add(rid)
                        await ready_q.put(("stream", idx, text, rid, gen))
                    else:
                        audio, _ = await self._speech_service._generate_audio_bytes(req)
                        await ready_q.put(("bytes", idx, text, None, audio))
                except Exception as exc:
                    await ready_q.put(("error", idx, text, None, exc))

        async def _sender():
            total = 0
            while True:
                item = await ready_q.get()
                if item is None:
                    await websocket.send_json({"type": "session.done", "total_sentences": total})
                    return
                kind, idx, text, extra, data = item
                if kind == "error":
                    logger.error("Generation failed for sentence %d: %s", idx, data)
                    await self._send_error(websocket, f"Generation failed for sentence {idx}: {data}")
                    total += 1
                    continue
                hdr = {"type": "audio.start", "sentence_index": idx,
                       "sentence_text": text, "format": response_format}
                if config.stream_audio and response_format == "pcm":
                    hdr["sample_rate"] = _PCM_SAMPLE_RATE
                await websocket.send_json(hdr)
                if config.stream_audio and response_format == "pcm" and idx == 0 and _LEADIN_SILENCE_BYTES:
                    await websocket.send_bytes(_LEADIN_SILENCE_BYTES)
                sent_bytes = 0
                failed = False
                try:
                    if kind == "stream":
                        async with aclosing(
                            self._speech_service._generate_pcm_chunks(data, extra)
                        ) as stream:
                            async for chunk in stream:
                                sent_bytes += len(chunk)
                                await websocket.send_bytes(chunk)
                        pending_rids.discard(extra)
                    else:
                        sent_bytes = len(data)
                        await websocket.send_bytes(data)
                    if response_format == "pcm" and _TRAILING_SILENCE_BYTES:
                        await websocket.send_bytes(_TRAILING_SILENCE_BYTES)
                except WebSocketDisconnect:
                    raise
                except Exception as exc:
                    failed = True
                    logger.error("Send failed for sentence %d: %s", idx, exc)
                    await self._send_error(websocket, f"Send failed for sentence {idx}: {exc}")
                finally:
                    try:
                        await websocket.send_json({
                            "type": "audio.done", "sentence_index": idx,
                            "total_bytes": sent_bytes, "error": failed,
                        })
                    except Exception:
                        pass
                total += 1

        synth_task = asyncio.create_task(_synthesizer())
        send_task = asyncio.create_task(_sender())
        sentence_count = 0
        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: no message received")
                    synth_task.cancel()
                    send_task.cancel()
                    await asyncio.gather(synth_task, send_task, return_exceptions=True)
                    await _abort_all()
                    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
                mtype = msg.get("type")
                if mtype == "input.text":
                    t = msg.get("text", "")
                    if not isinstance(t, str):
                        await self._send_error(websocket, "input.text requires a string value")
                        continue
                    for s in splitter.add_text(t):
                        await sentence_q.put((sentence_count, s))
                        sentence_count += 1
                elif mtype == "input.done":
                    rem = splitter.flush()
                    if rem:
                        await sentence_q.put((sentence_count, rem))
                        sentence_count += 1
                    await sentence_q.put(None)
                    break
                else:
                    await self._send_error(websocket, f"Unknown message type: {mtype}")
            await asyncio.gather(synth_task, send_task)
        except BaseException:
            synth_task.cancel()
            send_task.cancel()
            await asyncio.gather(synth_task, send_task, return_exceptions=True)
            await _abort_all()
            raise
    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:
            pass


OmniStreamingSpeechHandler.handle_session = _pipelined_handle_session
'''


def patch_pipeline_session() -> None:
    """Append pipelined handle_session to serving_speech_stream.py.

    Synthesis of sentence N+1 is submitted to vllm as soon as N's request
    is prepared, so vllm can start N+1 the moment N's GPU work finishes —
    overlapping with the network transfer of N's audio to the client.
    """
    text = STREAM.read_text()
    sentinel = 'OmniStreamingSpeechHandler.handle_session = _pipelined_handle_session'
    if sentinel in text:
        print('[runtime-patch] pipeline session handler already present', flush=True)
        return
    STREAM.write_text(text + _PIPELINE_CODE)
    print('[runtime-patch] pipeline session handler injected', flush=True)


def main() -> None:
    patch_seed_field()

    text = STREAM.read_text()
    original = text

    text = replace_once(
        text,
        '            if config.model and hasattr(self._speech_service, "_check_model"):\n',
        '            logger.info(\n'
        '                "Streaming speech config: model=%s voice=%s task_type=%s language=%s seed=%s split=%s stream_audio=%s instructions=%r speaker_embedding_dim=%s",\n'
        '                config.model,\n'
        '                config.voice,\n'
        '                config.task_type,\n'
        '                config.language,\n'
        '                config.seed,\n'
        '                config.split_granularity,\n'
        '                config.stream_audio,\n'
        '                config.instructions,\n'
        '                len(config.speaker_embedding) if config.speaker_embedding is not None else None,\n'
        '            )\n\n'
        '            if config.model and hasattr(self._speech_service, "_check_model"):\n',
        'config logging',
    )

    text = replace_once(
        text,
        '            stream=config.stream_audio,\n        )\n',
        '            stream=config.stream_audio,\n            seed=config.seed,\n        )\n'
        '        logger.info(\n'
        '            "Streaming speech segment: index=%s chars=%s voice=%s task_type=%s language=%s seed=%s speaker_embedding_dim=%s",\n'
        '            sentence_index,\n'
        '            len(sentence_text),\n'
        '            config.voice,\n'
        '            config.task_type,\n'
        '            config.language,\n'
        '            config.seed,\n'
        '            len(config.speaker_embedding) if config.speaker_embedding is not None else None,\n'
        '        )\n',
        'seed propagation and segment logging',
    )

    # --- Silence padding (PCM streaming only) ---
    # Lead-in silence at session start gives the client time to set up smooth
    # playback before real audio (mitigates start-of-stream stutter). Trailing
    # silence avoids the abrupt cut-off at the end. Both in ms via env vars.
    text = replace_once(
        text,
        '_PCM_SAMPLE_RATE = 24000\n',
        '_PCM_SAMPLE_RATE = 24000\n'
        'import os as _os\n'
        '_LEADIN_SILENCE_MS = int(_os.environ.get("QWEN3_TTS_LEADIN_SILENCE_MS", "250"))\n'
        '_TRAILING_SILENCE_MS = int(_os.environ.get("QWEN3_TTS_TRAILING_SILENCE_MS", "300"))\n'
        '_LEADIN_SILENCE_BYTES = b"\\x00" * (int(_PCM_SAMPLE_RATE * _LEADIN_SILENCE_MS / 1000) * 2)\n'
        '_TRAILING_SILENCE_BYTES = b"\\x00" * (int(_PCM_SAMPLE_RATE * _TRAILING_SILENCE_MS / 1000) * 2)\n',
        'silence padding constants',
    )

    text = replace_once(
        text,
        '        await websocket.send_json(start_payload)\n',
        '        await websocket.send_json(start_payload)\n'
        '        if config.stream_audio and response_format == "pcm" and sentence_index == 0 and _LEADIN_SILENCE_BYTES:\n'
        '            await websocket.send_bytes(_LEADIN_SILENCE_BYTES)\n',
        'lead-in silence',
    )

    text = replace_once(
        text,
        '                    # Send session.done\n'
        '                    await websocket.send_json(\n',
        '                    if config.stream_audio and config.response_format == "pcm" and _TRAILING_SILENCE_BYTES:\n'
        '                        await websocket.send_bytes(_TRAILING_SILENCE_BYTES)\n\n'
        '                    # Send session.done\n'
        '                    await websocket.send_json(\n',
        'trailing silence',
    )

    if text != original:
        STREAM.write_text(text)
        print('[runtime-patch] WebSocket seed propagation/logging + silence padding applied', flush=True)
    else:
        print('[runtime-patch] WebSocket seed propagation/logging already present', flush=True)

    patch_pipeline_session()


if __name__ == '__main__':
    main()