[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/utils/test_forced_aligner.py

@@ -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"