Spaces:
Paused
Paused
| import base64 | |
| import os | |
| import tempfile | |
| import uuid | |
| from pathlib import Path | |
| from threading import Lock | |
| from typing import Dict, Optional | |
| import requests | |
| import torch | |
| import torchaudio | |
| from torchaudio.transforms import Resample | |
| from fastapi import BackgroundTasks, Body, FastAPI, Header, HTTPException | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from pydantic import BaseModel, Field, HttpUrl | |
| SPACE_API_KEY = os.getenv("SPACE_API_KEY") | |
| DUB_API_KEY = os.getenv("DUB_API_KEY") | |
| AUTH_KEYS = {key for key in (SPACE_API_KEY, DUB_API_KEY) if key} | |
| HF_TOKEN = ( | |
| os.getenv("HUGGING_FACE_HUB_TOKEN") | |
| or os.getenv("HUGGINGFACEHUB_API_TOKEN") | |
| or os.getenv("HF_TOKEN") | |
| ) | |
| MODEL_REPO = "IndexTeam/IndexTTS-2" | |
| MODEL_DIR = os.getenv("MODEL_DIR", "/data/indextts2") | |
| MAX_TEXT_LENGTH = 1000 | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| SPEAKER_DIR = Path(tempfile.gettempdir(), "indextts2-speakers") | |
| SPEAKER_DIR.mkdir(parents=True, exist_ok=True) | |
| JOBS: Dict[str, Dict[str, str]] = {} | |
| JOB_LOCK = Lock() | |
| INFER_LOCK = Lock() | |
| if HF_TOKEN: | |
| os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN | |
| os.environ["HF_TOKEN"] = HF_TOKEN | |
| os.makedirs(MODEL_DIR, exist_ok=True) | |
| from huggingface_hub import snapshot_download | |
| if not Path(MODEL_DIR, "config.yaml").exists(): | |
| snapshot_download(repo_id=MODEL_REPO, local_dir=MODEL_DIR, token=HF_TOKEN) | |
| from indextts.infer_v2 import IndexTTS2 | |
| tts_model = IndexTTS2( | |
| cfg_path=os.path.join(MODEL_DIR, "config.yaml"), | |
| model_dir=MODEL_DIR, | |
| use_fp16=torch.cuda.is_available(), | |
| use_cuda_kernel=False, | |
| use_deepspeed=False, | |
| ) | |
| app = FastAPI(title="indextts2-api", version="2.0.0") | |
| class SpeakerRequest(BaseModel): | |
| speaker_wav: str = Field(..., description="HTTPS URL or base64 WAV") | |
| class GenerateRequest(BaseModel): | |
| text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH) | |
| speaker_id: Optional[str] = None | |
| speaker_wav: Optional[str] = None | |
| emotion_wav: Optional[str] = None | |
| emotion_alpha: float = Field(0.9, ge=0.0, le=1.0) | |
| def _require_api_key(x_api_key: Optional[str]): | |
| if AUTH_KEYS and x_api_key not in AUTH_KEYS: | |
| raise HTTPException(status_code=401, detail="Unauthorized") | |
| def _write_temp_audio_from_url(url: str) -> str: | |
| parsed = HttpUrl(url) | |
| response = requests.get(str(parsed), stream=True, timeout=60) | |
| if response.status_code >= 400: | |
| raise HTTPException(status_code=400, detail="Could not fetch audio") | |
| suffix = Path(parsed.path).suffix or ".wav" | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: | |
| for chunk in response.iter_content(chunk_size=1024 * 1024): | |
| if chunk: | |
| tmp.write(chunk) | |
| return tmp.name | |
| def _write_temp_audio_from_base64(payload: str) -> str: | |
| if payload.startswith("data:"): | |
| payload = payload.split(",", 1)[-1] | |
| try: | |
| raw = base64.b64decode(payload, validate=True) | |
| except Exception as exc: | |
| raise HTTPException(status_code=400, detail="Invalid base64 audio") from exc | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp: | |
| tmp.write(raw) | |
| return tmp.name | |
| def _temp_audio(value: str) -> str: | |
| if value.startswith("http://") or value.startswith("https://"): | |
| return _write_temp_audio_from_url(value) | |
| return _write_temp_audio_from_base64(value) | |
| def _preprocess_audio(path: str, target_sr: int = 24000, target_peak: float = 0.98) -> str: | |
| wav, sr = torchaudio.load(path) | |
| if wav.shape[0] > 1: | |
| wav = wav.mean(dim=0, keepdim=True) | |
| if sr != target_sr: | |
| wav = Resample(orig_freq=sr, new_freq=target_sr)(wav) | |
| sr = target_sr | |
| peak = wav.abs().max().item() if wav.numel() else 0.0 | |
| if peak > 0: | |
| wav = wav * min(target_peak / peak, 1.0) | |
| torchaudio.save(path, wav, sr, bits_per_sample=16) | |
| return path | |
| def _cleanup_files(*files: Optional[str]): | |
| for file_path in files: | |
| if file_path: | |
| Path(file_path).unlink(missing_ok=True) | |
| def _set_job(job_id: str, **kwargs): | |
| with JOB_LOCK: | |
| JOBS[job_id] = {**JOBS.get(job_id, {}), **kwargs} | |
| def _get_job(job_id: str) -> Optional[Dict[str, str]]: | |
| with JOB_LOCK: | |
| job = JOBS.get(job_id) | |
| return dict(job) if job else None | |
| def _pop_job(job_id: str) -> Optional[Dict[str, str]]: | |
| with JOB_LOCK: | |
| return JOBS.pop(job_id, None) | |
| def _registered_speaker_path(speaker_id: str) -> Path: | |
| if not speaker_id or any(ch not in "0123456789abcdef-" for ch in speaker_id.lower()): | |
| raise HTTPException(status_code=400, detail="Invalid speaker_id") | |
| path = SPEAKER_DIR / f"{speaker_id}.wav" | |
| if not path.exists(): | |
| raise HTTPException(status_code=404, detail="Speaker not found; register it again") | |
| return path | |
| def register_speaker(payload: SpeakerRequest, x_api_key: Optional[str] = Header(default=None)): | |
| _require_api_key(x_api_key) | |
| temp_path = _temp_audio(payload.speaker_wav) | |
| try: | |
| _preprocess_audio(temp_path) | |
| speaker_id = str(uuid.uuid4()) | |
| final_path = SPEAKER_DIR / f"{speaker_id}.wav" | |
| Path(temp_path).replace(final_path) | |
| return {"speaker_id": speaker_id} | |
| except Exception: | |
| _cleanup_files(temp_path) | |
| raise | |
| def delete_speaker(speaker_id: str, x_api_key: Optional[str] = Header(default=None)): | |
| _require_api_key(x_api_key) | |
| path = _registered_speaker_path(speaker_id) | |
| path.unlink(missing_ok=True) | |
| return {"deleted": speaker_id} | |
| def _run_generate_job(job_id: str, payload: dict): | |
| speaker_temp = None | |
| emotion_temp = None | |
| output_file = None | |
| _set_job(job_id, status="processing") | |
| try: | |
| if payload.get("speaker_id"): | |
| speaker_file = str(_registered_speaker_path(payload["speaker_id"])) | |
| elif payload.get("speaker_wav"): | |
| speaker_temp = _preprocess_audio(_temp_audio(payload["speaker_wav"])) | |
| speaker_file = speaker_temp | |
| else: | |
| raise HTTPException(status_code=400, detail="speaker_id or speaker_wav is required") | |
| if payload.get("emotion_wav"): | |
| emotion_temp = _preprocess_audio(_temp_audio(payload["emotion_wav"])) | |
| output_file = os.path.join(tempfile.gettempdir(), f"indextts2-{uuid.uuid4()}.wav") | |
| with INFER_LOCK: | |
| tts_model.infer( | |
| spk_audio_prompt=speaker_file, | |
| text=payload["text"], | |
| output_path=output_file, | |
| emo_audio_prompt=emotion_temp, | |
| emo_alpha=float(payload.get("emotion_alpha", 0.9)), | |
| use_random=False, | |
| interval_silence=80, | |
| verbose=False, | |
| ) | |
| _preprocess_audio(output_file) | |
| if not Path(output_file).exists(): | |
| raise RuntimeError("TTS output was not created") | |
| _cleanup_files(speaker_temp, emotion_temp) | |
| _set_job(job_id, status="completed", output_file=output_file) | |
| except Exception as exc: | |
| _cleanup_files(speaker_temp, emotion_temp, output_file) | |
| _set_job(job_id, status="error", error=str(exc)) | |
| def generate( | |
| payload: GenerateRequest = Body(...), | |
| background_tasks: BackgroundTasks = BackgroundTasks(), | |
| x_api_key: Optional[str] = Header(default=None), | |
| ): | |
| _require_api_key(x_api_key) | |
| if not payload.speaker_id and not payload.speaker_wav: | |
| raise HTTPException(status_code=400, detail="speaker_id or speaker_wav is required") | |
| job_id = str(uuid.uuid4()) | |
| _set_job(job_id, status="queued") | |
| data = payload.model_dump() if hasattr(payload, "model_dump") else payload.dict() | |
| background_tasks.add_task(_run_generate_job, job_id, data) | |
| return JSONResponse( | |
| status_code=202, | |
| content={ | |
| "job_id": job_id, | |
| "status": "queued", | |
| "status_url": f"/status/{job_id}", | |
| "result_url": f"/result/{job_id}", | |
| }, | |
| ) | |
| def job_status(job_id: str, x_api_key: Optional[str] = Header(default=None)): | |
| _require_api_key(x_api_key) | |
| job = _get_job(job_id) | |
| if not job: | |
| raise HTTPException(status_code=404, detail="Job not found") | |
| response = {"job_id": job_id, "status": job.get("status", "unknown")} | |
| if "error" in job: | |
| response["error"] = job["error"] | |
| return response | |
| def job_result( | |
| job_id: str, | |
| background_tasks: BackgroundTasks = BackgroundTasks(), | |
| x_api_key: Optional[str] = Header(default=None), | |
| ): | |
| _require_api_key(x_api_key) | |
| job = _get_job(job_id) | |
| if not job: | |
| raise HTTPException(status_code=404, detail="Job not found") | |
| if job.get("status") != "completed": | |
| raise HTTPException(status_code=409, detail=f"Job not ready (status={job.get('status')})") | |
| output_file = job.get("output_file") | |
| if not output_file or not Path(output_file).exists(): | |
| _pop_job(job_id) | |
| raise HTTPException(status_code=410, detail="Result expired or missing") | |
| _pop_job(job_id) | |
| background_tasks.add_task(_cleanup_files, output_file) | |
| return FileResponse(output_file, media_type="audio/wav", filename="output.wav") | |
| def health(x_api_key: Optional[str] = Header(default=None)): | |
| _require_api_key(x_api_key) | |
| return {"status": "ok", "model": "indextts2", "version": "2.0.0", "device": DEVICE} | |
| def root(): | |
| return {"name": "indextts2-api", "version": "2.0.0"} | |