feat: add OpenAI-compatible Chatterbox TTS service

app.py

@@ -0,0 +1,121 @@
+"""OpenAI-compatible speech endpoint backed by ResembleAI Chatterbox."""
+
+import gc
+import os
+import threading
+import time
+from io import BytesIO
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import Response
+from pydantic import BaseModel, Field
+
+
+VOICES_DIR = Path(__file__).parent / "voices"
+IDLE_UNLOAD_SECONDS = 300
+
+
+class SpeechRequest(BaseModel):
+    model: str = "chatterbox"
+    input: str = Field(max_length=2000)
+    voice: str
+    response_format: str = "wav"
+    speed: float = 1.0
+
+
+class ModelManager:
+    """Owns the GPU model and releases it after a period of inactivity."""
+
+    def __init__(self) -> None:
+        self._model = None
+        self._unload_timer: threading.Timer | None = None
+        self._lock = threading.Lock()
+        self.last_request_time: float | None = None
+
+    def _cancel_unload_timer(self) -> None:
+        if self._unload_timer is not None:
+            self._unload_timer.cancel()
+            self._unload_timer = None
+
+    def _schedule_unload(self) -> None:
+        self._cancel_unload_timer()
+        self._unload_timer = threading.Timer(IDLE_UNLOAD_SECONDS, self.unload)
+        self._unload_timer.daemon = True
+        self._unload_timer.start()
+
+    def load_model(self):
+        if self._model is None:
+            from chatterbox.tts import ChatterboxTTS
+
+            self._model = ChatterboxTTS.from_pretrained(device="cuda")
+        return self._model
+
+    def generate(self, text: str, voice_path: Path, speed: float) -> bytes:
+        with self._lock:
+            self._cancel_unload_timer()
+            try:
+                model = self.load_model()
+                # Chatterbox does not expose a native speed control.  Keep the
+                # OpenAI-compatible field accepted without changing speech speed.
+                del speed
+                audio = model.generate(text, audio_prompt_path=str(voice_path))
+
+                import torchaudio
+
+                buffer = BytesIO()
+                torchaudio.save(buffer, audio, model.sr, format="wav")
+                return buffer.getvalue()
+            finally:
+                if self._model is not None:
+                    self.last_request_time = time.monotonic()
+                    self._schedule_unload()
+
+    def unload(self) -> None:
+        with self._lock:
+            self._cancel_unload_timer()
+            if self._model is None:
+                return
+            self._model = None
+            gc.collect()
+            import torch
+
+            torch.cuda.empty_cache()
+
+
+def resolve_voice(voice: str) -> Path:
+    candidate = Path(voice)
+    if candidate.is_absolute():
+        raise HTTPException(status_code=400, detail="voice must be a relative path")
+
+    voices_root = VOICES_DIR.resolve()
+    voice_path = Path(VOICES_DIR, voice).resolve()
+    if not voice_path.is_relative_to(voices_root):
+        raise HTTPException(status_code=400, detail="voice path escapes the voices directory")
+    if not voice_path.is_file() or not os.access(voice_path, os.R_OK):
+        raise HTTPException(status_code=400, detail="voice must be a readable regular file")
+    return voice_path
+
+
+app = FastAPI()
+model_manager = ModelManager()
+
+
+@app.get("/health")
+def health() -> dict[str, str]:
+    return {"status": "ok"}
+
+
+@app.get("/v1/models")
+def list_models() -> dict[str, list[dict[str, str]]]:
+    return {"data": [{"id": "chatterbox", "object": "model"}]}
+
+
+@app.post("/v1/audio/speech")
+def speech(request: SpeechRequest) -> Response:
+    if request.response_format != "wav":
+        raise HTTPException(status_code=400, detail="only wav response_format is supported")
+
+    voice_path = resolve_voice(request.voice)
+    audio = model_manager.generate(request.input, voice_path, request.speed)
+    return Response(content=audio, media_type="audio/wav")