[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>
2026-06-22 04:51:45 +00:00
1 parent a60242092a
commit ad07f98b73
@@ -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))
@@ -0,0 +1,165 @@
+import numpy as np
+import pytest
+
+from vllm_omni.utils import forced_aligner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_decode_timestamps_maps_boundary_bins_to_words():
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0
+    logits[1, 2] = 1.0
+    logits[2, 2] = 1.0
+    logits[3, 4] = 1.0
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == [
+        forced_aligner.WordTimestamp("hello", 0, 400),
+        forced_aligner.WordTimestamp("world", 400, 800),
+    ]
+
+
+def test_decode_timestamps_repairs_non_monotonic_bins():
+    # Word 1's end bin (1) dips below its start bin (2); fix_timestamp should
+    # snap it back so the frame never carries end_ms < start_ms.
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0  # word0 start -> bin 0
+    logits[1, 2] = 1.0  # word0 end   -> bin 2
+    logits[2, 1] = 1.0  # word1 start -> bin 1 (out of order)
+    logits[3, 3] = 1.0  # word1 end   -> bin 3
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        timestamp_segment_time_ms=200,
+        audio_duration_ms=1000,
+    )
+
+    for ts in timestamps:
+        assert ts.end_ms >= ts.start_ms
+    # bins [0, 2, 1, 3] -> repaired [0, 2, 2, 3] -> ms x200
+    assert timestamps == [
+        forced_aligner.WordTimestamp("hello", 0, 400),
+        forced_aligner.WordTimestamp("world", 400, 600),
+    ]
+
+
+def test_decode_timestamps_clamps_to_audio_duration_and_stays_ordered():
+    # Bins run past the 1000 ms audio; the decoder must cap each bound at the
+    # audio length and keep words non-overlapping (start <= end, monotonic).
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0  # word0 start -> bin 0
+    logits[1, 3] = 1.0  # word0 end   -> bin 3 (1200 ms, past audio)
+    logits[2, 3] = 1.0  # word1 start -> bin 3 (1200 ms, past audio)
+    logits[3, 4] = 1.0  # word1 end   -> bin 4 (1600 ms, past audio)
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["a", "b"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        timestamp_segment_time_ms=400,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == [
+        forced_aligner.WordTimestamp("a", 0, 1000),
+        forced_aligner.WordTimestamp("b", 1000, 1000),
+    ]
+
+
+def test_decode_timestamps_rejects_marker_count_mismatch():
+    logits = np.zeros((2, 5), dtype=np.float32)
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1],
+        classify_num=5,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == []
+
+
+def test_build_config_from_yaml(tmp_path):
+    cfg = tmp_path / "forced_aligner.yaml"
+    cfg.write_text(
+        """
+forced_aligner:
+  model: Qwen/Qwen3-ForcedAligner-0.6B
+  gpu_memory_utilization: 0.42
+  dtype: float16
+  max_model_len: 2048
+  trust_remote_code: false
+""",
+        encoding="utf-8",
+    )
+    args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out == forced_aligner.ForcedAlignerConfig(
+        model="Qwen/Qwen3-ForcedAligner-0.6B",
+        runner="pooling",
+        architecture="Qwen3ASRForcedAlignerForTokenClassification",
+        pooling_task="token_classify",
+        gpu_memory_utilization=0.42,
+        dtype="float16",
+        max_model_len=2048,
+        trust_remote_code=False,
+    )
+
+
+def test_build_config_cli_model_overrides_yaml(tmp_path):
+    cfg = tmp_path / "forced_aligner.yaml"
+    cfg.write_text(
+        "forced_aligner:\n  model: old\n  gpu_memory_utilization: 0.2\n  dtype: float16\n",
+        encoding="utf-8",
+    )
+    args = type(
+        "Args",
+        (),
+        {
+            "forced_aligner": "new",
+            "forced_aligner_config": str(cfg),
+        },
+    )()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out is not None
+    # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype
+    # come from the user YAML (there is no longer a CLI flag for gpu mem).
+    assert out.model == "new"
+    assert out.gpu_memory_utilization == 0.2
+    assert out.dtype == "float16"
+    assert out.runner == "pooling"
+
+
+def test_build_config_from_cli_model_uses_default_yaml():
+    args = type(
+        "Args",
+        (),
+        {
+            "forced_aligner": "local-aligner",
+            "forced_aligner_config": None,
+        },
+    )()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out is not None
+    assert out.model == "local-aligner"
+    assert out.runner == "pooling"
+    assert out.architecture == "Qwen3ASRForcedAlignerForTokenClassification"
@@ -0,0 +1,65 @@
+import pytest
+
+from vllm_omni.utils import qwen3_force_align_processor as processor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_build_prompt_has_boundary_timestamp_markers():
+    prompt = processor.build_prompt(["hello", "world"])
+
+    assert prompt.count("<timestamp>") == 4
+    assert "hello<timestamp><timestamp>world" in prompt
+    assert prompt.endswith("<|im_start|>assistant\n")
+
+
+@pytest.mark.parametrize(
+    "text,expected",
+    [
+        # Punctuation inside a whitespace segment is stripped, not split on:
+        # the hand-rolled splitter used to break these into many fake words.
+        ("U.S.A", ["USA"]),
+        ("hello, world!", ["hello", "world"]),
+        ("don't stop", ["don't", "stop"]),
+        # CJK characters peel off as individual tokens, Latin runs stay whole.
+        ("你好world", ["你", "好", "world"]),
+        ("我爱 China", ["我", "爱", "China"]),
+        ("   spaced   out   ", ["spaced", "out"]),
+    ],
+)
+def test_tokenize_space_lang_matches_official_segmentation(text, expected):
+    assert processor._tokenize_space_lang(text) == expected
+
+
+def test_segment_words_falls_back_to_port_without_qwen_asr(monkeypatch):
+    # With qwen_asr unavailable, segmentation must use the built-in port and
+    # still produce the faithful result for the common (non-JP/KO) case.
+    monkeypatch.setattr(processor, "_get_official_processor", lambda: None)
+
+    assert processor.segment_words("U.S.A is here", "auto") == ["USA", "is", "here"]
+
+
+def test_fix_timestamp_repairs_dip_onto_monotonic_sequence():
+    # 200 dips below the preceding 400; it snaps back up to 400.
+    assert processor.fix_timestamp([0, 400, 200, 800]) == [0, 400, 400, 800]
+
+
+def test_fix_timestamp_passes_through_monotonic():
+    assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250]
+
+
+def test_resolve_timestamp_token_id_defaults_to_marker_token():
+    # Regression: the default must resolve the same <timestamp> marker that
+    # build_prompt inserts, not None (which would degrade every request).
+    seen = {}
+
+    class FakeTokenizer:
+        def convert_tokens_to_ids(self, token):
+            seen["token"] = token
+            return 151705
+
+    tid = processor.resolve_timestamp_token_id(FakeTokenizer())
+
+    assert tid == 151705
+    assert seen["token"] == processor.TIMESTAMP_TOKEN
+    assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"])