Update app.py

app.py

@@ -3,17 +3,14 @@ import io
 import random
 import numpy as np
 import torch
-import soundfile as sf
+from scipy.io import wavfile
 from fastapi import FastAPI, Form
-from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 
-from src.chatterbox.mtl_tts import (
-    ChatterboxMultilingualTTS,
-    SUPPORTED_LANGUAGES
-)
+from src.chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
 
 # ===============================
-# CPU ONLY PATCH
+# CPU-ONLY HARD PATCH
 # ===============================
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
@@ -23,57 +20,41 @@ def _cpu_only_torch_load(*args, **kwargs):
     return _original_torch_load(*args, **kwargs)
 torch.load = _cpu_only_torch_load
 
-
 # ===============================
 # LANGUAGE CONFIG
 # ===============================
 LANGUAGE_CONFIG = {
     "en": {
-        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac",
+        "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",
+        "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",
+        "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",
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/he_m1.flac"
     },
 }
 
 # ===============================
-# MODEL LOAD (CPU SAFE)
+# MODEL LOADING (SAFE)
 # ===============================
 MODEL = None
-DEVICE = "cpu"
 
 def get_or_load_model():
     global MODEL
-
     if MODEL is None:
-        print("🔄 Loading Chatterbox model (CPU-only)...")
-
+        print("🔄 Loading Chatterbox model (CPU-only)")
         MODEL = ChatterboxMultilingualTTS.from_pretrained("cpu")
-
-        # ⚠️ DO NOT call .to()
-        # ChatterboxMultilingualTTS is NOT a torch.nn.Module
-
         MODEL.eval()
-
-        # Disable grads if exposed
-        if hasattr(MODEL, "parameters"):
-            for p in MODEL.parameters():
-                p.requires_grad = False
-
-        print("✅ Model loaded successfully")
-
+        print("✅ Model loaded")
     return MODEL
 
-
+# Load at startup
 get_or_load_model()
 
-
 # ===============================
 # UTILITIES
 # ===============================
@@ -82,7 +63,6 @@ def set_seed(seed: int):
     random.seed(seed)
     np.random.seed(seed)
 
-
 def format_for_singing(lyrics: str) -> str:
     lines = []
     for line in lyrics.splitlines():
@@ -99,24 +79,31 @@ def format_for_singing(lyrics: str) -> str:
         lines.append(f"{line} ♪ ...")
     return "\n".join(lines)
 
-
 # ===============================
 # FASTAPI APP
 # ===============================
-app = FastAPI(title="Chatterbox TTS API", version="1.0")
-
+app = FastAPI(
+    title="Chatterbox Multilingual TTS",
+    version="1.0"
+)
 
+# ===============================
+# HEALTH API
+# ===============================
 @app.get("/health")
 def health():
-    return {"status": "ok", "device": DEVICE}
-
+    return {
+        "status": "ok",
+        "device": "cpu",
+        "languages": list(SUPPORTED_LANGUAGES.keys())
+    }
 
 # ===============================
 # TTS API
 # ===============================
 @app.post("/tts")
-def tts_api(
-    mode: str = Form("Speak"),
+def tts(
+    mode: str = Form("Speak"),          # Speak | Sing
     text: str = Form(""),
     lyrics: str = Form(""),
     language_id: str = Form("hi"),
@@ -130,13 +117,19 @@ def tts_api(
     if seed != 0:
         set_seed(seed)
 
-    if mode == "Sing":
+    if mode.lower() == "sing":
         if not lyrics.strip():
-            return JSONResponse({"error": "Lyrics required for Sing mode"}, status_code=400)
+            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)
+            return JSONResponse(
+                {"error": "Text required for Speak mode"},
+                status_code=400
+            )
         final_text = text
 
     kwargs = {
@@ -159,7 +152,7 @@ def tts_api(
     wav = wav.squeeze(0).cpu().numpy()
 
     buffer = io.BytesIO()
-    sf.write(buffer, wav, model.sr, format="WAV")
+    wavfile.write(buffer, model.sr, wav)
     buffer.seek(0)
 
     return StreamingResponse(
@@ -167,77 +160,3 @@ def tts_api(
         media_type="audio/wav",
         headers={"Content-Disposition": "inline; filename=output.wav"}
     )
-
-
-# ===============================
-# SIMPLE WEB UI
-# ===============================
-@app.get("/", response_class=HTMLResponse)
-def ui():
-    langs = "".join(
-        f"<option value='{k}'>{v}</option>"
-        for k, v in SUPPORTED_LANGUAGES.items()
-    )
-
-    return f"""
-<!DOCTYPE html>
-<html>
-<head>
-    <title>Chatterbox TTS</title>
-    <style>
-        body {{ font-family: Arial; max-width: 800px; margin: auto; }}
-        textarea {{ width: 100%; height: 120px; }}
-        select, button {{ padding: 6px; }}
-    </style>
-</head>
-<body>
-
-<h2>🎤 Chatterbox Multilingual TTS</h2>
-
-<label>Mode:</label>
-<select id="mode">
-    <option value="Speak">Speak</option>
-    <option value="Sing">Sing</option>
-</select>
-
-<br><br>
-
-<label>Language:</label>
-<select id="language">{langs}</select>
-
-<br><br>
-
-<label>Text (Speak):</label>
-<textarea id="text"></textarea>
-
-<label>Lyrics (Sing):</label>
-<textarea id="lyrics"></textarea>
-
-<br>
-
-<button onclick="run()">Generate</button>
-
-<br><br>
-<audio id="player" controls></audio>
-
-<script>
-async function run() {{
-    const form = new FormData();
-    form.append("mode", document.getElementById("mode").value);
-    form.append("language_id", document.getElementById("language").value);
-    form.append("text", document.getElementById("text").value);
-    form.append("lyrics", document.getElementById("lyrics").value);
-
-    const res = await fetch("/tts", {{
-        method: "POST",
-        body: form
-    }});
-
-    const blob = await res.blob();
-    document.getElementById("player").src = URL.createObjectURL(blob);
-}}
-</script>
-
-</body>
-</html>
-"""