[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>

examples/offline_inference/text_to_speech/README.md

@@ -369,6 +369,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
 ```
 Streaming requires `async_chunk: true` in the stage config.
 
+### Word Timestamps
+Generate a WAV offline and a JSON sidecar with word-level timestamps from
+`Qwen/Qwen3-ForcedAligner-0.6B`:
+```bash
+python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
+    --text "Hello world." \
+    --output-dir /tmp/qwen3_tts_timestamps
+```
+The script writes `qwen3_tts_word_timestamps.wav` and
+`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
+set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
+
 ### Batched decoding
 The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
 ```bash

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -0,0 +1,172 @@
+"""Offline Qwen3-TTS example with word-level timestamps.
+
+This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with
+the same shared forced aligner used by the streaming server path, then writes
+the WAV and JSON sidecar:
+
+    python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
+        --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+        --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
+
+On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0``
+to avoid FlashInfer sampler JIT during warmup.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+# The in-process offline aligner consumes token_classify outputs. Keeping V1
+# multiprocessing off avoids msgspec serialization issues for token-wise logits.
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+import numpy as np
+import soundfile as sf
+import torch
+from end2end import _estimate_prompt_len
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
+
+
+def _default_stage_config() -> str:
+    repo_root = Path(__file__).resolve().parents[4]
+    return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
+
+
+def _build_custom_voice_input(args: Any) -> dict[str, Any]:
+    additional_information = {
+        "task_type": ["CustomVoice"],
+        "text": [args.text],
+        "language": [args.language],
+        "speaker": [args.speaker],
+        "instruct": [args.instructions],
+        "max_new_tokens": [args.max_new_tokens],
+    }
+    return {
+        "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model),
+        "additional_information": additional_information,
+    }
+
+
+def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]:
+    audio_data = mm["audio"]
+    sr_raw = mm["sr"]
+    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+    sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val)
+    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
+    return audio_tensor.float().detach().cpu().flatten(), sample_rate
+
+
+def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes:
+    samples = audio.numpy()
+    pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2")
+    return pcm.tobytes()
+
+
+async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None:
+    aligner_config = build_forced_aligner_config(args)
+    if aligner_config is None:
+        raise ValueError("--forced-aligner or --forced-aligner-config is required")
+    timestamps = await align(
+        audio=_float_audio_to_pcm16_bytes(audio),
+        text=args.text,
+        sample_rate=sample_rate,
+        config=aligner_config,
+        language=args.language,
+    )
+    if timestamps is None:
+        return None
+    return [
+        {
+            "word": item.word,
+            "start_ms": item.start_ms,
+            "end_ms": item.end_ms,
+        }
+        for item in timestamps
+    ]
+
+
+def main(args: Any) -> None:
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    omni_kwargs = vars(args).copy()
+    for key in (
+        "forced_aligner",
+        "text",
+        "language",
+        "speaker",
+        "instructions",
+        "max_new_tokens",
+        "output_dir",
+    ):
+        omni_kwargs.pop(key, None)
+    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
+    omni_kwargs["log_stats"] = args.log_stats
+    omni = Omni(**omni_kwargs)
+
+    prompt = _build_custom_voice_input(args)
+    final_output = None
+    for stage_outputs in omni.generate([prompt]):
+        final_output = stage_outputs.request_output
+    if final_output is None:
+        raise RuntimeError("Qwen3-TTS did not produce an output.")
+
+    mm = final_output.outputs[0].multimodal_output
+    audio, sample_rate = _audio_tensor_and_sample_rate(mm)
+
+    wav_path = output_dir / "qwen3_tts_word_timestamps.wav"
+    # Align the in-memory tensor, not the saved WAV. The WAV is only a demo
+    # artifact so users can listen to the same audio referenced by the sidecar.
+    timestamps = asyncio.run(_run_alignment(args, audio, sample_rate))
+    sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV")
+
+    sidecar = {
+        "text": args.text,
+        "sample_rate": sample_rate,
+        "audio_path": str(wav_path),
+        "timestamps": timestamps,
+    }
+    json_path = output_dir / "qwen3_tts_word_timestamps.json"
+    json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+    print(f"Saved audio: {wav_path}")
+    print(f"Saved timestamps: {json_path}")
+    print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2))
+
+
+def parse_args() -> Any:
+    parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example")
+    parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name")
+    parser.add_argument(
+        "--forced-aligner",
+        default=None,
+        help="Qwen3 forced aligner model path/name",
+    )
+    parser.add_argument(
+        "--forced-aligner-config",
+        default=None,
+        help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
+    )
+    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
+    parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
+    parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
+    parser.add_argument("--instructions", default="", help="Optional speaking style instruction")
+    parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
+    parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
+    parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
+    nullify_stage_engine_defaults(parser)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/online_serving/text_to_speech/README.md

@@ -504,6 +504,35 @@ python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I
 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
 ```
 
+To receive word-level timestamps, launch the server with a forced aligner:
+```bash
+vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --omni \
+    --deploy-config vllm_omni/deploy/qwen3_tts.yaml \
+    --trust-remote-code \
+    --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
+```
+Then request PCM JSON sidecar chunks:
+```bash
+python qwen3_tts/streaming_speech_client.py \
+    --text "Hello world. How are you?" \
+    --stream-audio \
+    --response-format pcm \
+    --word-timestamps
+```
+The client writes one PCM file per sentence and a matching
+`sentence_XXX_timestamps.json` sidecar.
+
+To *see* the alignment instead of reading a JSON sidecar, run the
+word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
+```bash
+python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091
+```
+Each sentence's audio plays in an `<audio>` element while its text is rendered
+as inline word spans; the current word highlights as `audio.currentTime`
+crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and
+reports the last-spoken word, useful for the voice-agent barge-in case.
+
 ### Gradio demos
 ```bash
 ./qwen3_tts/run_gradio_demo.sh                              # CustomVoice (default)

examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py

@@ -18,6 +18,21 @@ Usage:
         --text "Hello world. How are you? I am fine." \
         --simulate-stt --stt-delay 0.1
 
+    # Receive JSON sidecar chunks with word-level timestamps
+    python streaming_speech_client.py \
+        --text "Hello world. How are you?" \
+        --stream-audio \
+        --response-format pcm \
+        --word-timestamps
+
+    # Also save every received audio.chunk frame as individual PCM/WAV files
+    python streaming_speech_client.py \
+        --text "Hello world. How are you?" \
+        --stream-audio \
+        --response-format pcm \
+        --word-timestamps \
+        --save-chunks
+
     # VoiceDesign task
     python streaming_speech_client.py \
         --text "Today is a great day. The weather is nice." \
@@ -37,8 +52,10 @@ Requirements:
 
 import argparse
 import asyncio
+import base64
 import json
 import os
+import wave
 
 try:
     import websockets
@@ -47,6 +64,37 @@ except ImportError:
     raise SystemExit(1)
 
 
+def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
+    audio_bytes = b"".join(chunks)
+    if response_format == "pcm":
+        raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
+        with open(raw_path, "wb") as f:
+            f.write(audio_bytes)
+
+        wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
+        save_pcm_wav(wav_path, audio_bytes)
+        return wav_path
+
+    filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
+    with open(filename, "wb") as f:
+        f.write(audio_bytes)
+    return filename
+
+
+def save_pcm_wav(path: str, audio_bytes: bytes, sample_rate: int = 24000) -> None:
+    with wave.open(path, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(audio_bytes)
+
+
+def timestamp_words_text(timestamps: list[dict] | None) -> str:
+    if not timestamps:
+        return ""
+    return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word"))
+
+
 async def stream_tts(
     url: str,
     text: str,
@@ -54,6 +102,7 @@ async def stream_tts(
     output_dir: str,
     simulate_stt: bool = False,
     stt_delay: float = 0.1,
+    save_chunks: bool = False,
 ) -> None:
     """Connect to the streaming TTS endpoint and process audio responses."""
     os.makedirs(output_dir, exist_ok=True)
@@ -99,8 +148,14 @@ async def stream_tts(
         sender_task = asyncio.create_task(send_text())
 
         response_format = config.get("response_format", "wav")
+        word_timestamps = bool(config.get("word_timestamps", False))
         current_sentence_index = 0
+        current_sentence_text = ""
         current_chunks: list[bytes] = []
+        current_timestamps: list[dict] = []
+        # Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
+        # the sentence is a failure only if no frame ever carried timestamps.
+        alignment_frame_seen = False
 
         try:
             while True:
@@ -116,22 +171,94 @@ async def stream_tts(
 
                     if msg_type == "audio.start":
                         current_sentence_index = msg["sentence_index"]
+                        current_sentence_text = msg.get("sentence_text", "")
                         current_chunks = []
-                        print(f"  [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}")
+                        current_timestamps = []
+                        alignment_frame_seen = False
+                        print(f"  [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}")
+                    elif msg_type == "audio.chunk":
+                        audio = base64.b64decode(msg["audio_b64"])
+                        current_chunks.append(audio)
+                        chunk_timestamps = msg.get("timestamps")
+                        # Sentence-level alignment sends a trailing timestamp-only
+                        # frame with empty audio; don't write an empty chunk file.
+                        if save_chunks and audio:
+                            chunk_base = os.path.join(
+                                output_dir,
+                                f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}",
+                            )
+                            chunk_pcm = f"{chunk_base}.pcm"
+                            chunk_wav = f"{chunk_base}.wav"
+                            chunk_json = f"{chunk_base}_timestamps.json"
+                            with open(chunk_pcm, "wb") as f:
+                                f.write(audio)
+                            save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
+                            with open(chunk_json, "w", encoding="utf-8") as f:
+                                json.dump(
+                                    {
+                                        "sentence_index": msg["sentence_index"],
+                                        "chunk_id": msg["chunk_id"],
+                                        "chunk_start_ms": msg.get("chunk_start_ms"),
+                                        "chunk_end_ms": msg.get("chunk_end_ms"),
+                                        "sentence_text": current_sentence_text,
+                                        "covered_text": timestamp_words_text(chunk_timestamps),
+                                        "timestamps": chunk_timestamps,
+                                    },
+                                    f,
+                                    ensure_ascii=False,
+                                    indent=2,
+                                )
+                        if chunk_timestamps is None:
+                            print(
+                                f"  [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
+                                f"{len(audio)} bytes, timestamps unavailable"
+                            )
+                        else:
+                            alignment_frame_seen = True
+                            current_timestamps.extend(chunk_timestamps)
+                            print(
+                                f"  [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
+                                f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
+                            )
+                        if save_chunks and audio:
+                            print(f"    saved chunk -> {chunk_wav}")
                     elif msg_type == "audio.done":
-                        filename = os.path.join(
+                        filename = save_audio_file(
                             output_dir,
-                            f"sentence_{msg['sentence_index']:03d}.{response_format}",
+                            msg["sentence_index"],
+                            response_format,
+                            current_chunks,
                         )
-                        with open(filename, "wb") as f:
-                            f.write(b"".join(current_chunks))
+                        if word_timestamps:
+                            # `null` on failure, else the accumulated list (`[]` for silence).
+                            sentence_timestamps = current_timestamps if alignment_frame_seen else None
+                            ts_filename = os.path.join(
+                                output_dir,
+                                f"sentence_{msg['sentence_index']:03d}_timestamps.json",
+                            )
+                            with open(ts_filename, "w", encoding="utf-8") as f:
+                                json.dump(
+                                    {
+                                        "sentence_index": msg["sentence_index"],
+                                        "sentence_text": current_sentence_text,
+                                        "covered_text": timestamp_words_text(sentence_timestamps),
+                                        "timestamps": sentence_timestamps,
+                                    },
+                                    f,
+                                    ensure_ascii=False,
+                                    indent=2,
+                                )
                         print(
                             f"  [sentence {msg['sentence_index']}] Done"
                             f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
                             f" error={msg.get('error', False)}"
                             f" -> {filename}"
                         )
+                        if word_timestamps:
+                            print(f"  [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}")
                         current_chunks = []
+                        current_timestamps = []
+                        alignment_frame_seen = False
                     elif msg_type == "session.done":
                         print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
                         break
@@ -189,6 +316,16 @@ def main():
         action="store_true",
         help="Receive one or more PCM chunks per sentence (requires --response-format pcm)",
     )
+    parser.add_argument(
+        "--word-timestamps",
+        action="store_true",
+        help="Receive JSON sidecar chunks with word-level timestamps (requires --stream-audio --response-format pcm)",
+    )
+    parser.add_argument(
+        "--save-chunks",
+        action="store_true",
+        help="Save each audio.chunk frame as sentence_XXX_chunk_YYY.{pcm,wav,json}",
+    )
     parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
     parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
 
@@ -238,6 +375,10 @@ def main():
             config[key] = val
     if args.stream_audio:
         config["stream_audio"] = True
+    if args.word_timestamps:
+        config["word_timestamps"] = True
+        config["stream_audio"] = True
+        config["response_format"] = "pcm"
     if args.x_vector_only_mode:
         config["x_vector_only_mode"] = True
 
@@ -249,6 +390,7 @@ def main():
             output_dir=args.output_dir,
             simulate_stt=args.simulate_stt,
             stt_delay=args.stt_delay,
+            save_chunks=args.save_chunks,
         )
     )
 

examples/online_serving/text_to_speech/qwen3_tts/word_timestamps_demo.py

@@ -0,0 +1,457 @@
+"""Gradio demo for streaming TTS word-level timestamps.
+
+Connects to the WebSocket endpoint ``/v1/audio/speech/stream`` with
+``word_timestamps: true`` and visualizes the alignment: each sentence's
+audio plays in a native ``<audio>`` element while its text is rendered as
+inline word spans, the current word highlighting as ``audio.currentTime``
+crosses each ``start_ms``.
+
+A "Stop (barge-in)" button cuts playback and reports the last-spoken word,
+demonstrating the voice-agent barge-in use case.
+
+Timestamps are sentence-level: audio streams in real time, then the forced
+aligner runs once per sentence and the complete word timestamps arrive in a
+trailing frame. The demo therefore plays one sentence block at a time.
+
+Usage:
+    # Launch the server with a forced aligner first, e.g.:
+    #   vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+    #       --deploy-config vllm_omni/deploy/qwen3_tts.yaml --trust-remote-code \
+    #       --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
+    # Then:
+    python word_timestamps_demo.py --api-base http://localhost:8000
+"""
+
+import argparse
+import asyncio
+import json
+import logging
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+try:
+    import websockets
+except ImportError:
+    raise ImportError("websockets is required to run this demo. Install it with: pip install websockets") from None
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from tts_common import (
+    PCM_SAMPLE_RATE,
+    SUPPORTED_LANGUAGES,
+    TASK_TYPES,
+    add_common_args,
+    fetch_voices,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _upstream_ws_url(api_base: str) -> str:
+    """Map an http(s) API base to its WebSocket streaming endpoint."""
+    base = api_base.rstrip("/")
+    if base.startswith("https://"):
+        base = "wss://" + base[len("https://") :]
+    elif base.startswith("http://"):
+        base = "ws://" + base[len("http://") :]
+    else:
+        base = "ws://" + base
+    return base + "/v1/audio/speech/stream"
+
+
+# ── Front-end ────────────────────────────────────────────────────────
+PLAYER_HTML = """
+<div id="wt-root">
+  <div style="display:flex; align-items:center; gap:10px;">
+    <div id="wt-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
+    <span id="wt-status" style="font-weight:600;font-size:1.05em;">Ready</span>
+    <button id="wt-stop" onclick="window.wtStop()"
+      style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
+             background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em;">Stop (barge-in)</button>
+  </div>
+  <div id="wt-barge" style="display:none; margin-top:10px; padding:8px 12px; border-radius:6px;
+       background:#fff5f5; border:1px solid #f3c2c2; color:#b02a2a; font-size:0.9em;"></div>
+  <div id="wt-sentences" style="margin-top:12px; display:flex; flex-direction:column; gap:14px;"></div>
+</div>
+"""
+
+
+def _build_head_js(sample_rate: int) -> str:
+    """Browser-side player: WebSocket -> per-sentence <audio> + word highlight."""
+    return f"""
+<script>
+const WT_SR = {sample_rate};
+let wtWS = null, wtQueue = [], wtPlaying = false, wtCur = null, wtLastWord = "";
+
+function wtSet(text, color) {{
+    const s = document.getElementById('wt-status');
+    const d = document.getElementById('wt-dot');
+    if (s) s.textContent = text;
+    if (d) d.style.background = color || '#ccc';
+}}
+
+function wtStopBtn(show) {{
+    const b = document.getElementById('wt-stop');
+    if (b) b.style.display = show ? 'inline-block' : 'none';
+}}
+
+// Build a WAV Blob (16-bit mono PCM) from concatenated Int16 samples.
+function wtPcmToWav(int16, sr) {{
+    const n = int16.length;
+    const buf = new ArrayBuffer(44 + n * 2);
+    const v = new DataView(buf);
+    const w = (off, s) => {{ for (let i = 0; i < s.length; i++) v.setUint8(off + i, s.charCodeAt(i)); }};
+    w(0, 'RIFF'); v.setUint32(4, 36 + n * 2, true); w(8, 'WAVE');
+    w(12, 'fmt '); v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
+    v.setUint32(24, sr, true); v.setUint32(28, sr * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
+    w(36, 'data'); v.setUint32(40, n * 2, true);
+    let off = 44;
+    for (let i = 0; i < n; i++) {{ v.setInt16(off, int16[i], true); off += 2; }}
+    return new Blob([buf], {{ type: 'audio/wav' }});
+}}
+
+function wtConcat(chunks) {{
+    let total = 0;
+    for (const c of chunks) total += c.length;
+    const out = new Int16Array(total);
+    let off = 0;
+    for (const c of chunks) {{ out.set(c, off); off += c.length; }}
+    return out;
+}}
+
+// Render one finished sentence: <audio> + inline word spans.
+function wtRenderSentence(sentence) {{
+    const wrap = document.getElementById('wt-sentences');
+    if (!wrap) return;
+    const block = document.createElement('div');
+    block.style.cssText = 'border:1px solid #e3e8ef;border-radius:8px;padding:10px 12px;';
+
+    const head = document.createElement('div');
+    head.style.cssText = 'font-size:0.75em;color:#888;margin-bottom:6px;';
+    head.textContent = 'Sentence ' + sentence.index +
+        (sentence.timestamps ? ' (' + sentence.timestamps.length + ' words)' : ' (no timestamps)');
+    block.appendChild(head);
+
+    const textEl = document.createElement('div');
+    textEl.style.cssText = 'font-size:1.15em;line-height:1.8;margin-bottom:8px;';
+    if (sentence.timestamps && sentence.timestamps.length) {{
+        sentence.spans = [];
+        for (const ts of sentence.timestamps) {{
+            const span = document.createElement('span');
+            span.textContent = ts.word;
+            span.dataset.start = ts.start_ms;
+            span.dataset.end = ts.end_ms;
+            span.style.cssText = 'padding:1px 2px;border-radius:3px;transition:background 0.1s;';
+            textEl.appendChild(span);
+            // Keep spacing for spaced languages; CJK chars sit flush.
+            textEl.appendChild(document.createTextNode(' '));
+            sentence.spans.push(span);
+        }}
+    }} else {{
+        textEl.textContent = sentence.text || '(timestamps unavailable)';
+    }}
+    block.appendChild(textEl);
+
+    const audio = document.createElement('audio');
+    audio.controls = true;
+    audio.style.cssText = 'width:100%;';
+    audio.src = URL.createObjectURL(wtPcmToWav(wtConcat(sentence.chunks), WT_SR));
+    block.appendChild(audio);
+    sentence.audio = audio;
+
+    // Highlight the current word. Driven by requestAnimationFrame (~60fps)
+    // rather than the <audio> 'timeupdate' event, which only fires ~4x/sec
+    // and is too coarse for short words. We light the latest word whose
+    // start_ms has passed, so there is always exactly one active word (no
+    // inter-word flicker) and wtLastWord stays well-defined for barge-in.
+    let raf = null;
+    function highlight() {{
+        if (!sentence.spans) return;
+        const t = audio.currentTime * 1000;
+        let idx = -1;
+        for (let i = 0; i < sentence.spans.length; i++) {{
+            if (t >= +sentence.spans[i].dataset.start) idx = i;
+            else break;
+        }}
+        sentence.spans.forEach((sp, i) => {{
+            const on = i === idx;
+            sp.style.background = on ? '#ffe680' : 'transparent';
+            sp.style.fontWeight = on ? '700' : '400';
+        }});
+        if (idx >= 0) wtLastWord = sentence.spans[idx].textContent;
+    }}
+    function loop() {{
+        highlight();
+        if (!audio.paused && !audio.ended) raf = requestAnimationFrame(loop);
+    }}
+    audio.addEventListener('play', () => {{
+        wtCur = audio;
+        wtSet('Playing sentence ' + sentence.index + '...', '#64dd17');
+        cancelAnimationFrame(raf);
+        raf = requestAnimationFrame(loop);
+    }});
+    audio.addEventListener('pause', () => cancelAnimationFrame(raf));
+    audio.addEventListener('seeked', highlight);
+    audio.addEventListener('ended', () => {{
+        cancelAnimationFrame(raf);
+        highlight();
+        wtPlaying = false;
+        wtPlayNext();
+    }});
+
+    wrap.appendChild(block);
+    wtQueue.push(audio);
+    wtPlayNext();
+}}
+
+function wtPlayNext() {{
+    if (wtPlaying) return;
+    const next = wtQueue.shift();
+    if (!next) {{
+        if (wtWS === null) {{ wtSet('Done', '#64dd17'); wtStopBtn(false); }}
+        return;
+    }}
+    wtPlaying = true;
+    next.play().catch(() => {{ wtPlaying = false; }});
+}}
+
+window.wtStop = function() {{
+    // Barge-in: cut playback, stop generation, report the last-spoken word.
+    if (wtCur) {{ try {{ wtCur.pause(); }} catch (e) {{}} }}
+    wtQueue = [];
+    wtPlaying = false;
+    if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} wtWS = null; }}
+    wtSet('Stopped (barge-in)', '#999');
+    wtStopBtn(false);
+    const b = document.getElementById('wt-barge');
+    if (b) {{
+        b.style.display = 'block';
+        b.textContent = wtLastWord
+            ? 'Barge-in at word: "' + wtLastWord + '"  (a voice agent would resume the conversation from here)'
+            : 'Barge-in before any word was spoken.';
+    }}
+}};
+
+window.wtGenerate = function(payload) {{
+    if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} }}
+    wtQueue = []; wtPlaying = false; wtCur = null; wtLastWord = "";
+    document.getElementById('wt-sentences').innerHTML = '';
+    const barge = document.getElementById('wt-barge');
+    if (barge) {{ barge.style.display = 'none'; barge.textContent = ''; }}
+    wtStopBtn(true);
+    wtSet('Connecting...', '#4A90D9');
+
+    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
+    const ws = new WebSocket(proto + '//' + location.host + '/proxy/stream');
+    wtWS = ws;
+    const sentences = {{}};
+
+    ws.onopen = () => {{
+        wtSet('Streaming...', '#4A90D9');
+        ws.send(JSON.stringify(Object.assign({{ type: 'session.config' }}, payload.config)));
+        ws.send(JSON.stringify({{ type: 'input.text', text: payload.text }}));
+        ws.send(JSON.stringify({{ type: 'input.done' }}));
+    }};
+
+    ws.onmessage = (ev) => {{
+        let msg;
+        try {{ msg = JSON.parse(ev.data); }} catch (e) {{ return; }}
+        if (msg.type === 'audio.start') {{
+            sentences[msg.sentence_index] = {{
+                index: msg.sentence_index, text: msg.sentence_text || '',
+                chunks: [], timestamps: null
+            }};
+        }} else if (msg.type === 'audio.chunk') {{
+            const s = sentences[msg.sentence_index];
+            if (!s) return;
+            if (msg.audio_b64) {{
+                const bin = atob(msg.audio_b64);
+                const bytes = new Uint8Array(bin.length);
+                for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
+                const usable = bytes.length - (bytes.length % 2);
+                if (usable > 0) s.chunks.push(new Int16Array(bytes.buffer, 0, usable / 2));
+            }}
+            if (msg.timestamps != null) s.timestamps = msg.timestamps;
+        }} else if (msg.type === 'audio.done') {{
+            const s = sentences[msg.sentence_index];
+            if (s) wtRenderSentence(s);
+        }} else if (msg.type === 'session.done') {{
+            wtWS = null;
+            try {{ ws.close(); }} catch (e) {{}}
+            wtPlayNext();
+        }} else if (msg.type === 'error') {{
+            wtSet('Error: ' + (msg.message || 'unknown'), '#EF5552');
+        }}
+    }};
+
+    ws.onerror = () => wtSet('WebSocket error', '#EF5552');
+    ws.onclose = () => {{ if (wtWS === ws) wtWS = null; }};
+}};
+</script>
+"""
+
+
+def create_app(api_base: str) -> FastAPI:
+    """FastAPI app with a same-origin WebSocket proxy + Gradio UI."""
+    fastapi_app = FastAPI()
+    upstream_url = _upstream_ws_url(api_base)
+
+    @fastapi_app.websocket("/proxy/stream")
+    async def proxy_stream(client_ws: WebSocket):
+        await client_ws.accept()
+        try:
+            async with websockets.connect(upstream_url, max_size=None) as upstream:
+
+                async def client_to_upstream():
+                    try:
+                        while True:
+                            await upstream.send(await client_ws.receive_text())
+                    except (WebSocketDisconnect, Exception):
+                        await upstream.close()
+
+                async def upstream_to_client():
+                    try:
+                        async for msg in upstream:
+                            if isinstance(msg, bytes):
+                                await client_ws.send_bytes(msg)
+                            else:
+                                await client_ws.send_text(msg)
+                    except Exception:
+                        pass
+                    finally:
+                        try:
+                            await client_ws.close()
+                        except Exception:
+                            pass
+
+                await asyncio.gather(client_to_upstream(), upstream_to_client())
+        except Exception as exc:
+            logger.exception("WS proxy error")
+            try:
+                await client_ws.send_text(json.dumps({"type": "error", "message": f"proxy: {exc}"}))
+                await client_ws.close()
+            except Exception:
+                pass
+
+    voices = fetch_voices(api_base)
+
+    with gr.Blocks(title="Qwen3-TTS Word Timestamps") as demo:
+        gr.HTML(f"""
+        <div style="margin-bottom:8px;">
+          <h1 style="margin:0; font-size:1.5em;">Qwen3-TTS Streaming Word Timestamps</h1>
+          <span style="font-size:0.85em; color:#666;">
+            Served by <a href="https://github.com/vllm-project/vllm-omni" target="_blank">vLLM-Omni</a>
+            &nbsp;&middot;&nbsp; <code>{api_base}</code> &nbsp;&middot;&nbsp;
+            launch the server with <code>--forced-aligner</code>.
+          </span>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    value="Hello world. How are you today? I am doing great.",
+                    lines=3,
+                )
+                with gr.Row():
+                    task_type = gr.Radio(choices=TASK_TYPES, value="CustomVoice", label="Task Type", scale=2)
+                    language = gr.Dropdown(choices=SUPPORTED_LANGUAGES, value="Auto", label="Language", scale=1)
+                voice = gr.Dropdown(
+                    choices=voices,
+                    value=voices[0] if voices else None,
+                    label="Speaker",
+                    allow_custom_value=True,
+                )
+                instructions = gr.Textbox(
+                    label="Instructions",
+                    placeholder="Optional style/emotion (required for VoiceDesign)",
+                    lines=1,
+                )
+                with gr.Column(visible=False) as ref_group:
+                    ref_audio_url = gr.Textbox(
+                        label="Reference Audio URL (Base task)",
+                        placeholder="https://example.com/reference.wav",
+                        lines=1,
+                    )
+                    ref_text = gr.Textbox(label="Reference Transcript", lines=1)
+                generate_btn = gr.Button("Generate + Align", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                gr.HTML(value=PLAYER_HTML)
+
+        hidden_payload = gr.Textbox(visible=False, elem_id="wt-payload")
+
+        def on_task_change(tt):
+            return gr.update(visible=(tt == "Base"))
+
+        task_type.change(fn=on_task_change, inputs=[task_type], outputs=[ref_group])
+
+        def on_generate(text, tt, voice_v, lang_v, instr_v, ref_url, ref_t):
+            if not text or not text.strip():
+                raise gr.Error("Please enter text to synthesize.")
+            config: dict = {
+                "task_type": tt,
+                "word_timestamps": True,
+                "stream_audio": True,
+                "response_format": "pcm",
+            }
+            if lang_v:
+                config["language"] = lang_v
+            if tt == "CustomVoice":
+                if voice_v:
+                    config["voice"] = voice_v
+                if instr_v and instr_v.strip():
+                    config["instructions"] = instr_v.strip()
+            elif tt == "VoiceDesign":
+                if not instr_v or not instr_v.strip():
+                    raise gr.Error("VoiceDesign requires voice style instructions.")
+                config["instructions"] = instr_v.strip()
+            elif tt == "Base":
+                if not ref_url or not ref_url.strip():
+                    raise gr.Error("Base (voice clone) needs a reference audio URL.")
+                config["ref_audio"] = ref_url.strip()
+                if ref_t and ref_t.strip():
+                    config["ref_text"] = ref_t.strip()
+            return json.dumps({"config": config, "text": text.strip()})
+
+        generate_btn.click(
+            fn=on_generate,
+            inputs=[text_input, task_type, voice, language, instructions, ref_audio_url, ref_text],
+            outputs=[hidden_payload],
+        ).then(
+            fn=lambda p: p,
+            inputs=[hidden_payload],
+            outputs=[hidden_payload],
+            js="(p) => { if (p && p.trim()) window.wtGenerate(JSON.parse(p)); return p; }",
+        )
+
+        demo.queue()
+
+    return gr.mount_gradio_app(
+        fastapi_app,
+        demo,
+        path="/",
+        head=_build_head_js(PCM_SAMPLE_RATE),
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gradio demo for streaming TTS word-level timestamps.")
+    add_common_args(parser)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    logging.basicConfig(level=logging.INFO)
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    app = create_app(args.api_base)
+
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()