121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
"""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")
|