perf: add TTS latency optimizations

Optimizations applied to voice_clone.py and tts.py:
- Auto-select bfloat16 (Ampere+) or float16 per GPU architecture
- FlashAttention-2 when available, SDPA fallback
- Streaming mode (non_streaming_mode=False)
- Reduced sampling: top_k=20, temperature=0.7, max_new_tokens=1024
- torch.set_float32_matmul_precision('high')
- torch.compile on decoder/predictor submodules (best-effort)
- Reference audio trimmed to 10s max for faster embedding extraction
- 0.6B CustomVoice model as fast non-cloned fallback (env: QWEN_TTS_MODE)
- Added librosa to requirements.txt
- FlashAttention-2 install instructions in requirements.txt
- 18-test suite for all optimization configs (test_latency_optimizations.py)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

requirements.txt

@@ -5,6 +5,7 @@ accelerate
 bitsandbytes
 soundfile
 numpy
+librosa
 
 # Supertonic TTS (fallback voice)
 supertonic>=1.3.1
@@ -13,3 +14,7 @@ huggingface-hub>=0.23.0
 
 # Qwen3-TTS (voice cloning)
 qwen-tts>=0.1.1
+
+# FlashAttention-2 — install separately on GPU machines:
+#   pip install flash-attn --no-build-isolation
+# Falls back to SDPA (PyTorch native) if not available.

test_modules/test_latency_optimizations.py

@@ -0,0 +1,189 @@
+"""
+Test latency optimizations in voice_clone.py and tts.py.
+
+Tests configuration logic, dtype selection, attention selection,
+audio trimming, and generation parameter values — without loading
+full models (no GPU required).
+"""
+import os
+import sys
+import tempfile
+import importlib
+
+import numpy as np
+import soundfile as sf
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+
+PASS = 0
+FAIL = 0
+
+
+def check(label, condition, detail=""):
+    global PASS, FAIL
+    if condition:
+        PASS += 1
+        print(f"  [OK] {label}")
+    else:
+        FAIL += 1
+        print(f"  [FAIL] {label} -- {detail}")
+
+
+def test_global_matmul_precision():
+    print("\n=== torch.set_float32_matmul_precision ===")
+    import voice_clone  # noqa: F401 — importing sets precision
+    # PyTorch doesn't expose a getter, but the call should not raise
+    check("Module imported without error (precision set)", True)
+
+
+def test_dtype_selection():
+    print("\n=== Dtype selection ===")
+    from voice_clone import _select_dtype
+
+    dtype = _select_dtype()
+    if torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability()
+        if cap[0] >= 8:
+            check("Ampere+ GPU -> bfloat16", dtype == torch.bfloat16,
+                  f"got {dtype}")
+        else:
+            check("Pre-Ampere GPU -> float16", dtype == torch.float16,
+                  f"got {dtype}")
+    else:
+        check("CPU -> float32", dtype == torch.float32, f"got {dtype}")
+
+
+def test_attn_selection():
+    print("\n=== Attention implementation selection ===")
+    from voice_clone import _select_attn_impl
+
+    impl = _select_attn_impl()
+    try:
+        import flash_attn  # noqa: F401
+        check("flash-attn installed -> flash_attention_2",
+              impl == "flash_attention_2", f"got {impl}")
+    except ImportError:
+        check("flash-attn not installed -> sdpa fallback",
+              impl == "sdpa", f"got {impl}")
+
+
+def test_generation_params():
+    print("\n=== Generation parameters ===")
+    from voice_clone import GENERATION_PARAMS
+
+    check("top_k reduced to 20",
+          GENERATION_PARAMS["top_k"] == 20,
+          f"got {GENERATION_PARAMS['top_k']}")
+    check("temperature reduced to 0.7",
+          GENERATION_PARAMS["temperature"] == 0.7,
+          f"got {GENERATION_PARAMS['temperature']}")
+    check("max_new_tokens capped at 1024",
+          GENERATION_PARAMS["max_new_tokens"] == 1024,
+          f"got {GENERATION_PARAMS['max_new_tokens']}")
+    check("subtalker_top_k reduced to 20",
+          GENERATION_PARAMS["subtalker_top_k"] == 20,
+          f"got {GENERATION_PARAMS['subtalker_top_k']}")
+    check("subtalker_temperature reduced to 0.7",
+          GENERATION_PARAMS["subtalker_temperature"] == 0.7,
+          f"got {GENERATION_PARAMS['subtalker_temperature']}")
+
+
+def test_ref_audio_trimming():
+    print("\n=== Reference audio trimming ===")
+    from voice_clone import _trim_reference_audio, REF_AUDIO_MAX_SEC
+
+    # Create a 20-second test WAV
+    sr = 24000
+    long_audio = np.random.randn(int(20 * sr)).astype(np.float32) * 0.1
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        sf.write(f.name, long_audio, sr)
+        long_path = f.name
+
+    # Create a 4-second test WAV (under limit)
+    short_audio = np.random.randn(int(4 * sr)).astype(np.float32) * 0.1
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        sf.write(f.name, short_audio, sr)
+        short_path = f.name
+
+    try:
+        # Long audio should be trimmed
+        trimmed = _trim_reference_audio(long_path)
+        check("Long audio (20s) is trimmed", trimmed != long_path)
+        if trimmed != long_path:
+            info = sf.info(trimmed)
+            check(f"Trimmed to <= {REF_AUDIO_MAX_SEC}s",
+                  info.duration <= REF_AUDIO_MAX_SEC + 0.1,
+                  f"got {info.duration:.1f}s")
+            os.unlink(trimmed)
+
+        # Short audio should NOT be trimmed
+        result = _trim_reference_audio(short_path)
+        check("Short audio (4s) is NOT trimmed", result == short_path)
+    finally:
+        os.unlink(long_path)
+        os.unlink(short_path)
+
+
+def test_model_mode_env():
+    print("\n=== Model mode env var ===")
+    from voice_clone import get_model_mode, BASE_MODEL_ID, CUSTOM_VOICE_MODEL_ID
+
+    mode = get_model_mode()
+    check("Default mode is 'base'", mode == "base", f"got '{mode}'")
+    check("BASE_MODEL_ID is 1.7B", "1.7B" in BASE_MODEL_ID, BASE_MODEL_ID)
+    check("CUSTOM_VOICE_MODEL_ID is 0.6B", "0.6B" in CUSTOM_VOICE_MODEL_ID,
+          CUSTOM_VOICE_MODEL_ID)
+
+
+def test_custom_voice_clone_blocked():
+    print("\n=== CustomVoice clone attempt blocked ===")
+    # Simulate QWEN_TTS_MODE=custom_voice by patching
+    import voice_clone
+    original = voice_clone._MODEL_MODE
+    voice_clone._MODEL_MODE = "custom_voice"
+    try:
+        voice_clone.create_voice_profile("dummy.wav")
+        check("Should have raised ValueError", False)
+    except ValueError as e:
+        check("create_voice_profile raises ValueError for custom_voice mode",
+              "Base model" in str(e), str(e))
+    except Exception as e:
+        check("Unexpected error type", False, str(e))
+    finally:
+        voice_clone._MODEL_MODE = original
+
+
+def test_tts_module_imports():
+    print("\n=== TTS module structure ===")
+    from tts import generate_audio_stream, split_into_chunks
+    import inspect
+
+    sig = inspect.signature(generate_audio_stream)
+    params = list(sig.parameters.keys())
+    check("generate_audio_stream has 'use_custom_voice' param",
+          "use_custom_voice" in params, f"params: {params}")
+    check("generate_audio_stream has 'custom_voice_speaker' param",
+          "custom_voice_speaker" in params, f"params: {params}")
+    check("generate_audio_stream has 'voice_profile_id' param",
+          "voice_profile_id" in params, f"params: {params}")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Latency Optimization Tests")
+    print("=" * 60)
+
+    test_global_matmul_precision()
+    test_dtype_selection()
+    test_attn_selection()
+    test_generation_params()
+    test_ref_audio_trimming()
+    test_model_mode_env()
+    test_custom_voice_clone_blocked()
+    test_tts_module_imports()
+
+    print("\n" + "=" * 60)
+    print(f"Results: {PASS} passed, {FAIL} failed")
+    print("=" * 60)
+    sys.exit(1 if FAIL > 0 else 0)

tts.py

@@ -1,9 +1,10 @@
 """
 TTS module — unified interface for text-to-speech synthesis.
 
-Supports two backends:
-  - Qwen3-TTS: voice-cloned synthesis using a cached voice profile
-  - Supertonic: fast stock voice fallback (no cloning)
+Supports three backends:
+  - Qwen3-TTS Base (1.7B): voice-cloned synthesis using a cached voice profile
+  - Qwen3-TTS CustomVoice (0.6B): fast predefined speakers (no cloning, lower latency)
+  - Supertonic: fast stock voice fallback (ONNX, no cloning)
 
 Usage:
     chunks = split_into_chunks(text)
@@ -44,20 +45,27 @@ def generate_audio_stream(
     chunks: list[str],
     voice_profile_id: str | None = None,
     voice_name: str = "F1",
+    use_custom_voice: bool = False,
+    custom_voice_speaker: str = "Chelsie",
 ):
     """
     Generator: synthesizes chunks in a background thread (maxsize=2 buffer).
     Yields (sample_rate, wav_array, chunk_idx, total_chunks, error_msg).
 
-    If voice_profile_id is provided, uses Qwen3-TTS with the cloned voice.
-    Otherwise falls back to Supertonic with the given voice_name.
+    Backend selection:
+      1. voice_profile_id provided → Qwen3-TTS Base (1.7B) with cloned voice
+      2. use_custom_voice=True → Qwen3-TTS CustomVoice (0.6B) with predefined speaker
+      3. Otherwise → Supertonic with given voice_name
     """
     n = len(chunks)
     chunk_q: queue.Queue = queue.Queue(maxsize=2)
 
     if voice_profile_id:
         _start_qwen_worker(chunks, voice_profile_id, chunk_q)
-        sample_rate = 24000  # Qwen3-TTS output rate
+        sample_rate = 24000
+    elif use_custom_voice:
+        _start_custom_voice_worker(chunks, custom_voice_speaker, chunk_q)
+        sample_rate = 24000
     else:
         sample_rate = _start_supertonic_worker(chunks, voice_name, chunk_q)
 
@@ -70,7 +78,7 @@ def generate_audio_stream(
 
 
 def _start_qwen_worker(chunks, profile_id, chunk_q):
-    """Background thread: synthesize chunks with Qwen3-TTS voice clone."""
+    """Background thread: synthesize chunks with Qwen3-TTS voice clone (Base 1.7B)."""
     def _worker():
         from voice_clone import synthesize_cloned
         for i, stmt in enumerate(chunks):
@@ -86,6 +94,23 @@ def _start_qwen_worker(chunks, profile_id, chunk_q):
     threading.Thread(target=_worker, daemon=True).start()
 
 
+def _start_custom_voice_worker(chunks, speaker, chunk_q):
+    """Background thread: synthesize chunks with Qwen3-TTS CustomVoice (0.6B)."""
+    def _worker():
+        from voice_clone import synthesize_custom_voice
+        for i, stmt in enumerate(chunks):
+            try:
+                wav, _sr = synthesize_custom_voice(stmt, speaker=speaker)
+                chunk_q.put((i, wav, None))
+            except Exception as exc:
+                logger.exception("CustomVoice synthesis failed on chunk %d", i)
+                chunk_q.put((i, None, str(exc)))
+                return
+        chunk_q.put(_SENTINEL)
+
+    threading.Thread(target=_worker, daemon=True).start()
+
+
 def _start_supertonic_worker(chunks, voice_name, chunk_q):
     """Background thread: synthesize chunks with Supertonic (stock voice)."""
     tts = _get_supertonic()

voice_clone.py

@@ -1,57 +1,189 @@
 """
 Voice cloning module — wraps Qwen3-TTS for zero-shot voice cloning.
 
+Supports two backends:
+  - Base 1.7B: zero-shot voice cloning from reference audio
+  - CustomVoice 0.6B: fast predefined speakers (no cloning, lower latency)
+
+Latency optimizations applied:
+  - bfloat16 / float16 precision (auto-detected per GPU arch)
+  - FlashAttention-2 when available
+  - Streaming mode (non_streaming_mode=False)
+  - Reduced sampling params (top_k=20, temperature=0.7)
+  - Capped max_new_tokens=1024
+  - torch.set_float32_matmul_precision('high')
+  - Reference audio trimmed to 3-5s for faster embedding extraction
+
 Usage:
     profile_id = create_voice_profile(ref_audio_path)
     wav, sr = synthesize_cloned(text, profile_id)
 """
 import logging
+import os
 import uuid
 import threading
 
 import numpy as np
+import soundfile as sf
 import torch
 
 logger = logging.getLogger(__name__)
 
+# ---------------------------------------------------------------------------
+# Global PyTorch optimizations
+# ---------------------------------------------------------------------------
+torch.set_float32_matmul_precision("high")
+
+# ---------------------------------------------------------------------------
+# Model configuration
+# ---------------------------------------------------------------------------
+# Base model for zero-shot voice cloning (1.7B)
+BASE_MODEL_ID = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+# CustomVoice model for fast predefined speakers (0.6B, ~3x faster)
+CUSTOM_VOICE_MODEL_ID = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
+
+# Choose model via env var: "base" (default) or "custom_voice"
+_MODEL_MODE = os.environ.get("QWEN_TTS_MODE", "base").lower()
+
+# Optimized generation parameters (reduced from defaults: top_k=50, temp=0.9, max=2048)
+GENERATION_PARAMS = dict(
+    top_k=20,
+    temperature=0.7,
+    subtalker_top_k=20,
+    subtalker_temperature=0.7,
+    max_new_tokens=1024,
+)
+
+# Reference audio limits (seconds) — 3-5s is optimal for Qwen3-TTS
+REF_AUDIO_MIN_SEC = 3.0
+REF_AUDIO_MAX_SEC = 10.0
+REF_AUDIO_TARGET_SR = 24000
+
+# ---------------------------------------------------------------------------
 # Server-side cache: { profile_id -> VoiceClonePromptItem list }
+# ---------------------------------------------------------------------------
 _PROFILE_CACHE: dict[str, list] = {}
 _cache_lock = threading.Lock()
 
 _qwen_tts_model = None
 _model_lock = threading.Lock()
 
-QWEN_TTS_MODEL_ID = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+
+def _select_dtype() -> torch.dtype:
+    """Pick optimal dtype based on GPU architecture."""
+    if not torch.cuda.is_available():
+        return torch.float32
+    cap = torch.cuda.get_device_capability()
+    # bfloat16 requires compute capability >= 8.0 (Ampere+)
+    if cap[0] >= 8:
+        return torch.bfloat16
+    return torch.float16
+
+
+def _select_attn_impl() -> str:
+    """Use FlashAttention-2 if available, else default."""
+    try:
+        import flash_attn  # noqa: F401
+        return "flash_attention_2"
+    except ImportError:
+        logger.info("flash-attn not installed — using default attention.")
+        return "sdpa"
 
 
 def get_qwen_tts():
-    """Lazy-load Qwen3-TTS model. Thread-safe, cached globally."""
+    """Lazy-load Qwen3-TTS model with latency optimizations. Thread-safe."""
     global _qwen_tts_model
     if _qwen_tts_model is None:
         with _model_lock:
             if _qwen_tts_model is None:
                 from qwen_tts import Qwen3TTSModel
 
-                logger.info("Loading %s...", QWEN_TTS_MODEL_ID)
+                model_id = CUSTOM_VOICE_MODEL_ID if _MODEL_MODE == "custom_voice" else BASE_MODEL_ID
+                dtype = _select_dtype()
+                attn_impl = _select_attn_impl()
+
+                logger.info(
+                    "Loading %s (dtype=%s, attn=%s)...",
+                    model_id, dtype, attn_impl,
+                )
                 _qwen_tts_model = Qwen3TTSModel.from_pretrained(
-                    QWEN_TTS_MODEL_ID,
+                    model_id,
                     device_map="cuda" if torch.cuda.is_available() else "cpu",
-                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                    torch_dtype=dtype,
+                    attn_implementation=attn_impl,
                 )
-                logger.info("Qwen3-TTS loaded.")
+
+                # Attempt torch.compile on decoder/predictor for extra speed
+                _try_torch_compile(_qwen_tts_model)
+
+                logger.info("Qwen3-TTS loaded (%s mode).", _MODEL_MODE)
     return _qwen_tts_model
 
 
+def _try_torch_compile(wrapper):
+    """Best-effort torch.compile on model submodules."""
+    if not hasattr(torch, "compile"):
+        return
+    model = wrapper.model
+    for name in ("decoder", "predictor", "speech_tokenizer"):
+        submod = getattr(model, name, None)
+        if submod is not None and isinstance(submod, torch.nn.Module):
+            try:
+                compiled = torch.compile(submod, mode="reduce-overhead")
+                setattr(model, name, compiled)
+                logger.info("torch.compile applied to model.%s", name)
+            except Exception as exc:
+                logger.debug("torch.compile skipped for %s: %s", name, exc)
+
+
+def _trim_reference_audio(audio_path: str) -> str:
+    """
+    Trim reference audio to REF_AUDIO_MAX_SEC seconds if longer.
+    Returns path to trimmed file (or original if already short enough).
+    """
+    try:
+        info = sf.info(audio_path)
+        duration = info.duration
+        if duration <= REF_AUDIO_MAX_SEC:
+            return audio_path
+
+        logger.info(
+            "Reference audio %.1fs exceeds %.0fs limit — trimming.",
+            duration, REF_AUDIO_MAX_SEC,
+        )
+        data, sr = sf.read(audio_path)
+        max_samples = int(REF_AUDIO_MAX_SEC * sr)
+        trimmed = data[:max_samples]
+
+        trimmed_path = audio_path + ".trimmed.wav"
+        sf.write(trimmed_path, trimmed, sr)
+        return trimmed_path
+    except Exception as exc:
+        logger.warning("Could not trim reference audio: %s", exc)
+        return audio_path
+
+
 def create_voice_profile(ref_audio_path: str) -> str:
     """
     Extract speaker embedding from reference audio and cache it.
     Returns a profile_id string for later synthesis.
+
+    Reference audio is trimmed to 3-10s for optimal latency.
+    Only supported with the Base model (1.7B).
     """
+    if _MODEL_MODE == "custom_voice":
+        raise ValueError(
+            "Voice cloning requires the Base model (1.7B). "
+            "The 0.6B CustomVoice model only supports predefined speakers. "
+            "Set QWEN_TTS_MODE=base or remove the env var."
+        )
+
+    trimmed_path = _trim_reference_audio(ref_audio_path)
     model = get_qwen_tts()
 
     logger.info("Creating voice profile from %s...", ref_audio_path)
     prompt_items = model.create_voice_clone_prompt(
-        ref_audio=ref_audio_path,
+        ref_audio=trimmed_path,
         x_vector_only_mode=True,
     )
 
@@ -79,6 +211,30 @@ def synthesize_cloned(text: str, profile_id: str) -> tuple[np.ndarray, int]:
         text=text,
         language="english",
         voice_clone_prompt=prompt_items,
+        non_streaming_mode=False,
+        **GENERATION_PARAMS,
+    )
+
+    wav = np.concatenate(audio_list) if audio_list else np.zeros(0, dtype=np.float32)
+    return wav, sample_rate
+
+
+def synthesize_custom_voice(
+    text: str, speaker: str = "Chelsie", language: str = "english"
+) -> tuple[np.ndarray, int]:
+    """
+    Synthesize text with the 0.6B CustomVoice model (predefined speakers).
+    Much faster than cloned synthesis but no voice cloning.
+    Returns (wav_array, sample_rate).
+    """
+    model = get_qwen_tts()
+
+    audio_list, sample_rate = model.generate_custom_voice(
+        text=text,
+        speaker=speaker,
+        language=language,
+        non_streaming_mode=False,
+        **GENERATION_PARAMS,
     )
 
     wav = np.concatenate(audio_list) if audio_list else np.zeros(0, dtype=np.float32)
@@ -99,3 +255,8 @@ def has_profile(profile_id: str | None) -> bool:
         return False
     with _cache_lock:
         return profile_id in _PROFILE_CACHE
+
+
+def get_model_mode() -> str:
+    """Return the current model mode ('base' or 'custom_voice')."""
+    return _MODEL_MODE