[Feat] Streaming TTS word-level timestamps via forced aligner (#4034)

Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -1,4 +1,6 @@
 import asyncio
+import base64
+from types import SimpleNamespace
 
 import pytest
 from fastapi import FastAPI, WebSocket
@@ -9,6 +11,7 @@ from starlette.websockets import WebSocketDisconnect
 from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
+from vllm_omni.utils.forced_aligner import WordTimestamp
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -25,10 +28,11 @@ def _build_test_app(
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
+        speech_service.forced_aligner_config = None
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
-                yield chunk
+                yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
         speech_service.engine_client = mocker.MagicMock()
@@ -82,6 +86,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
@@ -89,9 +94,9 @@ class TestStreamingSpeechWebSocket:
 
         speech_service._prepare_speech_generation = mock_prepare_speech_generation
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
-                yield chunk
+                yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
         app, _ = _build_test_app(speech_service)
@@ -130,6 +135,164 @@ class TestStreamingSpeechWebSocket:
         assert captured_requests[0].initial_codec_chunk_frames == 12
         assert speech_service._generate_audio_bytes.await_count == 0
 
+    def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
+        app, _ = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.config",
+                        "voice": "Vivian",
+                        "stream_audio": True,
+                        "response_format": "pcm",
+                        "word_timestamps": True,
+                    }
+                )
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+
+                error = ws.receive_json()
+                assert error["type"] == "error"
+                assert "without --forced-aligner" in error["message"]
+
+    def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture):
+        captured_requests = []
+        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
+        speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
+        speech_service.engine_client = mocker.MagicMock()
+        speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
+
+        async def mock_prepare_speech_generation(request):
+            captured_requests.append(request)
+            return "req-stream", object(), {}
+
+        speech_service._prepare_speech_generation = mock_prepare_speech_generation
+
+        first_chunk = b"\x01" * 1000
+        second_chunk = b"\x02" * 1000
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+            for chunk in (first_chunk, second_chunk):
+                yield (chunk, 1000) if include_sample_rate else chunk
+
+        speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
+        # Sentence-level: aligner runs once over the whole sentence audio.
+        mock_align = mocker.AsyncMock(
+            return_value=[
+                WordTimestamp("Hello", 0, 200),
+                WordTimestamp("world", 200, 900),
+            ]
+        )
+        mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
+        app, _ = _build_test_app(speech_service)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.config",
+                        "voice": "Vivian",
+                        "stream_audio": True,
+                        "response_format": "pcm",
+                        "word_timestamps": True,
+                    }
+                )
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+
+                start = ws.receive_json()
+                assert start["type"] == "audio.start"
+                assert start["word_timestamps"] is True
+
+                # Audio streams first; timestamps are null until the sentence
+                # is fully aligned.
+                chunk = ws.receive_json()
+                assert chunk["type"] == "audio.chunk"
+                assert chunk["chunk_id"] == 0
+                assert chunk["chunk_start_ms"] == 0
+                assert chunk["chunk_end_ms"] == 500
+                assert chunk["sample_rate"] == 1000
+                assert base64.b64decode(chunk["audio_b64"]) == first_chunk
+                assert chunk["timestamps"] is None
+
+                chunk = ws.receive_json()
+                assert chunk["type"] == "audio.chunk"
+                assert chunk["chunk_id"] == 1
+                assert chunk["chunk_start_ms"] == 500
+                assert chunk["chunk_end_ms"] == 1000
+                assert chunk["sample_rate"] == 1000
+                assert base64.b64decode(chunk["audio_b64"]) == second_chunk
+                assert chunk["timestamps"] is None
+
+                # Final frame: empty audio carrying the whole-sentence timestamps.
+                chunk = ws.receive_json()
+                assert chunk["type"] == "audio.chunk"
+                assert chunk["chunk_id"] == 2
+                assert chunk["chunk_start_ms"] == 0
+                assert chunk["chunk_end_ms"] == 1000
+                assert chunk["sample_rate"] == 1000
+                assert base64.b64decode(chunk["audio_b64"]) == b""
+                assert chunk["timestamps"] == [
+                    {"word": "Hello", "start_ms": 0, "end_ms": 200},
+                    {"word": "world", "start_ms": 200, "end_ms": 900},
+                ]
+
+                done = ws.receive_json()
+                assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False}
+
+        assert captured_requests[0].word_timestamps is True
+        assert mock_align.await_count == 1
+
+    def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture):
+        # The streaming layer forwards the aligner's (already monotonic,
+        # non-overlapping) words as JSON dicts in the trailing frame.
+        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
+        speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
+        speech_service.engine_client = mocker.MagicMock()
+        speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
+        speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+            chunk = b"\x01" * 1000
+            yield (chunk, 1000) if include_sample_rate else chunk
+
+        speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
+        mock_align = mocker.AsyncMock(
+            return_value=[
+                WordTimestamp("Hello", 0, 1000),
+                WordTimestamp("world", 1000, 1200),
+            ]
+        )
+        mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
+        app, _ = _build_test_app(speech_service)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.config",
+                        "voice": "Vivian",
+                        "stream_audio": True,
+                        "response_format": "pcm",
+                        "word_timestamps": True,
+                    }
+                )
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+
+                assert ws.receive_json()["type"] == "audio.start"
+                # Real-time audio chunk (timestamps null), then the timestamp frame.
+                assert ws.receive_json()["timestamps"] is None
+                final = ws.receive_json()
+                assert final["type"] == "audio.chunk"
+                timestamps = final["timestamps"]
+                assert timestamps == [
+                    {"word": "Hello", "start_ms": 0, "end_ms": 1000},
+                    {"word": "world", "start_ms": 1000, "end_ms": 1200},
+                ]
+                for ts in timestamps:
+                    assert ts["end_ms"] >= ts["start_ms"]
+
     def test_flush_on_input_done(self, mocker: MockerFixture):
         app, _ = _build_test_app(mocker=mocker)
 
@@ -249,6 +412,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_pcm_chunks = mocker.AsyncMock()
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
         app, _ = _build_test_app(speech_service)
 
         with TestClient(app) as client:
@@ -269,8 +433,9 @@ class TestStreamingSpeechWebSocket:
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {}))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -361,8 +526,9 @@ class TestStreamingSpeechWebSocket:
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {}))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
@@ -385,7 +551,9 @@ class TestStreamingSpeechWebSocket:
         config.ref_audio = None
         config.ref_text = None
         config.x_vector_only_mode = None
+        config.speaker_embedding = None
         config.stream_audio = True
+        config.word_timestamps = False
 
         with pytest.raises(WebSocketDisconnect):
             asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0))