Kozzzq/indextts2api

1import base642import os3import tempfile4import uuid5from pathlib import Path6from threading import Lock7from typing import Dict, Optional8 9import requests10import torch11import torchaudio12from torchaudio.transforms import Resample13from fastapi import BackgroundTasks, Body, FastAPI, Header, HTTPException14from fastapi.responses import FileResponse, JSONResponse15from pydantic import BaseModel, Field, HttpUrl16 17# Environment configuration18SPACE_API_KEY = os.getenv("SPACE_API_KEY")19HF_TOKEN = (20    os.getenv("HUGGING_FACE_HUB_TOKEN")21    or os.getenv("HUGGINGFACEHUB_API_TOKEN")22    or os.getenv("HF_TOKEN")23)24 25# Model configuration26MODEL_REPO = "IndexTeam/IndexTTS-2"27MODEL_DIR = os.getenv("MODEL_DIR", "/data/indextts2")28MAX_TEXT_LENGTH = 100029DEFAULT_LANGUAGE = "en"30DEVICE = "cuda" if torch.cuda.is_available() else "cpu"31 32# Job management33JOBS: Dict[str, Dict[str, str]] = {}34JOB_LOCK = Lock()35 36# Set token in environment before importing37if HF_TOKEN:38    os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN39    os.environ["HF_TOKEN"] = HF_TOKEN40    try:41        from huggingface_hub import login42        login(token=HF_TOKEN, add_to_git_credential=False)43    except ImportError:44        pass45 46# Download model checkpoints from Hugging Face47os.makedirs(MODEL_DIR, exist_ok=True)48 49try:50    from huggingface_hub import snapshot_download51    52    # Download model if not already present53    if not Path(MODEL_DIR, "config.yaml").exists():54        print(f"Downloading IndexTTS2 model from {MODEL_REPO}...")55        snapshot_download(56            repo_id=MODEL_REPO,57            local_dir=MODEL_DIR,58            token=HF_TOKEN,59        )60        print("Model download complete.")61except Exception as exc:62    print(f"Warning: Could not download model: {exc}")63    # Continue anyway - model might already be present64 65# Initialize IndexTTS266try:67    from indextts.infer_v2 import IndexTTS268    69    cfg_path = os.path.join(MODEL_DIR, "config.yaml")70    if not Path(cfg_path).exists():71        raise FileNotFoundError(72            f"Config file not found at {cfg_path}. Model may not be downloaded."73        )74    75    tts_model = IndexTTS2(76        cfg_path=cfg_path,77        model_dir=MODEL_DIR,78        use_fp16=False,  # CPU doesn't support FP1679        use_cuda_kernel=False,  # CPU mode80        use_deepspeed=False,  # CPU mode81    )82    print("IndexTTS2 model loaded successfully.")83except Exception as exc:84    raise RuntimeError(f"Failed to load IndexTTS2 model: {exc}") from exc85 86# Initialize FastAPI app87app = FastAPI(title="indextts2-api", version="1.0.0")88 89 90class GenerateRequest(BaseModel):91    text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)92    speaker_wav: str = Field(..., description="HTTPS URL or base64-encoded audio")93    language: Optional[str] = Field(DEFAULT_LANGUAGE, description="ISO code, default en")94 95 96def _require_api_key(x_api_key: Optional[str]):97    """Validate API key if configured."""98    if not SPACE_API_KEY:99        return100    if x_api_key != SPACE_API_KEY:101        raise HTTPException(status_code=401, detail="Unauthorized")102 103 104def _write_temp_audio_from_url(url: HttpUrl) -> str:105    """Download audio from URL to temporary file."""106    response = requests.get(url, stream=True, timeout=30)107    if response.status_code >= 400:108        raise HTTPException(109            status_code=400,110            detail=f"Could not fetch speaker audio: {response.status_code}"111        )112    113    suffix = Path(url.path).suffix or ".wav"114    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:115        for chunk in response.iter_content(chunk_size=8192):116            if chunk:117                tmp.write(chunk)118    return tmp.name119 120 121def _write_temp_audio_from_base64(payload: str) -> str:122    """Decode base64 audio to temporary file."""123    try:124        raw = base64.b64decode(payload)125    except Exception as exc:126        raise HTTPException(127            status_code=400,128            detail="Invalid base64 speaker_wav"129        ) from exc130    131    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:132        tmp.write(raw)133    return tmp.name134 135 136def _temp_speaker_file(speaker_wav: str) -> str:137    """Handle speaker audio input from URL or base64."""138    if speaker_wav.startswith("http://") or speaker_wav.startswith("https://"):139        return _write_temp_audio_from_url(HttpUrl(speaker_wav))140    return _write_temp_audio_from_base64(speaker_wav)141 142 143def _preprocess_audio_wav(144    path: str, 145    target_sr: int = 24000, 146    target_peak: float = 0.98147) -> str:148    """149    Light preprocessing to stabilize embeddings and output quality:150    - convert to mono151    - resample to target_sr152    - peak-normalize to target_peak (avoid clipping)153    """154    wav, sr = torchaudio.load(path)155    156    # Convert to mono157    if wav.shape[0] > 1:158        wav = wav.mean(dim=0, keepdim=True)159    160    # Resample if needed161    if sr != target_sr:162        resampler = Resample(orig_freq=sr, new_freq=target_sr)163        wav = resampler(wav)164        sr = target_sr165    166    # Peak normalize167    peak = wav.abs().max().item() if wav.numel() else 0.0168    if peak > 0:169        scale = min(target_peak / peak, 1.0)170        wav = wav * scale171    172    # Overwrite input file to avoid extra temp files173    torchaudio.save(path, wav, sr, bits_per_sample=16)174    return path175 176 177def _set_job(job_id: str, **kwargs):178    """Thread-safe job update."""179    with JOB_LOCK:180        JOBS[job_id] = {**JOBS.get(job_id, {}), **kwargs}181 182 183def _get_job(job_id: str) -> Optional[Dict[str, str]]:184    """Thread-safe job retrieval."""185    with JOB_LOCK:186        data = JOBS.get(job_id)187        return dict(data) if data else None188 189 190def _pop_job(job_id: str) -> Optional[Dict[str, str]]:191    """Thread-safe job removal."""192    with JOB_LOCK:193        return JOBS.pop(job_id, None)194 195 196def _cleanup_files(*files: str):197    """Background task to clean up temporary files after response is sent."""198    for file_path in files:199        if file_path and Path(file_path).exists():200            try:201                Path(file_path).unlink(missing_ok=True)202            except Exception:203                pass  # Ignore cleanup errors204 205 206def _run_generate_job(job_id: str, payload: Dict[str, str]):207    """Background job for TTS generation."""208    speaker_file = None209    output_file = None210    _set_job(job_id, status="processing")211    212    try:213        speaker_file = _temp_speaker_file(payload["speaker_wav"])214        speaker_file = _preprocess_audio_wav(speaker_file)215        output_file = os.path.join(216            tempfile.gettempdir(), 217            f"indextts2-{uuid.uuid4()}.wav"218        )219        emot_happy = float(payload.get("emot_happy", 0.0))220        emot_angry = float(payload.get("emot_angry", 0.0))221        emot_sad = float(payload.get("emot_sad", 0.0))222        emot_scared = float(payload.get("emot_scared", 0.0))223        emot_raged = float(payload.get("emot_raged", 0.0))224        emot_melancholic = float(payload.get("emot_melancholic", 0.0))225        emot_surprised = float(payload.get("emot_surprised", 0.0))226        emot_calm = float(payload.get("emot_calm", 0.0))227        tts_model.infer(228            spk_audio_prompt=speaker_file,229            text=payload["text"],230            output_path=output_file,231            emo_vector=[emot_happy, emot_angry, emot_sad, emot_scared, emot_raged, emot_melancholic, emot_surprised, emot_calm],232            use_random=False,233            emo_alpha=1.0,234            verbose=False,235        )236        237        output_file = _preprocess_audio_wav(output_file)238        239        if not Path(output_file).exists():240            raise RuntimeError(241                f"TTS generation failed: output file was not created at {output_file}"242            )243        244        _cleanup_files(speaker_file)245        _set_job(job_id, status="completed", output_file=output_file)246    except Exception as exc:247        _cleanup_files(speaker_file, output_file)248        _set_job(job_id, status="error", error=str(exc))249 250 251@app.post("/health")252def health(x_api_key: Optional[str] = Header(default=None)):253    """Health check endpoint."""254    _require_api_key(x_api_key)255    return {"status": "ok", "model": "indextts2", "device": DEVICE}256 257 258@app.post("/generate")259def generate(260    payload: GenerateRequest = Body(...),261    background_tasks: BackgroundTasks = BackgroundTasks(),262    x_api_key: Optional[str] = Header(default=None),263):264    """265    Generate speech from text using voice cloning.266    Returns job information for async processing.267    """268    _require_api_key(x_api_key)269    270    job_id = str(uuid.uuid4())271    _set_job(job_id, status="queued")272    273    # Offload the long-running synthesis so the HTTP request stays fast (<100s)274    background_tasks.add_task(_run_generate_job, job_id, payload.dict())275    276    return JSONResponse(277        status_code=202,278        content={279            "job_id": job_id,280            "status": "queued",281            "status_url": f"/status/{job_id}",282            "result_url": f"/result/{job_id}",283        },284    )285 286 287@app.get("/status/{job_id}")288def job_status(job_id: str, x_api_key: Optional[str] = Header(default=None)):289    """Check the status of a generation job."""290    _require_api_key(x_api_key)291    292    job = _get_job(job_id)293    if not job:294        raise HTTPException(status_code=404, detail="Job not found")295    296    payload: Dict[str, str] = {297        "job_id": job_id,298        "status": job.get("status", "unknown")299    }300    301    if "error" in job:302        payload["error"] = job["error"]303    304    return payload305 306 307@app.get("/result/{job_id}")308def job_result(309    job_id: str,310    background_tasks: BackgroundTasks = BackgroundTasks(),311    x_api_key: Optional[str] = Header(default=None),312):313    """Retrieve the result of a completed generation job."""314    _require_api_key(x_api_key)315    316    job = _get_job(job_id)317    if not job:318        raise HTTPException(status_code=404, detail="Job not found")319    320    status = job.get("status")321    if status != "completed":322        raise HTTPException(323            status_code=409, 324            detail=f"Job not ready (status={status})"325        )326    327    output_file = job.get("output_file")328    if not output_file or not Path(output_file).exists():329        _pop_job(job_id)330        raise HTTPException(status_code=410, detail="Result expired or missing")331    332    # Remove job from memory and cleanup output after sending333    _pop_job(job_id)334    background_tasks.add_task(_cleanup_files, output_file)335    336    return FileResponse(337        output_file, 338        media_type="audio/wav", 339        filename="output.wav"340    )341 342 343@app.get("/")344def root():345    """API root with available endpoints."""346    return {347        "name": "indextts2-api",348        "endpoints": [349            "/health", 350            "/generate", 351            "/status/{job_id}", 352            "/result/{job_id}"353        ],354    }355