Update app.py

app.py

@@ -1,7 +1,10 @@
 import base64
+import hashlib
 import os
 import tempfile
 import uuid
+import time
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from threading import Lock
 from typing import Dict, Optional
@@ -14,85 +17,158 @@ from fastapi import BackgroundTasks, Body, FastAPI, Header, HTTPException
 from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel, Field, HttpUrl
 
+# ----------------------------
+# Config / Tunables
+# ----------------------------
 SPACE_API_KEY = os.getenv("SPACE_API_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"
-MAX_TEXT_LENGTH = 1000
-DEFAULT_LANGUAGE = "en"
+MODEL_REPO = os.getenv("MODEL_REPO", "IndexTeam/IndexTTS-2")
+MODEL_DIR = os.getenv("MODEL_DIR", "/data/indextts2")
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+MAX_TEXT_LENGTH = int(os.getenv("MAX_TEXT_LENGTH", "1000"))
+DEFAULT_LANGUAGE = os.getenv("DEFAULT_LANGUAGE", "en")
+TARGET_SR = int(os.getenv("TARGET_SR", "16000"))  # lowered to 16 kHz for speed
+TORCH_NUM_THREADS = int(os.getenv("TORCH_NUM_THREADS", "2"))
+
+# Embedding cache settings
+EMBED_CACHE_MAX = int(os.getenv("EMBED_CACHE_MAX", "128"))  # max entries
+EMBED_CACHE_TTL = int(os.getenv("EMBED_CACHE_TTL", str(60 * 60 * 24)))  # 24h by default
+
+# Threadpool for bounded parallel jobs (keeps worker threads limited)
+WORKER_COUNT = int(os.getenv("WORKER_COUNT", "1"))  # keep low on CPU
+
+# ----------------------------
+# Torch settings
+# ----------------------------
+torch.set_num_threads(TORCH_NUM_THREADS)
+try:
+    # optional: limit interop threads
+    torch.set_num_interop_threads(max(1, TORCH_NUM_THREADS // 2))
+except Exception:
+    pass
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-JOBS: Dict[str, Dict[str, str]] = {}
-JOB_LOCK = Lock()
 
-# Set token in environment before importing
+# ----------------------------
+# Hugging Face login (if token)
+# ----------------------------
 if HF_TOKEN:
     os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
     os.environ["HF_TOKEN"] = HF_TOKEN
     try:
         from huggingface_hub import login
+
         login(token=HF_TOKEN, add_to_git_credential=False)
-    except ImportError:
+    except Exception:
         pass
 
-# Download model checkpoints from Hugging Face
-MODEL_DIR = os.getenv("MODEL_DIR", "/data/indextts2")
-os.makedirs(MODEL_DIR, exist_ok=True)
-
+# ----------------------------
+# Optionally download model snapshot (only if missing)
+# ----------------------------
 try:
     from huggingface_hub import snapshot_download
-    
-    # Download model if not already present
-    if not Path(MODEL_DIR, "config.yaml").exists():
-        print(f"Downloading IndexTTS2 model from {MODEL_REPO}...")
-        snapshot_download(
-            repo_id=MODEL_REPO,
-            local_dir=MODEL_DIR,
-            token=HF_TOKEN,
-        )
-        print("Model download complete.")
+
+    cfg_path = Path(MODEL_DIR) / "config.yaml"
+    if not cfg_path.exists():
+        print(f"Config missing; downloading model snapshot {MODEL_REPO} to {MODEL_DIR} ...")
+        snapshot_download(repo_id=MODEL_REPO, local_dir=MODEL_DIR, token=HF_TOKEN)
+        print("Download complete.")
 except Exception as exc:
-    print(f"Warning: Could not download model: {exc}")
-    # Continue anyway - model might already be present
+    print(f"Warning: snapshot_download skipped or failed: {exc}")
 
-# Initialize IndexTTS2
+# ----------------------------
+# Load IndexTTS2 model (CPU mode safe defaults)
+# ----------------------------
 try:
     from indextts.infer_v2 import IndexTTS2
-    
-    cfg_path = os.path.join(MODEL_DIR, "config.yaml")
-    if not Path(cfg_path).exists():
-        raise FileNotFoundError(f"Config file not found at {cfg_path}. Model may not be downloaded.")
-    
-    tts_model = IndexTTS2(
-        cfg_path=cfg_path,
-        model_dir=MODEL_DIR,
-        use_fp16=False,  # CPU doesn't support FP16
-        use_cuda_kernel=False,  # CPU mode
-        use_deepspeed=False,  # CPU mode
-    )
-    print("IndexTTS2 model loaded successfully.")
 except Exception as exc:
-    raise RuntimeError(f"Failed to load IndexTTS2 model: {exc}") from exc
-
-app = FastAPI(title="indextts2-api", version="1.0.0")
-
-
-class GenerateRequest(BaseModel):
-    text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)
-    speaker_wav: str = Field(..., description="HTTPS URL or base64-encoded audio")
-    language: Optional[str] = Field(DEFAULT_LANGUAGE, description="ISO code, default en")
-
-
-def _require_api_key(x_api_key: Optional[str]):
-    if not SPACE_API_KEY:
-        return
-    if x_api_key != SPACE_API_KEY:
-        raise HTTPException(status_code=401, detail="Unauthorized")
+    raise RuntimeError("indextts.infer_v2 import failed. Make sure IndexTTS2 is installed.") from exc
+
+cfg_path = os.path.join(MODEL_DIR, "config.yaml")
+if not Path(cfg_path).exists():
+    raise FileNotFoundError(f"Config file not found at {cfg_path}. Place model files in {MODEL_DIR}.")
+
+# Use CPU-safe options. If GPU becomes available, you can toggle use_fp16/use_cuda_kernel.
+tts_model = IndexTTS2(
+    cfg_path=cfg_path,
+    model_dir=MODEL_DIR,
+    use_fp16=False,  # CPU doesn't support FP16 reliably
+    use_cuda_kernel=False,
+    use_deepspeed=False,
+)
+print("IndexTTS2 loaded.")
 
+# ----------------------------
+# App + job state
+# ----------------------------
+app = FastAPI(title="indextts2-api-optimized", version="1.0.0")
+JOBS: Dict[str, Dict[str, str]] = {}
+JOB_LOCK = Lock()
 
+# Threadpool for running TTS jobs; limits concurrency to WORKER_COUNT
+EXECUTOR = ThreadPoolExecutor(max_workers=WORKER_COUNT)
+
+# ----------------------------
+# Simple LRU-like embedding cache (in-memory)
+# ----------------------------
+class _EmbedCacheEntry:
+    def __init__(self, emb_tensor: torch.Tensor):
+        self.emb = emb_tensor.detach().cpu()  # keep on CPU, detached
+        self.ts = time.time()
+
+
+EMBED_CACHE: Dict[str, _EmbedCacheEntry] = {}
+EMBED_CACHE_LOCK = Lock()
+
+
+def _evict_cache_if_needed():
+    with EMBED_CACHE_LOCK:
+        if len(EMBED_CACHE) <= EMBED_CACHE_MAX:
+            return
+        # Simple eviction: remove oldest entries
+        items = sorted(EMBED_CACHE.items(), key=lambda kv: kv[1].ts)
+        for key, _ in items[: max(1, len(items) - EMBED_CACHE_MAX)]:
+            EMBED_CACHE.pop(key, None)
+
+
+def _get_cache_key_for_file(path: str) -> str:
+    # Hash the file contents (fast enough for short audio)
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        while True:
+            chunk = f.read(8192)
+            if not chunk:
+                break
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _cache_get(key: str) -> Optional[torch.Tensor]:
+    with EMBED_CACHE_LOCK:
+        entry = EMBED_CACHE.get(key)
+        if not entry:
+            return None
+        if (time.time() - entry.ts) > EMBED_CACHE_TTL:
+            EMBED_CACHE.pop(key, None)
+            return None
+        # update timestamp for LRU-ish behavior
+        entry.ts = time.time()
+        return entry.emb.clone()
+
+
+def _cache_set(key: str, emb: torch.Tensor):
+    with EMBED_CACHE_LOCK:
+        EMBED_CACHE[key] = _EmbedCacheEntry(emb)
+    _evict_cache_if_needed()
+
+# ----------------------------
+# Utilities for audio input handling
+# ----------------------------
 def _write_temp_audio_from_url(url: HttpUrl) -> str:
     response = requests.get(url, stream=True, timeout=30)
     if response.status_code >= 400:
@@ -108,7 +184,7 @@ def _write_temp_audio_from_url(url: HttpUrl) -> str:
 def _write_temp_audio_from_base64(payload: str) -> str:
     try:
         raw = base64.b64decode(payload)
-    except Exception as exc:  # pragma: no cover
+    except Exception as exc:
         raise HTTPException(status_code=400, detail="Invalid base64 speaker_wav") from exc
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
         tmp.write(raw)
@@ -121,12 +197,10 @@ def _temp_speaker_file(speaker_wav: str) -> str:
     return _write_temp_audio_from_base64(speaker_wav)
 
 
-def _preprocess_audio_wav(path: str, target_sr: int = 24000, target_peak: float = 0.98) -> str:
+def _preprocess_audio_wav(path: str, target_sr: int = TARGET_SR, target_peak: float = 0.98) -> str:
     """
-    Light preprocessing to stabilize embeddings and output quality:
-    - convert to mono
-    - resample to target_sr
-    - peak-normalize to target_peak (avoid clipping)
+    Convert to mono, resample to target_sr, and peak-normalize.
+    Overwrites the input file.
     """
     wav, sr = torchaudio.load(path)
 
@@ -146,11 +220,64 @@ def _preprocess_audio_wav(path: str, target_sr: int = 24000, target_peak: float
         scale = min(target_peak / peak, 1.0)
         wav = wav * scale
 
-    # Overwrite input file to avoid extra temp files
+    # Overwrite file in 16-bit PCM
     torchaudio.save(path, wav, sr, bits_per_sample=16)
     return path
 
 
+# ----------------------------
+# Embedding extraction helper (tries multiple API variants)
+# ----------------------------
+def _compute_spk_embedding(speaker_path: str) -> torch.Tensor:
+    """
+    Returns a CPU tensor containing the speaker embedding.
+    Tries multiple methods to extract embedding (get_spk_emb, extract_spk_emb, etc.)
+    """
+    # Key: use hash of file contents
+    key = _get_cache_key_for_file(speaker_path)
+    cached = _cache_get(key)
+    if cached is not None:
+        return cached
+
+    # Ensure audio preprocessed (mono/resample/normalize)
+    _preprocess_audio_wav(speaker_path, target_sr=TARGET_SR)
+
+    # Try known wrapper method names (depending on IndexTTS2 version)
+    emb = None
+    try:
+        if hasattr(tts_model, "get_spk_emb"):
+            emb = tts_model.get_spk_emb(speaker_path)
+        elif hasattr(tts_model, "extract_spk_emb"):
+            emb = tts_model.extract_spk_emb(speaker_path)
+        elif hasattr(tts_model, "spk_encoder") and hasattr(tts_model.spk_encoder, "embed_utterance"):
+            # some wrappers expose internal encoders
+            wav, sr = torchaudio.load(speaker_path)
+            if wav.shape[0] > 1:
+                wav = wav.mean(dim=0, keepdim=True)
+            wav = wav.squeeze(0).numpy()  # expected shape for some encoders
+            emb = tts_model.spk_encoder.embed_utterance(wav)
+            emb = torch.from_numpy(emb)
+        else:
+            raise RuntimeError("No known speaker embedding method available on tts_model.")
+    except Exception as exc:
+        # If the model doesn't provide a direct API or something fails, fallback to infer path
+        # where infer() might internally compute embedding. In that case we return None to indicate
+        # that caller should call infer with spk_audio_prompt.
+        raise RuntimeError(f"Failed to compute speaker embedding: {exc}") from exc
+
+    # Normalize & store on CPU as float32
+    if isinstance(emb, torch.Tensor):
+        emb_cpu = emb.detach().cpu().float()
+    else:
+        emb_cpu = torch.tensor(emb, dtype=torch.float32, device="cpu")
+
+    _cache_set(key, emb_cpu)
+    return emb_cpu
+
+
+# ----------------------------
+# Job helpers
+# ----------------------------
 def _set_job(job_id: str, **kwargs):
     with JOB_LOCK:
         JOBS[job_id] = {**JOBS.get(job_id, {}), **kwargs}
@@ -168,48 +295,111 @@ def _pop_job(job_id: str) -> Optional[Dict[str, str]]:
 
 
 def _cleanup_files(*files: str):
-    """Background task to clean up temporary files after response is sent."""
     for file_path in files:
         if file_path and Path(file_path).exists():
             try:
                 Path(file_path).unlink(missing_ok=True)
             except Exception:
-                pass  # Ignore cleanup errors
+                pass
 
 
 def _run_generate_job(job_id: str, payload: Dict[str, str]):
+    """
+    Worker function that computes (or reuses) embedding and performs TTS.
+    """
     speaker_file = None
     output_file = None
     _set_job(job_id, status="processing")
     try:
+        # prepare speaker audio
         speaker_file = _temp_speaker_file(payload["speaker_wav"])
-        speaker_file = _preprocess_audio_wav(speaker_file)
-        output_file = os.path.join(tempfile.gettempdir(), f"indextts2-{uuid.uuid4()}.wav")
+        # preprocess (mono + resample + normalize)
+        speaker_file = _preprocess_audio_wav(speaker_file, target_sr=TARGET_SR)
+
+        # compute or fetch embedding (cached)
+        try:
+            spk_emb = _compute_spk_embedding(speaker_file)
+            use_spk_emb = True
+        except Exception as exc_emb:
+            # If embedding extraction fails, fall back to passing audio path to infer
+            spk_emb = None
+            use_spk_emb = False
+            print(f"Warning: embedding extraction failed, falling back to audio prompt: {exc_emb}")
 
-        tts_model.infer(
-            spk_audio_prompt=speaker_file,
-            text=payload["text"],
-            output_path=output_file,
-            use_random=False,
-            verbose=False,
-        )
-
-        output_file = _preprocess_audio_wav(output_file)
+        output_file = os.path.join(tempfile.gettempdir(), f"indextts2-{uuid.uuid4()}.wav")
 
+        # Call inference: prefer spk_emb if available.
+        infer_kwargs = {
+            "text": payload["text"],
+            "output_path": output_file,
+            "use_random": False,
+            "verbose": False,
+        }
+        # include sample_rate if supported by this wrapper
+        try:
+            infer_kwargs["sample_rate"] = TARGET_SR
+        except Exception:
+            pass
+
+        if use_spk_emb and spk_emb is not None:
+            # Use embedding path - many wrappers accept spk_emb or spk_embedding
+            try:
+                tts_model.infer(spk_emb=spk_emb, **infer_kwargs)
+            except TypeError:
+                # fallback argument name
+                tts_model.infer(speaker_emb=spk_emb, **infer_kwargs)
+        else:
+            # pass the audio file as prompt (slower, model will compute embedding internally)
+            tts_model.infer(spk_audio_prompt=speaker_file, **infer_kwargs)
+
+        # Minimal validation: ensure file created
         if not Path(output_file).exists():
-            raise RuntimeError(f"TTS generation failed: output file was not created at {output_file}")
+            raise RuntimeError(f"TTS generation failed: output file not created at {output_file}")
+
+        # Do NOT re-run heavy preprocess; only resample if the model returned a different sr (rare)
+        try:
+            out_wav, out_sr = torchaudio.load(output_file)
+            if out_sr != TARGET_SR:
+                resampler = Resample(orig_freq=out_sr, new_freq=TARGET_SR)
+                out_wav = resampler(out_wav)
+                torchaudio.save(output_file, out_wav, TARGET_SR, bits_per_sample=16)
+        except Exception:
+            # If this fails, still return the original output file
+            pass
+
+        # cleanup speaker temp (we keep output until client downloads)
+        if speaker_file:
+            try:
+                Path(speaker_file).unlink(missing_ok=True)
+            except Exception:
+                pass
 
-        _cleanup_files(speaker_file)
         _set_job(job_id, status="completed", output_file=output_file)
     except Exception as exc:
         _cleanup_files(speaker_file, output_file)
         _set_job(job_id, status="error", error=str(exc))
 
 
+# ----------------------------
+# FastAPI endpoints
+# ----------------------------
+class GenerateRequest(BaseModel):
+    text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)
+    speaker_wav: str = Field(..., description="HTTPS URL or base64-encoded audio")
+    language: Optional[str] = Field(DEFAULT_LANGUAGE, description="ISO code, default en")
+
+
+def _require_api_key(x_api_key: Optional[str]):
+    if not SPACE_API_KEY:
+        return
+    if x_api_key != SPACE_API_KEY:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
+
 @app.post("/health")
 def health(x_api_key: Optional[str] = Header(default=None)):
     _require_api_key(x_api_key)
-    return {"status": "ok", "model": "indextts2", "device": DEVICE}
+    return {"status": "ok", "model": "indextts2", "device": DEVICE, "torch_threads": torch.get_num_threads()}
 
 
 @app.post("/generate")
@@ -222,8 +412,8 @@ def generate(
     job_id = str(uuid.uuid4())
     _set_job(job_id, status="queued")
 
-    # Offload the long-running synthesis so the HTTP request stays fast (<100s)
-    background_tasks.add_task(_run_generate_job, job_id, payload.dict())
+    # Submit to bounded threadpool to avoid uncontrolled concurrency on CPU
+    EXECUTOR.submit(_run_generate_job, job_id, payload.dict())
 
     return JSONResponse(
         status_code=202,
@@ -276,7 +466,4 @@ def job_result(
 
 @app.get("/")
 def root():
-    return {
-        "name": "indextts2-api",
-        "endpoints": ["/health", "/generate", "/status/{job_id}", "/result/{job_id}"],
-    }
+    return {"name": "indextts2-api-optimized", "endpoints": ["/health", "/generate", "/status/{job_id}", "/result/{job_id}"]}