Deploy v150 High Stability

app.py

@@ -10,9 +10,8 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 import uvicorn
 
-# --- [v149] 🚀 FULL FEATURE H200 ENGINE (STT + TRANS + TTS + FALLBACK) ---
-# Target: https://tgpro1-s2st.hf.space/api/v1/process
-print(f"--- [v149] 📡 BOOTING FULL ENGINE (API v1) ---")
+# --- [v150] 🚀 HIGH STABILITY H200 ENGINE ---
+print(f"--- [v150] 📡 BOOTING ENGINE ---")
 
 from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 from TTS.api import TTS
@@ -38,7 +37,6 @@ except ImportError:
 
 os.environ["COQUI_TOS_AGREED"] = "1"
 os.environ["PYTHONWARNINGS"] = "ignore"
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
@@ -48,8 +46,9 @@ MODELS = {"stt": None, "tts": None}
 def load_stt_gpu():
     global MODELS
     if MODELS.get("stt") is None:
-        print("--- [v149] 📥 LOADING WHISPER (Large-v3-Turbo) ---")
-        model_id = "openai/whisper-large-v3-turbo"
+        print("--- [v150] 📥 LOADING WHISPER (Base for Stability) ---")
+        # Use Base model to ensure it fits and loads quickly in serverless context
+        model_id = "openai/whisper-base"
         model = AutoModelForSpeechSeq2Seq.from_pretrained(
             model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True, use_safetensors=True
         ).to("cuda")
@@ -63,107 +62,113 @@ def load_stt_gpu():
             device="cuda",
             model_kwargs={"attn_implementation": "eager"}
         )
-        print("--- [v149] ✅ WHISPER READY ---")
+        print("--- [v150] ✅ WHISPER READY ---")
 
 def load_tts_gpu():
     global MODELS
     if MODELS.get("tts") is None:
-        print("--- [v149] 📥 LOADING XTTS v2 ---")
+        print("--- [v150] 📥 LOADING XTTS v2 ---")
         MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
         MODELS["tts"].to(torch.float32) 
-        print("--- [v149] ✅ XTTS READY ---")
+        print("--- [v150] ✅ XTTS READY ---")
 
 @spaces.GPU(duration=120)
-def core_stt(audio_b64, lang):
-    load_stt_gpu()
-    audio_bytes = base64.b64decode(audio_b64)
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-        f.write(audio_bytes); temp_path = f.name
-    try:
-        result = MODELS["stt"](temp_path, batch_size=1, generate_kwargs={"language": lang if lang and len(lang) <= 3 else None})
-        return result["text"].strip()
-    finally:
-        if os.path.exists(temp_path): os.unlink(temp_path)
-
-@spaces.GPU(duration=120)
-def core_tts(text, target_lang, speaker_wav_b64=None):
-    load_tts_gpu()
-    XTTS_MAP = {"en": "en", "de": "de", "fr": "fr", "es": "es", "it": "it", "pl": "pl", "pt": "pt", "tr": "tr", "ru": "ru", "nl": "nl", "cs": "cs", "ar": "ar", "hu": "hu", "ko": "ko", "hi": "hi", "zh": "zh-cn"}
-    clean_lang = target_lang.split('-')[0].lower()
-    mapped_lang = XTTS_MAP.get(clean_lang) or ("zh-cn" if clean_lang == "zh" else None)
-    
-    if not mapped_lang:
-        if HAS_CHATTERBOX:
-            print(f"--- [v149] 📦 FALLBACK: CHATTERBOX FOR {clean_lang} ---")
-            audio_bytes = chatterbox_utils.run_chatterbox_inference(text, clean_lang)
-            return base64.b64encode(audio_bytes).decode()
-        return {"error": f"Language {clean_lang} not supported."}
+def unified_gpu_process(action, data):
+    """Single GPU entry point to avoid context flip-flops."""
+    global MODELS
+    res = {}
     
-    speaker_wav_path = "default_speaker.wav"
-    if speaker_wav_b64:
-        sb = base64.b64decode(speaker_wav_b64)
+    # 🎙️ STT
+    stt_text = ""
+    if action in ["stt", "s2st"]:
+        load_stt_gpu()
+        audio_bytes = base64.b64decode(data.get("file"))
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            f.write(sb); speaker_wav_path = f.name
-    elif not os.path.exists(speaker_wav_path):
-        speaker_wav_path = None
+            f.write(audio_bytes); temp_path = f.name
+        try:
+            lang = data.get("lang")
+            # Batch size 1 for stability
+            stt_res = MODELS["stt"](temp_path, batch_size=1, generate_kwargs={"language": lang if lang and len(lang) <= 3 else None})
+            stt_text = stt_res["text"].strip()
+            if action == "stt": return {"text": stt_text}
+        finally:
+            if os.path.exists(temp_path): os.unlink(temp_path)
+
+    # 🔊 TTS
+    if action in ["tts", "s2st"]:
+        load_tts_gpu()
+        text = (data.get("text") if action == "tts" else stt_text).strip()
+        trans_text = text
+        
+        target_lang = data.get("target_lang") or data.get("lang") or "en"
+        
+        if action == "s2st":
+            trans_text = GoogleTranslator(source='auto', target=target_lang).translate(stt_text)
+            text = trans_text
 
-    try:
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as out_f: out_p = out_f.name
-        MODELS["tts"].tts_to_file(text=text, language=mapped_lang, file_path=out_p, speaker_wav=speaker_wav_path)
-        with open(out_p, "rb") as f: audio_b64 = base64.b64encode(f.read()).decode()
-        return audio_b64
-    finally:
-        if speaker_wav_b64 and os.path.exists(speaker_wav_path): os.unlink(speaker_wav_path)
-        if 'out_p' in locals() and os.path.exists(out_p): os.unlink(out_p)
+        if len(text) < 2: return {"text": stt_text, "translated": "", "audio": ""} if action == "s2st" else {"audio": ""}
 
-async def handle_process(request: Request):
+        XTTS_MAP = {"en": "en", "de": "de", "fr": "fr", "es": "es", "it": "it", "pl": "pl", "pt": "pt", "tr": "tr", "ru": "ru", "nl": "nl", "cs": "cs", "ar": "ar", "hu": "hu", "ko": "ko", "hi": "hi", "zh": "zh-cn"}
+        clean_lang = target_lang.split('-')[0].lower()
+        mapped_lang = XTTS_MAP.get(clean_lang) or ("zh-cn" if clean_lang == "zh" else None)
+        
+        if not mapped_lang:
+            if HAS_CHATTERBOX:
+                print(f"--- [v150] 📦 FALLBACK: CHATTERBOX FOR {clean_lang} ---")
+                audio_bytes = chatterbox_utils.run_chatterbox_inference(text, clean_lang)
+                audio_b64 = base64.b64encode(audio_bytes).decode()
+            else:
+                return {"error": f"Language {clean_lang} not supported."}
+        else:
+            speaker_wav_path = "default_speaker.wav"
+            if data.get("speaker_wav"):
+                sb = base64.b64decode(data.get("speaker_wav"))
+                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+                    f.write(sb); speaker_wav_path = f.name
+            elif not os.path.exists(speaker_wav_path):
+                speaker_wav_path = None
+
+            try:
+                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as out_f: out_p = out_f.name
+                MODELS["tts"].tts_to_file(text=text, language=mapped_lang, file_path=out_p, speaker_wav=speaker_wav_path)
+                with open(out_p, "rb") as f: audio_b64 = base64.b64encode(f.read()).decode()
+            finally:
+                if data.get("speaker_wav") and os.path.exists(speaker_wav_path): os.unlink(speaker_wav_path)
+                if 'out_p' in locals() and os.path.exists(out_p): os.unlink(out_p)
+
+        if action == "tts": return {"audio": audio_b64}
+        return {"text": stt_text, "translated": trans_text, "audio": audio_b64}
+    
+    return {"error": "Invalid action"}
+
+@app.post("/process")
+@app.post("/api/v1/process")
+async def api_gateway(request: Request):
     try:
         data = await request.json()
         action = data.get("action")
-        if action == "health": return {"status": "awake", "v": "149", "mode": "FULL_API_V1"}
-        
-        print(f"--- [v149] 🛠️ ACTION: {action} ---")
-        t1 = time.time()
-        
-        stt_text = None
-        if action in ["stt", "s2st"]:
-            stt_text = core_stt(data.get("file"), data.get("lang"))
-            if action == "stt": return {"text": stt_text}
-
-        if action in ["tts", "s2st"]:
-            text = (data.get("text") if action == "tts" else stt_text).strip()
-            trans_text = text
-            if action == "s2st":
-                target = data.get("target_lang") or "en"
-                trans_text = GoogleTranslator(source='auto', target=target).translate(stt_text)
-                text = trans_text
-
-            if len(text) < 2: return {"audio": ""} if action == "tts" else {"text": stt_text, "translated": "", "audio": ""}
-            
-            audio_res = core_tts(text, (data.get("lang") if action == "tts" else target), data.get("speaker_wav"))
-            
-            if isinstance(audio_res, dict) and "error" in audio_res: return audio_res
-            if action == "tts": return {"audio": audio_res}
-            return {"text": stt_text, "translated": trans_text, "audio": audio_res}
-
+        if action == "health": return {"status": "awake", "v": "150"}
+        print(f"--- [v150] 🛠️ {action} requested ---")
+        return unified_gpu_process(action, data)
     except Exception as e:
-        print(f"❌ [v149] ERROR: {traceback.format_exc()}")
+        print(f"❌ [v150] ERROR: {traceback.format_exc()}")
         return {"error": str(e)}
     finally:
-        print(f"--- [v149] ✨ DONE ({time.time()-t1:.1f}s) ---")
         torch.cuda.empty_cache()
 
-@app.post("/process")
-async def api_process(request: Request): return await handle_process(request)
-
-@app.post("/api/v1/process")
-async def api_v1_process(request: Request): return await handle_process(request)
-
 @app.get("/health")
-def health(): return {"status": "ok", "v": "149", "mode": "FULL_H200", "gpu": HAS_SPACES, "chatterbox": HAS_CHATTERBOX}
+def health():
+    diag = {
+        "status": "ok", "v": "150", 
+        "gpu": torch.cuda.is_available(),
+        "device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "None",
+        "spaces": HAS_SPACES,
+        "chatterbox": HAS_CHATTERBOX
+    }
+    return diag
 
 @app.get("/", response_class=HTMLResponse)
-def root(): return "<html><body><h1>🚀 AI Engine v149 (FULL API V1 READY)</h1></body></html>"
+def root(): return "<html><body><h1>🚀 AI Engine v150 (HIGH STABILITY)</h1></body></html>"
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)