Upload app.py with huggingface_hub

app.py

@@ -5,13 +5,14 @@ import base64
 import torch
 import tempfile
 import traceback
+import gc
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 import uvicorn
 
-# --- [v151] 🚀 TTS DEBUG ENGINE ---
-print(f"--- [v151] 📡 BOOTING DEBUG ENGINE ---")
+# --- [v152] 🚀 REFINED HYBRID ENGINE (CPU-STT + CPU-TRANS + GPU-TTS) ---
+print(f"--- [v152] 📡 BOOTING REFINED HYBRID ENGINE ---")
 
 from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 from TTS.api import TTS
@@ -43,128 +44,132 @@ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], all
 
 MODELS = {"stt": None, "tts": None}
 
-def load_stt_gpu():
+def load_stt_cpu():
+    """STT on CPU is stable and fast for Whisper Base."""
     global MODELS
     if MODELS.get("stt") is None:
-        print("--- [v151] 📥 LOADING WHISPER (Base) ---")
+        print("--- [v152] 📥 LOADING WHISPER (Base) ON CPU ---")
         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")
-        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="cuda",
-            model_kwargs={"attn_implementation": "eager"}
+            model=model_id,
+            device="cpu", # Force CPU for stability
+            torch_dtype=torch.float32
         )
-        print("--- [v151] ✅ WHISPER READY ---")
+        print("--- [v152] ✅ WHISPER READY (CPU) ---")
 
 def load_tts_gpu():
     global MODELS
     if MODELS.get("tts") is None:
-        print("--- [v151] 📥 LOADING XTTS V2 ---")
-        try:
-            # Try loading once and keeping in VRAM if possible (ZeroGPU might clear it)
-            MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
-            MODELS["tts"].to(torch.float32)
-            print("--- [v151] ✅ XTTS LOADED SUCCESSFULLY ---")
-        except Exception as e:
-            print(f"--- [v151] ❌ XTTS FAILED TO LOAD: {e} ---")
-            raise e
-
-@spaces.GPU(duration=180) # Longer duration for XTTS
-def process_full(action, data):
+        print("--- [v152] 📥 LOADING XTTS V2 ---")
+        # Load to CPU first, then move to CUDA inside the decorated function if needed, 
+        # or just load directly if ZeroGPU allows.
+        MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
+        print("--- [v152] ✅ XTTS READY (CPU MEMORY) ---")
+
+@spaces.GPU(duration=120)
+def gpu_tts_inference(text, mapped_lang, speaker_path):
+    """Isolated GPU inference for XTTS."""
     global MODELS
-    print(f"--- [v151] 🚀 STARTING {action} on GPU ---")
+    if MODELS["tts"] is None:
+        load_tts_gpu()
+    
+    # Move to GPU inside the decorated scope
+    MODELS["tts"].to("cuda")
+    
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as out_f:
+        out_p = out_f.name
     
-    # 🎙️ STT
-    stt_text = ""
-    if action in ["stt", "s2st"]:
-        load_stt_gpu()
-        try:
+    try:
+        MODELS["tts"].tts_to_file(text=text, language=mapped_lang, file_path=out_p, speaker_wav=speaker_path)
+        with open(out_p, "rb") as f:
+            audio_b64 = base64.b64encode(f.read()).decode()
+        return audio_b64
+    finally:
+        # Move back to CPU to release VRAM as per ZeroGPU guidelines
+        MODELS["tts"].to("cpu")
+        if os.path.exists(out_p): os.unlink(out_p)
+        torch.cuda.empty_cache()
+
+async def handle_process(request: Request):
+    try:
+        data = await request.json()
+        action = data.get("action")
+        if action == "health": return {"status": "awake", "v": "152"}
+        
+        print(f"--- [v152] 🛠️ ENGINE ACTION: {action} ---")
+        t1 = time.time()
+        
+        # 🎙️ STT (CPU)
+        stt_text = None
+        if action in ["stt", "s2st"]:
+            load_stt_cpu()
             audio_bytes = base64.b64decode(data.get("file"))
             with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                 f.write(audio_bytes); temp_path = f.name
-            lang = data.get("lang")
-            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()
-            print(f"--- [v151] 🎙️ STT: {stt_text[:50]}... ---")
-            if action == "stt": return {"text": stt_text}
-        finally:
-            if 'temp_path' in locals() and os.path.exists(temp_path): os.unlink(temp_path)
-
-    # 🔊 TTS
-    if action in ["tts", "s2st"]:
-        text = (data.get("text") if action == "tts" else stt_text).strip()
-        trans_text = text
-        target = data.get("target_lang") or data.get("lang") or "en"
-        
-        if action == "s2st":
-            print(f"--- [v151] 🌏 TRANSLATING TO {target}... ---")
-            trans_text = GoogleTranslator(source='auto', target=target).translate(stt_text)
-            text = trans_text
-            print(f"--- [v151] 📝 TRANS: {trans_text[:50]}... ---")
-
-        if len(text) < 2: return {"text": stt_text, "translated": "", "audio": ""} if action == "s2st" else {"audio": ""}
+            try:
+                lang = data.get("lang")
+                res = MODELS["stt"](temp_path, generate_kwargs={"language": lang if lang and len(lang) <= 3 else None})
+                stt_text = res["text"].strip()
+                if action == "stt": return {"text": stt_text}
+            finally:
+                if os.path.exists(temp_path): os.unlink(temp_path)
+
+        # 🔊 TTS (Hybrid GPU)
+        if action in ["tts", "s2st"]:
+            text = (data.get("text") if action == "tts" else stt_text).strip()
+            trans_text = text
+            target = data.get("target_lang") or data.get("lang") or "en"
+            
+            if action == "s2st":
+                trans_text = GoogleTranslator(source='auto', target=target).translate(stt_text)
+                text = trans_text
+
+            if len(text) < 2: return {"text": stt_text, "translated": "", "audio": ""} if action == "s2st" else {"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"}
+            clean_lang = target.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:
+                    audio_bytes = chatterbox_utils.run_chatterbox_inference(text, clean_lang)
+                    audio_b64 = base64.b64encode(audio_bytes).decode()
+                else: return {"error": f"Lang {clean_lang} unsupported"}
+            else:
+                speaker_wav_b64 = data.get("speaker_wav")
+                speaker_path = None
+                if speaker_wav_b64:
+                    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+                        f.write(base64.b64decode(speaker_wav_b64)); speaker_path = f.name
+                else:
+                    speaker_path = "default_speaker.wav"
+                    if not os.path.exists(speaker_path): speaker_path = None
+                
+                try:
+                    audio_b64 = gpu_tts_inference(text, mapped_lang, speaker_path)
+                finally:
+                    if speaker_wav_b64 and speaker_path and os.path.exists(speaker_path): os.unlink(speaker_path)
+            
+            if action == "tts": return {"audio": audio_b64}
+            return {"text": stt_text, "translated": trans_text, "audio": audio_b64}
 
-        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.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"--- [v151] 📦 FALLBACK: CHATTERBOX ---")
-                audio_bytes = chatterbox_utils.run_chatterbox_inference(text, clean_lang)
-                return {"text": stt_text, "translated": trans_text, "audio": base64.b64encode(audio_bytes).decode()}
-            return {"error": f"Lang {clean_lang} unsupported"}
-
-        print(f"--- [v151] 📥 LOADING XTTS... ---")
-        load_tts_gpu()
-        
-        speaker_wav = data.get("speaker_wav")
-        speaker_path = None
-        if speaker_wav:
-            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-                f.write(base64.b64decode(speaker_wav)); speaker_path = f.name
-        else:
-            speaker_path = "default_speaker.wav"
-            if not os.path.exists(speaker_path): speaker_path = None
-
-        print(f"--- [v151] 🔊 RUNNING XTTS INFERENCE... ---")
-        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_path)
-            with open(out_p, "rb") as f: audio_b64 = base64.b64encode(f.read()).decode()
-            print(f"--- [v151] ✅ TTS SUCCESS! ---")
-        finally:
-            if speaker_wav and speaker_path and os.path.exists(speaker_path): os.unlink(speaker_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}
+    except Exception as e:
+        print(f"❌ [v152] ERROR: {traceback.format_exc()}")
+        return {"error": str(e)}
+    finally:
+        print(f"--- [v152] ✨ DONE ({time.time()-t1:.1f}s) ---")
 
 @app.post("/process")
 @app.post("/api/v1/process")
-async def api_process(request: Request):
-    try:
-        data = await request.json()
-        action = data.get("action")
-        if action == "health": return {"status": "awake", "v": "151"}
-        return process_full(action, data)
-    except Exception as e:
-        print(f"❌ [v151] CRASH: {traceback.format_exc()}")
-        return {"error": str(e)}
+async def api_process(request: Request): return await handle_process(request)
 
 @app.get("/health")
 def health():
-    return {"status": "ok", "v": "151", "gpu": torch.cuda.is_available(), "chatterbox": HAS_CHATTERBOX}
+    return {"status": "ok", "v": "152", "gpu": torch.cuda.is_available()}
 
 @app.get("/", response_class=HTMLResponse)
-def root(): return "<html><body><h1>🚀 AI Engine v151 (DEBUG)</h1></body></html>"
+def root(): return "<h1>🚀 AI Engine v152 (REFINED HYBRID)</h1>"
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)