Upload app.py with huggingface_hub

app.py

@@ -11,8 +11,9 @@ from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 
-# --- [v161] 🚀 ULTIMATE GPU ENGINE (STABILITY MISSION) ---
-print(f"--- [v161] 📡 BOOTING ULTIMATE ENGINE ---")
+# --- [v162] 🚀 PRO GPU ENGINE (SPEED & PRECISION) ---
+# This version enforces GPU usage for STT and TTS to meet high-performance requirements.
+print(f"--- [v162] 📡 BOOTING ENGINE ---")
 
 # 🛠️ CRITICAL: TORCHAUDIO MONKEYPATCH 🛠️
 import torchaudio
@@ -26,13 +27,13 @@ def HeroLoad(filepath, **kwargs):
             data = data.T
         return torch.from_numpy(data).float(), samplerate
     except Exception as e:
-        print(f"--- [v161] ❌ PATCHED LOAD FAILED: {e} ---")
+        print(f"--- [v162] ❌ PATCHED LOAD FAILED: {e} ---")
         return torchaudio.load_orig(filepath, **kwargs)
 
 if not hasattr(torchaudio, 'load_orig'):
     torchaudio.load_orig = torchaudio.load
     torchaudio.load = HeroLoad
-    print("--- [v161] 🩹 TORCHAUDIO HERO PATCH APPLIED ---")
+    print("--- [v162] 🩹 TORCHAUDIO PATCH APPLIED ---")
 
 from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 from TTS.api import TTS
@@ -62,62 +63,61 @@ os.environ["PYTHONWARNINGS"] = "ignore"
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-MODELS = {"stt": None, "tts": None, "current_gpu": "auto"}
+MODELS = {"stt": None, "tts": None, "gpu_id": 0}
 
-def select_best_gpu():
-    """Logic to select the best available GPU (Switch)."""
-    if torch.cuda.is_available():
-        count = torch.cuda.device_count()
-        print(f"--- [v161] 🖥️ DETECTED {count} GPUs ---")
-        # In ZeroGPU, we usually have 1 MIG instance assigned.
-        return "cuda:0"
-    return "cpu"
+def get_best_gpu():
+    """Architecture for multi-GPU support (Switch)."""
+    if not torch.cuda.is_available(): return "cpu"
+    # Select GPU with most free memory if multiple exist
+    # For ZeroGPU, this defaults to the allocated MIG instance.
+    return f"cuda:{MODELS['gpu_id']}"
 
 @spaces.GPU(duration=120)
-def gpu_stt_large(temp_path, lang):
+def gpu_stt_full(temp_path, lang):
     global MODELS
-    device = select_best_gpu() if MODELS["current_gpu"] == "auto" else MODELS["current_gpu"]
+    device = get_best_gpu()
     
     if MODELS.get("stt") is None:
-        print(f"--- [v161] 📥 LOADING WHISPER (Large-v3-Turbo) ON {device} ---")
+        print(f"--- [v162] 📥 LOADING WHISPER LARGE ON {device} ---")
         model_id = "openai/whisper-large-v3-turbo"
-        model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch.float16 if "cuda" in device else torch.float32).to(device)
+        model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
         processor = AutoProcessor.from_pretrained(model_id)
         MODELS["stt"] = pipeline(
             "automatic-speech-recognition",
             model=model,
             tokenizer=processor.tokenizer,
             feature_extractor=processor.feature_extractor,
+            chunk_length_s=30, # Fixes >30s audio support
+            batch_size=8,      # Accelerated batch processing
             device=device
         )
     
-    print(f"--- [v161] 🎙️ WHISPER INFERENCE (Temp=0, TS=True) ---")
-    # Added return_timestamps=True to handle > 30s audio
+    print(f"--- [v162] 🎙️ WHISPER INFERENCE (TEMP 0) ---")
     res = MODELS["stt"](
         temp_path, 
         generate_kwargs={
             "language": lang if lang and len(lang) <= 3 else None,
-            "temperature": 0.0,
-            "return_timestamps": True
-        }
+            "temperature": 0.0 # High precision as requested
+        },
+        return_timestamps=True
     )
     return res["text"].strip()
 
-@spaces.GPU(duration=120)
-def gpu_tts_inference(text, mapped_lang, speaker_path):
+@spaces.GPU(duration=180)
+def gpu_tts_full(text, mapped_lang, speaker_path):
     global MODELS
-    device = "cuda" # Always use CUDA inside the decorator
+    device = "cuda"
     
     if MODELS.get("tts") is None:
-        print(f"--- [v161] 📥 LOADING XTTS V2 ON {device} ---")
-        # We load directly to GPU inside the decorator for max stability
+        print(f"--- [v162] 📥 LOADING XTTS V2 ON GPU ---")
+        # Pre-load to RAM if possible or just load directly
         MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
-    
-    # Ensure model is on GPU (sometimes ZeroGPU logic moves it back to CPU)
-    if hasattr(MODELS["tts"], "to"):
-        MODELS["tts"].to(device)
-        
-    print(f"--- [v161] 🔊 XTTS INFERENCE ON GPU ---")
+    else:
+        # Hybrid management: Ensure model is on CUDA during inference
+        try: MODELS["tts"].to(device)
+        except: pass
+
+    print(f"--- [v162] 🔊 XTTS GPU INFERENCE ---")
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as out_f:
         out_p = out_f.name
     
@@ -126,9 +126,12 @@ def gpu_tts_inference(text, mapped_lang, speaker_path):
     with open(out_p, "rb") as f:
         audio_b64 = base64.b64encode(f.read()).decode()
     
-    if os.path.exists(out_p):
-        os.unlink(out_p)
-        
+    if os.path.exists(out_p): os.unlink(out_p)
+    
+    # Cleanup to prevent ZeroGPU worker errors
+    torch.cuda.empty_cache()
+    gc.collect()
+    
     return audio_b64
 
 async def handle_process(request: Request):
@@ -136,37 +139,40 @@ async def handle_process(request: Request):
     try:
         data = await request.json()
         action = data.get("action")
-        if action == "health": return {"status": "awake", "v": "161"}
+        if action == "health": return {"status": "awake", "v": "162"}
         
-        print(f"--- [v161] 🛠️ {action.upper()} REQUESTED ---")
+        print(f"--- [v162] 🛠️ API REQUEST: {action.upper()} ---")
         
         stt_text = ""
-        # 🟢 STT
+        # 🟢 SPEECH-TO-TEXT
         if action in ["stt", "s2st"]:
             audio_b64 = data.get("file")
-            if not audio_b64: return {"error": "Missing audio file"}
+            if not audio_b64: return {"error": "Missing audio data"}
+            
             audio_bytes = base64.b64decode(audio_b64)
             with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                 f.write(audio_bytes); temp_path = f.name
             try:
-                stt_text = gpu_stt_large(temp_path, data.get("lang"))
-                print(f"--- [v161] 🎙️ TRANSCRIPT: {stt_text[:100]}... ---")
+                stt_text = gpu_stt_full(temp_path, data.get("lang"))
+                print(f"--- [v162] 🎙️ TEXT: {stt_text[:100]}... ---")
             finally:
                 if os.path.exists(temp_path): os.unlink(temp_path)
+            
             if action == "stt": return {"text": stt_text}
 
-        # 🔵 TTS
+        # 🔵 TEXT-TO-SPEECH
         if action in ["tts", "s2st"]:
             text = (data.get("text") if action == "tts" else stt_text).strip()
             if not text: return {"error": "Input text is empty"}
             
             target = data.get("target_lang") or data.get("lang") or "en"
             trans_text = text
+            
             if action == "s2st":
-                print(f"--- [v161] 🌏 TRANSLATING TO {target} ---")
+                print(f"--- [v162] 🌏 TRANSLATING TO {target} ---")
                 trans_text = GoogleTranslator(source='auto', target=target).translate(stt_text)
                 text = trans_text
-                print(f"--- [v161] 📝 TRANSLATED: {text[:100]}... ---")
+                print(f"--- [v162] 📝 TRANS: {text[:100]}... ---")
 
             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()
@@ -176,7 +182,7 @@ async def handle_process(request: Request):
                 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: return {"error": f"Language {clean_lang} not supported by XTTS/Chatterbox"}
             else:
                 speaker_wav_b64 = data.get("speaker_wav")
                 speaker_path = None
@@ -188,8 +194,8 @@ async def handle_process(request: Request):
                     if not os.path.exists(speaker_path): speaker_path = None
                 
                 try:
-                    # RUN ON GPU
-                    audio_b64 = gpu_tts_inference(text, mapped_lang, speaker_path)
+                    # EXECUTE GPU TTS
+                    audio_b64 = gpu_tts_full(text, mapped_lang, speaker_path)
                 finally:
                     if speaker_wav_b64 and speaker_path and os.path.exists(speaker_path): os.unlink(speaker_path)
             
@@ -197,10 +203,10 @@ async def handle_process(request: Request):
             return {"text": stt_text, "translated": trans_text, "audio": audio_b64}
 
     except Exception as e:
-        print(f"❌ [v161] ULTIMATE ERROR: {traceback.format_exc()}")
+        print(f"❌ [v162] ENGINE ERROR: {traceback.format_exc()}")
         return {"error": str(e)}
     finally:
-        print(f"--- [v161] ✨ DONE ({time.time()-t1:.1f}s) ---")
+        print(f"--- [v162] ✨ MISSION COMPLETED ({time.time()-t1:.1f}s) ---")
 
 @app.post("/process")
 @app.post("/api/v1/process")
@@ -209,15 +215,14 @@ async def api_process(request: Request): return await handle_process(request)
 @app.get("/health")
 def health(): 
     return {
-        "status": "ok", 
-        "v": "161", 
-        "mode": "ULTIMATE_GPU", 
-        "gpu_available": torch.cuda.is_available(),
-        "gpu_count": torch.cuda.device_count() if torch.cuda.is_available() else 0
+        "status": "ready", 
+        "v": "162", 
+        "gpu": torch.cuda.is_available(),
+        "devices": torch.cuda.device_count()
     }
 
 @app.get("/", response_class=HTMLResponse)
-def root(): return "<h1>🚀 AI Engine v161 (ULTIMATE GPU)</h1>"
+def root(): return "<h1>🚀 PRO AI Engine v162 (GPU MODE)</h1>"
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)