Upload 4 files

app.py

@@ -1,3 +1,19 @@
+"""
+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
@@ -219,26 +235,6 @@ 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
 # ═══════════════════════════════════════════════════════════════════
@@ -246,19 +242,12 @@ def _helper_complete_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:
@@ -270,6 +259,31 @@ 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)
@@ -315,47 +329,9 @@ async def text_to_speech(
 # ═══════════════════════════════════════════════════════════════════
 
 _active_streams: dict[str, threading.Event] = {}
-# stream_id -> expires_at epoch seconds
-_internal_cancelled_streams: dict[str, float] = {}
+_internal_cancelled_streams: set[str] = set()
 _internal_cancel_lock = threading.Lock()
-# 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)
+_internal_stream_voice_keys: dict[str, set[str]] = {}
 
 
 # ═══════════════════════════════════════════════════════════════════
@@ -481,7 +457,6 @@ 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:
@@ -633,17 +608,9 @@ 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()
-        if stream_completed:
-            _helper_complete_stream(helper_base_url, stream_id)
-        else:
-            _helper_cancel_stream(helper_base_url, stream_id)
+        _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)
@@ -810,13 +777,8 @@ 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:
-            _purge_internal_stream_state_locked()
-            keys, _ = _internal_stream_voice_keys.get(stream_id, (set(), 0.0))
+            keys = _internal_stream_voice_keys.setdefault(stream_id, set())
             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}
 
@@ -833,10 +795,8 @@ 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:
@@ -845,9 +805,6 @@ 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
@@ -888,28 +845,11 @@ async def internal_chunk_cancel(stream_id: str, http_request: Request):
             raise HTTPException(403, "Forbidden")
 
     with _internal_cancel_lock:
-        _purge_internal_stream_state_locked()
-        _internal_cancelled_streams[stream_id] = (
-            time.time() + max(1, Config.INTERNAL_CANCEL_TTL_SEC)
-        )
+        _internal_cancelled_streams.add(stream_id)
         _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,7 +27,6 @@ import tempfile
 import time
 from collections import OrderedDict
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Callable, Generator, Optional
 
 import librosa
@@ -49,21 +48,6 @@ _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
 # ═══════════════════════════════════════════════════════════════════
@@ -219,15 +203,8 @@ class ChatterboxWrapper:
             ttl_seconds=self.cfg.VOICE_CACHE_TTL_SEC,
         )
 
-        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("Encoding default reference voice …")
+        self.default_voice = self._load_default_voice()
 
         logger.info("✅ ChatterboxWrapper ready")
 
@@ -283,185 +260,16 @@ class ChatterboxWrapper:
         opts.enable_mem_reuse = True
         return opts
 
-    # ─── Built-in voices (HF default + local samples) ────────────
+    # ─── Default voice ────────────────────────────────────────────
 
-    def _download_hf_default_voice_bytes(self) -> bytes:
+    def _load_default_voice(self) -> VoiceProfile:
         path = hf_hub_download(
             self.cfg.DEFAULT_VOICE_REPO,
             filename=self.cfg.DEFAULT_VOICE_FILE,
             cache_dir=self.cfg.MODELS_DIR,
         )
-        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())
+        audio, _ = librosa.load(path, sr=self.cfg.SAMPLE_RATE)
+        return self._encode_audio_array(audio, audio_hash="__default__")
 
     # ─── Voice encoding ──────────────────────────────────────────
 

config.py

@@ -17,12 +17,6 @@ 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")
@@ -73,12 +67,6 @@ 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
@@ -89,19 +77,11 @@ 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 (3-way split: primary + helper1 + helper2).
+    # Additive parallel mode (odd/even split across primary/helper).
     ENABLE_PARALLEL_MODE: bool = _get_bool("CB_ENABLE_PARALLEL_MODE", True)
-    HELPER_BASE_URL: str = os.getenv("CB_HELPER_BASE_URL", "https://shadowhunter222-chab2.hf.space").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", "https://shadowhunter222-chab3.hf.space").strip()
+    HELPER_BASE_URL: str = os.getenv("CB_HELPER_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()