Upload app.py with huggingface_hub

app.py

@@ -6,145 +6,143 @@ import torch
 import tempfile
 import traceback
 import gc
-from fastapi import FastAPI, Request, Response
+from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 import uvicorn
 from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 from TTS.api import TTS
 
-# --- [v140] 🚀 H200 CORE STABILIZATION (Dynamic Workspace + Explicit FP32) ---
-print(f"--- [v140] 📡 BOOTING CORE ENGINE ---")
+# --- [v141] 🚀 HYBRID STABLE ENGINE (CPU STT + GPU TTS) ---
+print(f"--- [v141] 📡 BOOTING HYBRID ENGINE ---")
 
 try:
     import spaces
-    HAS_SPACES = True
 except ImportError:
-    HAS_SPACES = False
     class spaces:
         @staticmethod
         def GPU(duration=60, f=None):
             if f is None: return lambda x: x
             return f
 
-# --- System Config ---
+# --- Global Sync ---
 os.environ["COQUI_TOS_AGREED"] = "1"
 os.environ["PYTHONWARNINGS"] = "ignore"
-# REMOVED: CUBLAS_WORKSPACE_CONFIG (Let the driver decide)
-torch.backends.cuda.matmul.allow_tf32 = True # Allow TF32 for better alignment on Hopper
-torch.backends.cudnn.allow_tf32 = True
-torch.backends.cudnn.benchmark = False # Avoid erratic kernel selection
 
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 MODELS = {"stt": None, "tts": None}
 
-def load_gpu_models():
-    """Persistent loading with explicit casting."""
+def load_cpu_stt():
+    """Loads Whisper on CPU for maximum stability."""
     global MODELS
-    device = "cuda"
-    
     if MODELS.get("stt") is None:
-        print("--- [v140] 📥 LOADING WHISPER (EXPLICIT FP32) ---")
+        print("--- [v141] 📥 LOADING WHISPER ON CPU ---")
         model_id = "openai/whisper-large-v3-turbo"
-        # Load and force cast to float()
+        # CPU loading (fp32)
         model = AutoModelForSpeechSeq2Seq.from_pretrained(
-            model_id, low_cpu_mem_usage=True, use_safetensors=True
-        ).to(device).float() # FORCE FLOAT32
-        
+            model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True
+        )
         processor = AutoProcessor.from_pretrained(model_id)
-        
         MODELS["stt"] = pipeline(
             "automatic-speech-recognition",
             model=model,
             tokenizer=processor.tokenizer,
             feature_extractor=processor.feature_extractor,
-            torch_dtype=torch.float32,
-            device=device,
-            model_kwargs={"attn_implementation": "eager"}
+            device="cpu"
         )
-        print("--- [v140] ✅ WHISPER LOADED (FORCED FP32) ---")
+        print("--- [v141] ✅ WHISPER LOADED (CPU) ---")
 
-    if MODELS.get("tts") is None:
-        print("--- [v140] 📥 LOADING XTTS ---")
-        MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
-        print("--- [v140] ✅ XTTS LOADED ---")
-
-@spaces.GPU(duration=120)
-def core_process(request_dict):
+@spaces.GPU(duration=90)
+def load_gpu_tts():
+    """Loads XTTS on GPU for maximum speed."""
     global MODELS
-    action = request_dict.get("action")
-    print(f"--- [v140] 🛠️ CORE ENGINE: {action} ---")
-    t1 = time.time()
+    if MODELS.get("tts") is None:
+        print("--- [v141] 📥 LOADING XTTS ON GPU ---")
+        MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
+        print("--- [v141] ✅ XTTS LOADED (GPU) ---")
+
+def stt_process(audio_b64, lang):
+    load_cpu_stt()
+    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, 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=90)
+def tts_process(text, target_lang, speaker_wav_b64=None):
+    load_gpu_tts()
+    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: return {"error": f"Language {clean_lang} not supported."}
     
+    speaker_wav_path = "default_speaker.wav"
+    if speaker_wav_b64:
+        sb = base64.b64decode(speaker_wav_b64)
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            f.write(sb); speaker_wav_path = f.name
+
     try:
-        load_gpu_models()
+        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)
 
+@app.post("/process")
+async def api_process(request: Request):
+    try:
+        data = await request.json()
+        action = data.get("action")
+        if action == "health": return {"status": "awake", "v": "141"}
+        
+        print(f"--- [v141] 🛠️ HYBRID ENGINE: {action} ---")
+        t1 = time.time()
+        
+        # 🎙️ STT (CPU)
+        stt_text = None
         if action in ["stt", "s2st"]:
-            audio_bytes = base64.b64decode(request_dict.get("file"))
-            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-                f.write(audio_bytes); temp_path = f.name
-            try:
-                lang = request_dict.get("lang")
-                result = MODELS["stt"](temp_path, batch_size=1, generate_kwargs={"language": lang if lang and len(lang) <= 3 else None})
-                stt_text = result["text"].strip()
-            finally:
-                if os.path.exists(temp_path): os.unlink(temp_path)
-            
+            stt_text = stt_process(data.get("file"), data.get("lang"))
             if action == "stt": return {"text": stt_text}
 
+        # 🔊 TTS (GPU)
         if action in ["tts", "s2st"]:
-            text = (request_dict.get("text") if action == "tts" else stt_text).strip()
+            text = (data.get("text") if action == "tts" else stt_text).strip()
             trans_text = text
             if action == "s2st":
                 from deep_translator import GoogleTranslator
-                target = request_dict.get("target_lang") or "en"
-                text = GoogleTranslator(source='auto', target=target).translate(stt_text)
-                trans_text = text
+                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": ""}
-
-            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"}
-            raw_lang = (request_dict.get("lang") if action == "tts" else target).strip().lower()
-            clean_lang = raw_lang.split('-')[0]
-            mapped_lang = XTTS_MAP.get(clean_lang) or ("zh-cn" if clean_lang == "zh" else None)
             
-            if mapped_lang:
-                speaker_wav_path = "default_speaker.wav" 
-                if 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 'out_p' in locals() and os.path.exists(out_p): os.unlink(out_p)
-            else: return {"error": f"Language {clean_lang} not supported."}
+            audio_res = tts_process(text, (data.get("lang") if action == "tts" else data.get("target_lang")), data.get("speaker_wav"))
+            if isinstance(audio_res, dict) and "error" in audio_res: return audio_res
             
-            if action == "tts": return {"audio": audio_b64}
-            return {"text": stt_text, "translated": trans_text, "audio": audio_b64}
+            if action == "tts": return {"audio": audio_res}
+            return {"text": stt_text, "translated": trans_text, "audio": audio_res}
 
     except Exception as e:
-        print(f"❌ [v140] ERROR: {traceback.format_exc()}")
+        print(f"❌ [v141] ERROR: {traceback.format_exc()}")
         return {"error": str(e)}
     finally:
-        print(f"--- [v140] ✨ DONE ({time.time()-t1:.1f}s) ---")
-        torch.cuda.empty_cache() 
-
-@app.post("/process")
-async def api_process(request: Request):
-    try:
-        data = await request.json()
-        if data.get("action") == "health": return {"status": "awake", "v": "140"}
-        return core_process(data)
-    except Exception as e: return {"error": str(e)}
+        print(f"--- [v141] ✨ DONE ---")
 
 @app.get("/health")
-def health(): return {"status": "ok", "v": "140", "gpu": HAS_SPACES}
+def health(): return {"status": "ok", "v": "141", "mode": "HYBRID"}
 
 @app.get("/", response_class=HTMLResponse)
-def root():
-    return f"<html><body><h1>🚀 AI Engine v140 (STABLE BASELINE)</h1><p>GPU: {HAS_SPACES}</p></body></html>"
+def root(): return "<html><body><h1>🚀 AI Engine v141 (HYBRID STABLE)</h1></body></html>"
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)