Upload 6 files

app.py

@@ -1,19 +1,3 @@
-"""
-Chatterbox Turbo TTS -- FastAPI Server
-======================================
-Production-ready API with true real-time MP3 streaming,
-in-memory voice cloning, and fully non-blocking inference.
-
-Endpoints:
-  GET  /health              -> health check + optional warmup
-  GET  /info                -> model info, supported tags, parameters
-  POST /tts                 -> full audio response (WAV/MP3/FLAC)
-  POST /tts/stream          -> chunked MP3 streaming (MediaSource-ready)
-  POST /tts/true-stream     -> alias for /tts/stream (Kokoro compat)
-  POST /tts/stop/{stream_id}-> cancel a specific active stream
-  POST /tts/stop            -> cancel ALL active streams
-  POST /v1/audio/speech     -> OpenAI-compatible streaming
-"""
 import asyncio
 import io
 import json
@@ -235,6 +219,26 @@ def _helper_cancel_stream(helper_base_url: str, stream_id: str):
         pass
 
 
+def _helper_complete_stream(helper_base_url: str, stream_id: str):
+    """Best-effort stream completion cleanup on helper.
+
+    Falls back to cancel for backwards compatibility if helper does not expose
+    the completion endpoint yet.
+    """
+    try:
+        url = _build_helper_endpoint(helper_base_url, f"/internal/chunk/complete/{stream_id}")
+        req = urllib.request.Request(
+            url=url,
+            data=b"",
+            headers=_internal_headers(),
+            method="POST",
+        )
+        with urllib.request.urlopen(req, timeout=3.0):
+            pass
+    except Exception:
+        _helper_cancel_stream(helper_base_url, stream_id)
+
+
 # ═══════════════════════════════════════════════════════════════════
 # Endpoints
 # ═══════════════════════════════════════════════════════════════════
@@ -242,12 +246,19 @@ def _helper_cancel_stream(helper_base_url: str, stream_id: str):
 @app.get("/health")
 async def health(warm_up: bool = False):
     wrapper: ChatterboxWrapper = getattr(app.state, "wrapper", None)
+    with _internal_cancel_lock:
+        _purge_internal_stream_state_locked()
+        cancelled_count = len(_internal_cancelled_streams)
+        voice_state_count = len(_internal_stream_voice_keys)
+
     status = {
         "status": "healthy" if wrapper else "loading",
         "model_loaded": wrapper is not None,
         "model_dtype": Config.MODEL_DTYPE,
         "streaming_supported": True,
         "voice_cache_entries": wrapper._voice_cache.size if wrapper else 0,
+        "internal_cancelled_streams": cancelled_count,
+        "internal_stream_voice_states": voice_state_count,
     }
     if warm_up and wrapper:
         try:
@@ -259,31 +270,6 @@ async def health(warm_up: bool = False):
     return status
 
 
-@app.get("/info")
-async def info():
-    return {
-        "model": Config.MODEL_ID,
-        "dtype": Config.MODEL_DTYPE,
-        "sample_rate": Config.SAMPLE_RATE,
-        "paralinguistic_tags": list(Config.PARALINGUISTIC_TAGS),
-        "tag_usage": "Insert tags directly in text, e.g. 'That is so funny! [laugh] Anyway…'",
-        "parameters": {
-            "max_new_tokens": {"default": Config.MAX_NEW_TOKENS, "range": "64–2048"},
-            "repetition_penalty": {"default": Config.REPETITION_PENALTY, "range": "1.0–2.0"},
-        },
-        "voice_cloning": {
-            "description": "Upload 3–30s reference WAV/MP3 as 'voice_ref' field",
-            "max_upload_mb": Config.MAX_VOICE_UPLOAD_BYTES // (1024 * 1024),
-        },
-        "parallel_mode": {
-            "enabled": Config.ENABLE_PARALLEL_MODE,
-            "helper_configured": bool(Config.HELPER_BASE_URL),
-            "helper_base_url": Config.HELPER_BASE_URL or None,
-            "supports_voice_ref": True,
-        },
-    }
-
-
 # ── POST /tts ─────────────────────────────────────────────────────
 
 @app.post("/tts", response_class=Response)
@@ -329,9 +315,47 @@ async def text_to_speech(
 # ═══════════════════════════════════════════════════════════════════
 
 _active_streams: dict[str, threading.Event] = {}
-_internal_cancelled_streams: set[str] = set()
+# stream_id -> expires_at epoch seconds
+_internal_cancelled_streams: dict[str, float] = {}
 _internal_cancel_lock = threading.Lock()
-_internal_stream_voice_keys: dict[str, set[str]] = {}
+# stream_id -> (voice_keys, expires_at)
+_internal_stream_voice_keys: dict[str, tuple[set[str], float]] = {}
+
+
+def _purge_internal_stream_state_locked(now: Optional[float] = None):
+    now_ts = now if now is not None else time.time()
+
+    expired_cancel_ids = [
+        sid for sid, expires_at in _internal_cancelled_streams.items()
+        if expires_at <= now_ts
+    ]
+    for sid in expired_cancel_ids:
+        _internal_cancelled_streams.pop(sid, None)
+
+    expired_voice_state_ids = [
+        sid for sid, (_, expires_at) in _internal_stream_voice_keys.items()
+        if expires_at <= now_ts
+    ]
+    for sid in expired_voice_state_ids:
+        _internal_stream_voice_keys.pop(sid, None)
+
+
+def _touch_internal_stream_voice_keys_locked(stream_id: str):
+    if not stream_id:
+        return
+    entry = _internal_stream_voice_keys.get(stream_id)
+    if entry is None:
+        return
+    keys, _ = entry
+    _internal_stream_voice_keys[stream_id] = (
+        keys,
+        time.time() + max(1, Config.INTERNAL_STREAM_STATE_TTL_SEC),
+    )
+
+
+def _clear_internal_stream_state_locked(stream_id: str):
+    _internal_cancelled_streams.pop(stream_id, None)
+    _internal_stream_voice_keys.pop(stream_id, None)
 
 
 # ═══════════════════════════════════════════════════════════════════
@@ -457,6 +481,7 @@ def _parallel_odd_even_stream_generator(
     ready: dict[int, bytes] = {}
     first_error: Optional[Exception] = None
     workers_done = 0
+    stream_completed = False
 
     def _publish(idx: int, data: bytes):
         with cond:
@@ -608,9 +633,17 @@ def _parallel_odd_even_stream_generator(
 
             yield data
             next_idx += 1
+        stream_completed = (
+            next_idx >= total_chunks
+            and first_error is None
+            and not cancel_event.is_set()
+        )
     finally:
         cancel_event.set()
-        _helper_cancel_stream(helper_base_url, stream_id)
+        if stream_completed:
+            _helper_complete_stream(helper_base_url, stream_id)
+        else:
+            _helper_cancel_stream(helper_base_url, stream_id)
         odd_thread.join(timeout=1.0)
         even_thread.join(timeout=1.0)
         _active_streams.pop(stream_id, None)
@@ -777,8 +810,13 @@ async def internal_voice_register(http_request: Request):
     stream_id = (http_request.query_params.get("stream_id") or "").strip()
     if stream_id:
         with _internal_cancel_lock:
-            keys = _internal_stream_voice_keys.setdefault(stream_id, set())
+            _purge_internal_stream_state_locked()
+            keys, _ = _internal_stream_voice_keys.get(stream_id, (set(), 0.0))
             keys.add(voice_key)
+            _internal_stream_voice_keys[stream_id] = (
+                keys,
+                time.time() + max(1, Config.INTERNAL_STREAM_STATE_TTL_SEC),
+            )
 
     return {"status": "registered", "voice_key": voice_key}
 
@@ -795,8 +833,10 @@ async def internal_chunk_synthesize(
             raise HTTPException(403, "Forbidden")
 
     with _internal_cancel_lock:
+        _purge_internal_stream_state_locked()
         if request.stream_id in _internal_cancelled_streams:
             raise HTTPException(409, "Stream already cancelled")
+        _touch_internal_stream_voice_keys_locked(request.stream_id)
 
     wrapper: ChatterboxWrapper = getattr(app.state, "wrapper", None)
     if not wrapper:
@@ -805,6 +845,9 @@ async def internal_chunk_synthesize(
     voice_profile = wrapper.default_voice
     if request.voice_key:
         cached_voice = wrapper._voice_cache.get(request.voice_key)
+        if cached_voice is None:
+            # Built-in voices are permanent in wrapper registry even if TTL cache entry expired.
+            cached_voice = wrapper.get_builtin_voice_by_hash(request.voice_key)
         if cached_voice is None:
             raise HTTPException(409, "Voice key expired or not found")
         voice_profile = cached_voice
@@ -845,11 +888,28 @@ async def internal_chunk_cancel(stream_id: str, http_request: Request):
             raise HTTPException(403, "Forbidden")
 
     with _internal_cancel_lock:
-        _internal_cancelled_streams.add(stream_id)
+        _purge_internal_stream_state_locked()
+        _internal_cancelled_streams[stream_id] = (
+            time.time() + max(1, Config.INTERNAL_CANCEL_TTL_SEC)
+        )
         _internal_stream_voice_keys.pop(stream_id, None)
     return {"status": "cancelled", "stream_id": stream_id}
 
 
+@app.post("/internal/chunk/complete/{stream_id}")
+async def internal_chunk_complete(stream_id: str, http_request: Request):
+    """Best-effort immediate cleanup after stream completes normally."""
+    if Config.INTERNAL_SHARED_SECRET:
+        provided = http_request.headers.get("X-Internal-Secret", "")
+        if provided != Config.INTERNAL_SHARED_SECRET:
+            raise HTTPException(403, "Forbidden")
+
+    with _internal_cancel_lock:
+        _purge_internal_stream_state_locked()
+        _clear_internal_stream_state_locked(stream_id)
+    return {"status": "completed", "stream_id": stream_id}
+
+
 @app.post("/v1/audio/speech")
 async def openai_compatible_tts(request: TTSJsonRequest):
     """OpenAI-compatible streaming endpoint (JSON body, no file upload).

chatterbox_wrapper.py

@@ -27,6 +27,7 @@ import tempfile
 import time
 from collections import OrderedDict
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Callable, Generator, Optional
 
 import librosa
@@ -48,6 +49,21 @@ _SUPPORTED_AUDIO_EXTENSIONS = {
 }
 
 
+def _slugify(text: str) -> str:
+    buf = []
+    prev_underscore = False
+    for ch in text.strip().lower():
+        if ch.isalnum():
+            buf.append(ch)
+            prev_underscore = False
+        else:
+            if not prev_underscore:
+                buf.append("_")
+                prev_underscore = True
+    slug = "".join(buf).strip("_")
+    return slug or "voice"
+
+
 # ═══════════════════════════════════════════════════════════════════
 # Data Structures
 # ═══════════════════════════════════════════════════════════════════
@@ -203,8 +219,15 @@ class ChatterboxWrapper:
             ttl_seconds=self.cfg.VOICE_CACHE_TTL_SEC,
         )
 
-        logger.info("Encoding default reference voice …")
-        self.default_voice = self._load_default_voice()
+        self._builtin_voice_profiles: dict[str, VoiceProfile] = {}
+        self._builtin_voice_bytes: dict[str, bytes] = {}
+        self._builtin_voice_by_hash: dict[str, VoiceProfile] = {}
+        self._voice_alias_to_id: dict[str, str] = {}
+        self._builtin_voice_catalog: list[dict] = []
+        self._default_voice_id: str = "default"
+
+        logger.info("Loading built-in voices (HF default + local samples) …")
+        self.default_voice = self._load_builtin_voices()
 
         logger.info("✅ ChatterboxWrapper ready")
 
@@ -260,16 +283,185 @@ class ChatterboxWrapper:
         opts.enable_mem_reuse = True
         return opts
 
-    # ─── Default voice ────────────────────────────────────────────
+    # ─── Built-in voices (HF default + local samples) ────────────
 
-    def _load_default_voice(self) -> VoiceProfile:
+    def _download_hf_default_voice_bytes(self) -> bytes:
         path = hf_hub_download(
             self.cfg.DEFAULT_VOICE_REPO,
             filename=self.cfg.DEFAULT_VOICE_FILE,
             cache_dir=self.cfg.MODELS_DIR,
         )
-        audio, _ = librosa.load(path, sr=self.cfg.SAMPLE_RATE)
-        return self._encode_audio_array(audio, audio_hash="__default__")
+        return Path(path).read_bytes()
+
+    def _list_local_voice_paths(self) -> list[Path]:
+        wrapper_dir = Path(__file__).resolve().parent
+
+        # Support both module-level and repo-root deployment layouts.
+        candidates = []
+        for d in (wrapper_dir, Path.cwd().resolve(), wrapper_dir.parent):
+            try:
+                resolved = d.resolve()
+            except Exception:
+                continue
+            if resolved.is_dir() and resolved not in candidates:
+                candidates.append(resolved)
+
+        voices: list[Path] = []
+        seen_real_paths: set[str] = set()
+        for root in candidates:
+            try:
+                entries = sorted(root.iterdir(), key=lambda x: x.name.lower())
+            except Exception:
+                continue
+
+            for p in entries:
+                if not p.is_file():
+                    continue
+                if p.suffix.lower() not in _SUPPORTED_AUDIO_EXTENSIONS:
+                    continue
+                real_path = str(p.resolve())
+                if real_path in seen_real_paths:
+                    continue
+                seen_real_paths.add(real_path)
+                voices.append(p)
+
+        return voices
+
+    def _make_unique_voice_id(self, preferred: str) -> str:
+        base = _slugify(preferred)
+        candidate = base
+        idx = 2
+        while candidate in self._builtin_voice_profiles:
+            candidate = f"{base}_{idx}"
+            idx += 1
+        return candidate
+
+    def _register_builtin_voice(
+        self,
+        *,
+        preferred_id: str,
+        display_name: str,
+        source: str,
+        source_ref: str,
+        audio_bytes: bytes,
+        is_default: bool = False,
+    ) -> str:
+        if not audio_bytes:
+            raise ValueError("Voice file is empty")
+
+        voice_id = self._make_unique_voice_id(preferred_id)
+        audio_hash = hashlib.md5(audio_bytes).hexdigest()
+
+        profile = self._voice_cache.get(audio_hash)
+        if profile is None:
+            audio = _load_audio_bytes(audio_bytes, sr=self.cfg.SAMPLE_RATE)
+            profile = self._encode_audio_array(audio, audio_hash=audio_hash)
+            self._voice_cache.put(audio_hash, profile)
+        else:
+            # Keep hash attached to cached profile for metadata/voice-key usage.
+            profile.audio_hash = audio_hash
+
+        self._builtin_voice_profiles[voice_id] = profile
+        self._builtin_voice_bytes[voice_id] = audio_bytes
+        if audio_hash:
+            self._builtin_voice_by_hash[audio_hash] = profile
+
+        aliases: list[str] = []
+        for alias in (voice_id, _slugify(Path(display_name).stem)):
+            if alias not in self._voice_alias_to_id:
+                self._voice_alias_to_id[alias] = voice_id
+                aliases.append(alias)
+
+        if is_default:
+            self._default_voice_id = voice_id
+            self._voice_alias_to_id["default"] = voice_id
+            if "default" not in aliases:
+                aliases.append("default")
+
+        self._builtin_voice_catalog.append(
+            {
+                "id": voice_id,
+                "display_name": display_name,
+                "source": source,
+                "source_ref": source_ref,
+                "aliases": aliases,
+                "voice_key": audio_hash,
+            }
+        )
+        return voice_id
+
+    def _load_builtin_voices(self) -> VoiceProfile:
+        # 1) HF default voice (kept as true default fallback)
+        hf_bytes = self._download_hf_default_voice_bytes()
+        self._register_builtin_voice(
+            preferred_id="default_hf_voice",
+            display_name=self.cfg.DEFAULT_VOICE_FILE,
+            source="huggingface",
+            source_ref=f"{self.cfg.DEFAULT_VOICE_REPO}:{self.cfg.DEFAULT_VOICE_FILE}",
+            audio_bytes=hf_bytes,
+            is_default=True,
+        )
+
+        # 2) Local voice samples placed next to app files
+        for path in self._list_local_voice_paths():
+            # Avoid duplicate entry if someone also copied default_voice.wav locally.
+            if path.name == self.cfg.DEFAULT_VOICE_FILE:
+                continue
+            try:
+                self._register_builtin_voice(
+                    preferred_id=path.stem,
+                    display_name=path.name,
+                    source="local",
+                    source_ref=str(path.name),
+                    audio_bytes=path.read_bytes(),
+                    is_default=False,
+                )
+            except Exception as e:
+                logger.warning(f"Skipping local voice {path.name}: {e}")
+
+        default_profile = self._builtin_voice_profiles.get(self._default_voice_id)
+        if default_profile is None:
+            raise RuntimeError("Default built-in voice could not be initialized")
+
+        logger.info(
+            f"Built-in voices loaded: {len(self._builtin_voice_catalog)} "
+            f"(default={self._default_voice_id})"
+        )
+        return default_profile
+
+    def list_builtin_voices(self) -> list[dict]:
+        """Return metadata for startup-preloaded voices."""
+        return [dict(v) for v in self._builtin_voice_catalog]
+
+    @property
+    def default_voice_name(self) -> str:
+        return self._default_voice_id
+
+    def resolve_voice_id(self, voice_name: Optional[str]) -> str:
+        if voice_name is None:
+            return self._default_voice_id
+        key = _slugify(str(voice_name))
+        if not key:
+            return self._default_voice_id
+        voice_id = self._voice_alias_to_id.get(key)
+        if voice_id is None:
+            available = ", ".join(sorted(self._voice_alias_to_id.keys()))
+            raise ValueError(f"Unknown voice '{voice_name}'. Available: {available}")
+        return voice_id
+
+    def get_builtin_voice(self, voice_name: Optional[str]) -> VoiceProfile:
+        voice_id = self.resolve_voice_id(voice_name)
+        profile = self._builtin_voice_profiles[voice_id]
+        if profile.audio_hash:
+            self._voice_cache.put(profile.audio_hash, profile)
+        return profile
+
+    def get_builtin_voice_bytes(self, voice_name: Optional[str]) -> Optional[bytes]:
+        voice_id = self.resolve_voice_id(voice_name)
+        return self._builtin_voice_bytes.get(voice_id)
+
+    def get_builtin_voice_by_hash(self, audio_hash: str) -> Optional[VoiceProfile]:
+        return self._builtin_voice_by_hash.get((audio_hash or "").strip())
 
     # ─── Voice encoding ──────────────────────────────────────────
 

config.py

@@ -17,6 +17,12 @@ def _get_bool(name: str, default: bool) -> bool:
     return raw.strip().lower() in {"1", "true", "yes", "on"}
 
 
+def _get_csv(name: str, default: str) -> tuple[str, ...]:
+    raw = os.getenv(name, default)
+    items = [x.strip() for x in raw.split(",")]
+    return tuple(x for x in items if x)
+
+
 class Config:
     # ── Model ────────────────────────────────────────────────────
     MODEL_ID: str = os.getenv("CB_MODEL_ID", "ResembleAI/chatterbox-turbo-ONNX")
@@ -67,6 +73,12 @@ class Config:
     # (not a model — just a reference WAV, safe to use from any source).
     DEFAULT_VOICE_REPO: str = "onnx-community/chatterbox-ONNX"
     DEFAULT_VOICE_FILE: str = "default_voice.wav"
+    DEFAULT_VOICE_REPOS: tuple[str, ...] = _get_csv(
+        "CB_DEFAULT_VOICE_REPOS",
+        DEFAULT_VOICE_REPO,
+    )
+    PRELOAD_BUILTIN_VOICES: bool = _get_bool("CB_PRELOAD_BUILTIN_VOICES", True)
+    MAX_PRELOAD_BUILTIN_VOICES: int = int(os.getenv("CB_MAX_PRELOAD_BUILTIN_VOICES", "64"))
     MAX_VOICE_UPLOAD_BYTES: int = 10 * 1024 * 1024   # 10 MB
     MIN_REF_DURATION_SEC: float = 1.5
     MAX_REF_DURATION_SEC: float = 30.0
@@ -77,11 +89,19 @@ class Config:
     # Smaller chunks = faster TTFB (first audio arrives sooner)
     # ~200 chars ≈ 1–2 sentences ≈ fastest first-chunk on 2 vCPU
     MAX_CHUNK_CHARS: int = int(os.getenv("CB_MAX_CHUNK_CHARS", "100"))
-    # Additive parallel mode (odd/even split across primary/helper).
+    # Additive parallel mode (3-way split: primary + helper1 + helper2).
     ENABLE_PARALLEL_MODE: bool = _get_bool("CB_ENABLE_PARALLEL_MODE", True)
     HELPER_BASE_URL: str = os.getenv("CB_HELPER_BASE_URL", "").strip()
+    HELPER1_BASE_URL: str = os.getenv(
+        "CB_HELPER1_BASE_URL",
+        HELPER_BASE_URL,
+    ).strip()
+    HELPER2_BASE_URL: str = os.getenv("CB_HELPER2_BASE_URL", "").strip()
     HELPER_TIMEOUT_SEC: float = float(os.getenv("CB_HELPER_TIMEOUT_SEC", "45"))
     HELPER_RETRY_ONCE: bool = _get_bool("CB_HELPER_RETRY_ONCE", True)
+    # Internal housekeeping TTLs to avoid retaining stream metadata indefinitely.
+    INTERNAL_CANCEL_TTL_SEC: int = int(os.getenv("CB_INTERNAL_CANCEL_TTL_SEC", "120"))
+    INTERNAL_STREAM_STATE_TTL_SEC: int = int(os.getenv("CB_INTERNAL_STREAM_STATE_TTL_SEC", "600"))
     # Optional shared secret for internal chunk endpoints.
     INTERNAL_SHARED_SECRET: str = os.getenv("CB_INTERNAL_SHARED_SECRET", "").strip()
 

text_processor.py

@@ -222,11 +222,10 @@ def sanitize(text: str) -> str:
     text = _RE_MULTI_NEWLINE.sub("\n\n", text)
     text = text.strip()
 
-    # 9. Strip abbreviation dots (Mr. → Mr, Dr. → Dr, etc.)
-    #    The dot is not needed for correct TTS pronunciation and removing it
-    #    prevents false sentence-boundary splits in split_for_streaming().
+    # 9. Restore abbreviations
     for idx, original in abbrevs_found:
-        text = text.replace(f"§ABR{idx}§", "")
+        # Restore the full abbreviation with its period
+        text = text.replace(f"§ABR{idx}§", ".")
 
     # 10. Restore paralinguistic tags
     for idx, original in tags_found: