Update app.py

app.py

@@ -1,162 +1,99 @@
 import os
-import io
-import random
-import numpy as np
+import uuid
 import torch
-from scipy.io import wavfile
-from fastapi import FastAPI, Form
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi import FastAPI
+from pydantic import BaseModel
+from fastapi.responses import FileResponse, HTMLResponse
 
 from src.chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
 
-# ===============================
-# CPU-ONLY HARD PATCH
-# ===============================
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
+# -------------------------------------------------
+# App
+# -------------------------------------------------
+app = FastAPI(title="Chatterbox Multilingual TTS")
 
-_original_torch_load = torch.load
-def _cpu_only_torch_load(*args, **kwargs):
-    kwargs.setdefault("map_location", torch.device("cpu"))
-    return _original_torch_load(*args, **kwargs)
-torch.load = _cpu_only_torch_load
+# -------------------------------------------------
+# Globals (model loaded once)
+# -------------------------------------------------
+MODEL = None
+OUTPUT_DIR = "/tmp/tts_outputs"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# ===============================
-# LANGUAGE CONFIG
-# ===============================
-LANGUAGE_CONFIG = {
-    "en": {
-        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac"
-    },
-    "hi": {
-        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/hi_f1.flac"
-    },
-    "fr": {
-        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fr_f1.flac"
-    },
-    "he": {
-        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/he_m1.flac"
-    },
-}
+# -------------------------------------------------
+# Request schema
+# -------------------------------------------------
+class TTSRequest(BaseModel):
+    text: str
+    language: str = "en"   # "en" or "hi"
+    speaker: str | None = None
 
-# ===============================
-# MODEL LOADING (SAFE)
-# ===============================
-MODEL = None
 
+# -------------------------------------------------
+# Model loader (NO .eval())
+# -------------------------------------------------
 def get_or_load_model():
     global MODEL
     if MODEL is None:
-        print("🔄 Loading Chatterbox model (CPU-only)")
-        MODEL = ChatterboxMultilingualTTS.from_pretrained("cpu")
-        MODEL.eval()
-        print("✅ Model loaded")
-    return MODEL
-
-# Load at startup
-get_or_load_model()
-
-# ===============================
-# UTILITIES
-# ===============================
-def set_seed(seed: int):
-    torch.manual_seed(seed)
-    random.seed(seed)
-    np.random.seed(seed)
-
-def format_for_singing(lyrics: str) -> str:
-    lines = []
-    for line in lyrics.splitlines():
-        line = line.strip()
-        if not line:
-            continue
-        line = (
-            line.replace("a", "aa")
-                .replace("e", "ee")
-                .replace("i", "ii")
-                .replace("o", "oo")
-                .replace("u", "uu")
+        MODEL = ChatterboxMultilingualTTS(
+            device="cuda" if torch.cuda.is_available() else "cpu"
         )
-        lines.append(f"{line} ♪ ...")
-    return "\n".join(lines)
-
-# ===============================
-# FASTAPI APP
-# ===============================
-app = FastAPI(
-    title="Chatterbox Multilingual TTS",
-    version="1.0"
-)
+    return MODEL
 
-# ===============================
-# HEALTH API
-# ===============================
-@app.get("/health")
-def health():
-    return {
-        "status": "ok",
-        "device": "cpu",
-        "languages": list(SUPPORTED_LANGUAGES.keys())
-    }
 
-# ===============================
-# TTS API
-# ===============================
+# -------------------------------------------------
+# API: TTS
+# -------------------------------------------------
 @app.post("/tts")
-def tts(
-    mode: str = Form("Speak"),          # Speak | Sing
-    text: str = Form(""),
-    lyrics: str = Form(""),
-    language_id: str = Form("hi"),
-    exaggeration: float = Form(0.5),
-    temperature: float = Form(0.8),
-    cfg_weight: float = Form(0.5),
-    seed: int = Form(0),
-):
-    model = get_or_load_model()
+def tts(req: TTSRequest):
+    if req.language not in SUPPORTED_LANGUAGES:
+        return {
+            "error": f"Unsupported language. Supported: {SUPPORTED_LANGUAGES}"
+        }
 
-    if seed != 0:
-        set_seed(seed)
-
-    if mode.lower() == "sing":
-        if not lyrics.strip():
-            return JSONResponse(
-                {"error": "Lyrics required for Sing mode"},
-                status_code=400
-            )
-        final_text = format_for_singing(lyrics)
-    else:
-        if not text.strip():
-            return JSONResponse(
-                {"error": "Text required for Speak mode"},
-                status_code=400
-            )
-        final_text = text
-
-    kwargs = {
-        "exaggeration": exaggeration,
-        "temperature": temperature,
-        "cfg_weight": cfg_weight,
-    }
-
-    prompt = LANGUAGE_CONFIG.get(language_id, {}).get("audio")
-    if prompt:
-        kwargs["audio_prompt_path"] = prompt
-
-    with torch.no_grad():
-        wav = model.generate(
-            final_text[:300],
-            language_id=language_id,
-            **kwargs
+    model = get_or_load_model()
+    out_path = os.path.join(OUTPUT_DIR, f"{uuid.uuid4().hex}.wav")
+
+    # ✅ Correct inference pattern
+    with torch.inference_mode():
+        audio = model.tts(
+            text=req.text,
+            language=req.language,
+            speaker=req.speaker,
+            output_path=out_path,
         )
 
-    wav = wav.squeeze(0).cpu().numpy()
-
-    buffer = io.BytesIO()
-    wavfile.write(buffer, model.sr, wav)
-    buffer.seek(0)
-
-    return StreamingResponse(
-        buffer,
+    return FileResponse(
+        out_path,
         media_type="audio/wav",
-        headers={"Content-Disposition": "inline; filename=output.wav"}
+        filename="speech.wav",
     )
+
+
+# -------------------------------------------------
+# Simple UI (for quick testing)
+# -------------------------------------------------
+@app.get("/", response_class=HTMLResponse)
+def ui():
+    return """
+    <html>
+        <body>
+            <h2>Chatterbox Multilingual TTS</h2>
+            <form action="/tts" method="post">
+                <textarea name="text" rows="4" cols="60">Hello, how are you?</textarea><br><br>
+                <select name="language">
+                    <option value="en">English</option>
+                    <option value="hi">Hindi</option>
+                </select><br><br>
+                <button type="submit">Generate Speech</button>
+            </form>
+        </body>
+    </html>
+    """
+
+
+# -------------------------------------------------
+# Warm-up (optional, safe)
+# -------------------------------------------------
+@app.on_event("startup")
+def warmup():
+    get_or_load_model()