[Feature] timestamps for TTS via a forced-aligner pooling (#4795)

Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: xurui <xurui1@modelbest.onaliyun.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -3,6 +3,7 @@ import base64
 from types import SimpleNamespace
 
 import pytest
+import torch
 from fastapi import FastAPI, WebSocket
 from pytest_mock import MockerFixture
 from starlette.testclient import TestClient
@@ -11,11 +12,24 @@ 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
+from vllm_omni.model_executor.stage_input_processors.forced_aligner import ALIGNER_WORDS_KEY
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _fake_aligner_res(pairs, words):
+    """Build a stand-in for the forced-aligner stage's pooling output.
+
+    Mirrors what rides the generator in production: ``res.outputs.data`` is a
+    ``[n_words, 2]`` int32 tensor of ``[start_ms, end_ms]`` and the word strings
+    travel in ``additional_information``. Decoded by ``extract_word_timestamps``.
+    """
+    return SimpleNamespace(
+        outputs=SimpleNamespace(data=torch.tensor(pairs, dtype=torch.int32), additional_information=None),
+        additional_information={ALIGNER_WORDS_KEY: list(words)},
+    )
+
+
 def _build_test_app(
     speech_service=None,
     *,
@@ -28,9 +42,9 @@ 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
+        speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -235,7 +249,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
+        speech_service.forced_aligner_enabled = False
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
@@ -311,13 +325,13 @@ class TestStreamingSpeechWebSocket:
                 assert error["type"] == "error"
                 assert "without --forced-aligner" in error["message"]
 
-    def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture):
+    def test_word_timestamps_emit_pipeline_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")
+        speech_service.forced_aligner_enabled = True
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
@@ -328,19 +342,15 @@ class TestStreamingSpeechWebSocket:
         first_chunk = b"\x01" * 1000
         second_chunk = b"\x02" * 1000
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        # The forced-aligner stage rides the same generator: its pooling output
+        # is surfaced via the ``collect`` channel once the audio has streamed.
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
             for chunk in (first_chunk, second_chunk):
                 yield (chunk, 1000) if include_sample_rate else chunk
+            if collect is not None:
+                collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"])
 
         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:
@@ -406,7 +416,6 @@ class TestStreamingSpeechWebSocket:
                 }
 
         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,
@@ -415,21 +424,16 @@ 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 = SimpleNamespace(model="aligner")
+        speech_service.forced_aligner_enabled = True
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
             chunk = b"\x01" * 1000
             yield (chunk, 1000) if include_sample_rate else chunk
+            if collect is not None:
+                collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"])
 
         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:
@@ -572,7 +576,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
+        speech_service.forced_aligner_enabled = False
         app, _ = _build_test_app(speech_service)
 
         with TestClient(app) as client:
@@ -602,7 +606,7 @@ 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
+        speech_service.forced_aligner_enabled = False
 
         async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"
@@ -696,7 +700,7 @@ 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
+        speech_service.forced_aligner_enabled = False
 
         async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"