Upload app.py with huggingface_hub

app.py

@@ -13,8 +13,7 @@ import torchaudio
 import chatterbox_utils
 import gc
 
-# 🛡️ ZeroGPU Support (v69)
-# CRITICAL: @spaces.GPU MUST only be used on synchronous functions (def, not async def)
+# 🛡️ ZeroGPU Support
 try:
     import spaces
     print("✅ ZeroGPU/Spaces detected")
@@ -40,7 +39,7 @@ if "torchaudio.backend" not in sys.modules:
     sys.modules["torchaudio.backend"] = backend
     sys.modules["torchaudio.backend.common"] = common
 
-# 🛡️ Torchaudio Compatibility Fix (v60)
+# 🛡️ Torchaudio Compatibility Fix
 if not hasattr(torchaudio, "info"):
     print("🛠️ Mocking torchaudio.info for compatibility...")
     def mock_info(filepath, **kwargs):
@@ -61,8 +60,8 @@ if not hasattr(torchaudio, "info"):
 
 from df.enhance import enhance, init_df, load_audio, save_audio
 
-# FORCE BUILD TRIGGER: 10:00:00 Jan 21 2026
-# v85: ZeroGPU Warmup & Pre-caching (Prevents session timeouts)
+# FORCE BUILD TRIGGER: 10:10:00 Jan 21 2026
+# v86: Pre-load to CPU RAM + Fast Transfer to GPU (Prevents ZeroGPU timeouts)
 
 # 🛠️ Monkeypatch torchaudio.load
 try:
@@ -87,88 +86,90 @@ except Exception as e:
 
 os.environ["COQUI_TOS_AGREED"] = "1"
 
-# Global models
-MODELS = {"stt": None, "translate": None, "tts": None, "tokenizer": None, "denoiser": None}
+# Global models (Resident in RAM)
+MODELS = {"stt": None, "translate": None, "tts": None, "denoiser": None}
 
 def load_models():
+    """Fast GPU Activation: Moves pre-loaded CPU models to GPU (v86)"""
     global MODELS
+    
+    # 1. Faster-Whisper (Must re-init for device change, but disk cache is hot)
     if MODELS["stt"] is None:
-        print("🎙️ Loading Faster-Whisper large-v3 into Engine...")
+        print("🎙️ Initializing Faster-Whisper...")
         from faster_whisper import WhisperModel
-        if torch.cuda.is_available():
-            print(f"🚀 GPU Detected: {torch.cuda.get_device_name(0)}")
-            MODELS["stt"] = WhisperModel("large-v3", device="cuda", compute_type="float16")
-        else:
-            print("⚠️ Falling back to CPU (int8)")
-            MODELS["stt"] = WhisperModel("large-v3", device="cpu", compute_type="int8")
+        dev = "cuda" if torch.cuda.is_available() else "cpu"
+        ct = "float16" if dev == "cuda" else "int8"
+        MODELS["stt"] = WhisperModel("large-v3", device=dev, compute_type=ct)
     
-    # 🧹 Proactive Memory Cleanup
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    # 2. XTTS-v2 (Efficient .to("cuda") transfer)
+    if MODELS["tts"] is None:
+        print("🔊 Loading XTTS-v2 into Engine (CPU Base)...")
+        from TTS.api import TTS
+        MODELS["tts"] = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)
     
-    # Initialize Chatterbox ONNX (High-Speed Fallback)
+    if torch.cuda.is_available():
+        # Move XTTS to GPU if it's currently on CPU
+        # We check the device of the underlying model to avoid redundant moves
+        try:
+            current_dev = str(next(MODELS["tts"].synthesizer.tts_model.parameters()).device)
+            if "cuda" not in current_dev:
+                print("🚀 Moving XTTS-v2 to GPU...")
+                MODELS["tts"].to("cuda")
+        except: 
+            # Fallback if structure differs
+            MODELS["tts"].to("cuda")
+
+    # 3. DeepFilterNet
+    if MODELS["denoiser"] is None:
+        try: MODELS["denoiser"] = init_df()
+        except: pass
+
+    # 4. Chatterbox ONNX
     chatterbox_utils.load_chatterbox(device="cuda" if torch.cuda.is_available() else "cpu")
 
+    # 5. Translate
     if MODELS["translate"] is None:
-        print("🌍 Loading Google Translate...")
         MODELS["translate"] = "active"
 
-    if MODELS["denoiser"] is None:
-        print("🧹 Loading DeepFilterNet...")
-        try:
-            MODELS["denoiser"] = init_df()
-            print("✨ DeepFilterNet Loaded")
-        except: pass
-
-    if MODELS["tts"] is None:
-        print("🔊 Loading XTTS-v2 into Engine...")
-        from TTS.api import TTS
-        try:
-            MODELS["tts"] = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", gpu=torch.cuda.is_available())
-            print(f"✨ XTTS-v2 Loaded (GPU={torch.cuda.is_available()})")
-        except Exception as e:
-            print(f"❌ Failed to load XTTS: {e}")
-            raise e
+    # 🧹 Proactive Memory Cleanup
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
 def warmup_models():
-    """Download and cache all models on CPU at startup (Prevents GPU timeouts)"""
-    print("\n🔥 --- SYSTEM WARMUP STARTING (CPU) ---")
+    """PRE-LOAD EVERYTHING INTO SYSTEM RAM (CPU)"""
+    print("\n🔥 --- SYSTEM WARMUP: RESIDENT RAM LOADING (v86) ---")
     start = time.time()
     try:
-        # 1. Warmup Whisper
-        print("📥 Initializing Whisper large-v3 cache...")
+        # Load Whisper into RAM
+        print("��� Pre-loading Whisper to RAM...")
         from faster_whisper import WhisperModel
-        # Use simple init to trigger download
-        _ = WhisperModel("large-v3", device="cpu", compute_type="int8")
+        MODELS["stt"] = WhisperModel("large-v3", device="cpu", compute_type="int8")
         
-        # 2. Warmup XTTS-v2 (This takes the longest)
-        print("📥 Initializing XTTS-v2 cache...")
+        # Load XTTS into RAM (The heaviest part)
+        print("📥 Pre-loading XTTS-v2 to RAM...")
         from TTS.api import TTS
-        # Initialize once on CPU just to force download
-        _ = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)
+        MODELS["tts"] = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)
         
-        # 3. Warmup DeepFilterNet
-        print("📥 Initializing DeepFilterNet cache...")
-        try: init_df()
+        # Load Denoiser
+        print("📥 Pre-loading Denoiser...")
+        try: MODELS["denoiser"] = init_df()
         except: pass
         
-        # 4. Warmup Chatterbox
+        # Pre-download ONNX weights
         chatterbox_utils.warmup_chatterbox()
         
-        print(f"✅ --- SYSTEM WARMUP COMPLETE (Time: {time.time()-start:.2f}s) --- \n")
+        print(f"✅ --- WARMUP COMPLETE: All models resident in RAM ({time.time()-start:.2f}s) --- \n")
     except Exception as e:
         print(f"⚠️ Warmup warning: {e}")
 
 def _stt_logic(request_dict):
-    """STT Logic (Runs on GPU when called via core_process)"""
     audio_bytes = base64.b64decode(request_dict.get("file"))
     lang = request_dict.get("lang")
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
         f.write(audio_bytes)
         temp_path = f.name
     try:
-        # Transcribe (Uses GPU if device="cuda" in MODELS)
         segments, _ = MODELS["stt"].transcribe(temp_path, language=lang, beam_size=1)
         text = " ".join([s.text for s in segments]).strip()
         return {"text": text}
@@ -176,100 +177,68 @@ def _stt_logic(request_dict):
         if os.path.exists(temp_path): os.unlink(temp_path)
 
 def _translate_logic(text, target_lang):
-    """Translation (CPU/Network)"""
     from deep_translator import GoogleTranslator
-    translated = GoogleTranslator(source='auto', target=target_lang).translate(text)
-    return translated
+    return GoogleTranslator(source='auto', target=target_lang).translate(text)
 
 def _tts_logic(text, lang, speaker_wav_b64):
-    """TTS Logic (Runs on GPU when called via core_process)"""
-    if not text or not text.strip():
-        return {"error": "TTS Error: Input text is empty"}
+    if not text or not text.strip(): return {"error": "Input empty"}
         
     XTTS_MAP = {
-        "en": "en", "en-us": "en", "en-gb": "en",
-        "de": "de", "de-de": "de",
-        "fr": "fr", "fr-fr": "fr",
-        "es": "es", "es-es": "es",
-        "it": "it", "it-it": "it",
-        "pl": "pl", "pl-pl": "pl",
-        "pt": "pt", "pt-pt": "pt", "pt-br": "pt",
-        "tr": "tr", "tr-tr": "tr",
-        "ru": "ru", "ru-ru": "ru",
-        "nl": "nl", "nl-nl": "nl",
-        "cs": "cs", "cs-cz": "cs",
-        "ar": "ar", "ar-sa": "ar", "ar-eg": "ar",
-        "hu": "hu", "hu-hu": "hu",
-        "ko": "ko", "ko-kr": "ko",
-        "hi": "hi", "hi-in": "hi",
-        "zh": "zh-cn", "zh-cn": "zh-cn", "zh-tw": "zh-cn"
+        "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"
     }
     
-    XTTS_LANG_CODES = set(XTTS_MAP.values())
-    mapped_lang = None
-    if lang:
-        lang_key = lang.strip().lower()
-        mapped_lang = XTTS_MAP.get(lang_key) or XTTS_MAP.get(lang_key.split('-')[0])
+    clean_lang = lang.strip().lower().split('-')[0]
+    mapped_lang = XTTS_MAP.get(clean_lang)
+    if clean_lang == "zh": mapped_lang = "zh-cn"
         
-    print(f"[v85] TTS Request - Original: {lang}, Mapped: {mapped_lang}")
+    print(f"[v86] TTS: {lang} -> {mapped_lang}")
     
-    if mapped_lang and mapped_lang in XTTS_LANG_CODES:
-        print(f"[v85] Using XTTS-v2 for '{mapped_lang}'")
+    if mapped_lang:
+        print(f"[v86] GPU Mode: XTTS-v2")
         speaker_wav_path = None
         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
-        else:
-            speaker_wav_path = "default_speaker.wav"
+                f.write(sb); speaker_wav_path = f.name
+        else: speaker_wav_path = "default_speaker.wav"
             
         try:
             with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as output_file:
                 output_path = output_file.name
-            
             MODELS["tts"].tts_to_file(text=text, language=mapped_lang, file_path=output_path, speaker_wav=speaker_wav_path)
-            
-            with open(output_path, "rb") as f:
-                audio_b64 = base64.b64encode(f.read()).decode()
+            with open(output_path, "rb") as f: audio_b64 = base64.b64encode(f.read()).decode()
             return {"audio": audio_b64}
         finally:
             if speaker_wav_path and "default_speaker" not in speaker_wav_path:
                 if os.path.exists(speaker_wav_path): os.unlink(speaker_wav_path)
             if 'output_path' in locals() and os.path.exists(output_path): os.unlink(output_path)
             
-    print(f"[v85] Using Chatterbox ONNX Fallback for '{lang}'")
+    # Fallback to Chatterbox
+    print(f"[v86] Fallback Mode: Chatterbox ONNX")
     try:
         temp_ref = None
         if speaker_wav_b64:
              sb = base64.b64decode(speaker_wav_b64)
              with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                  f.write(sb); temp_ref = f.name
-        
-        chatter_lang = lang.strip().lower().split('-')[0]
-        audio_bytes = chatterbox_utils.run_chatterbox_inference(text, chatter_lang, speaker_wav_path=temp_ref)
+        audio_bytes = chatterbox_utils.run_chatterbox_inference(text, clean_lang, speaker_wav_path=temp_ref)
         if temp_ref and os.path.exists(temp_ref): os.unlink(temp_ref)
-        audio_b64 = base64.b64encode(audio_bytes).decode()
-        return {"audio": audio_b64}
+        return {"audio": base64.b64encode(audio_bytes).decode()}
     except Exception as e:
-        print(f"❌ Chatterbox Fallback failed: {e}")
-        return {"error": f"TTS Failure: '{lang}' not supported by XTTS or Chatterbox."}
+        return {"error": f"TTS Failure: {str(e)}"}
 
 @spaces.GPU
 def core_process(request_dict):
-    """Unified GPU Entry Point (v85)"""
     action = request_dict.get("action")
     t0 = time.time()
-    print(f"--- [v85] 🚀 GPU SESSION START: {action} at {time.ctime()} ---")
+    print(f"--- [v86] 🚀 GPU SESSION START: {action} ---")
     load_models() 
-    
     try:
-        if action == "stt":
-            res = _stt_logic(request_dict)
-        elif action == "translate":
-            res = {"translated": _translate_logic(request_dict.get("text"), request_dict.get("target_lang", "en"))}
-        elif action == "tts":
-            res = _tts_logic(request_dict.get("text"), request_dict.get("lang"), request_dict.get("speaker_wav"))
+        if action == "stt": res = _stt_logic(request_dict)
+        elif action == "translate": res = {"translated": _translate_logic(request_dict.get("text"), request_dict.get("target_lang", "en"))}
+        elif action == "tts": res = _tts_logic(request_dict.get("text"), request_dict.get("lang"), request_dict.get("speaker_wav"))
         elif action == "s2st":
             stt_res = _stt_logic({"file": request_dict.get("file"), "lang": request_dict.get("source_lang")})
             text = stt_res.get("text", "")
@@ -277,19 +246,15 @@ def core_process(request_dict):
             translated = _translate_logic(text, request_dict.get("target_lang"))
             tts_res = _tts_logic(translated, request_dict.get("target_lang"), request_dict.get("speaker_wav"))
             res = {"text": text, "translated": translated, "audio": tts_res.get("audio")}
-        elif action == "health":
-            res = {"status": "awake", "time": time.ctime()}
-        else:
-            res = {"error": f"Unknown action: {action}"}
+        elif action == "health": res = {"status": "awake"}
+        else: res = {"error": f"Unknown action: {action}"}
     finally:
-        print(f"--- [v85] ✨ SESSION END: {action} (Total: {time.time()-t0:.2f}s) ---")
+        print(f"--- [v86] ✨ SESSION END: {action} ({time.time()-t0:.2f}s) ---")
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
     return res
 
 def create_wav_header(sample_rate=24000, channels=1, bit_depth=16):
-    """Returns a standard WAV header as standard BYTES"""
     header = bytearray(b'RIFF')
     header.extend((1000000000).to_bytes(4, 'little')) 
     header.extend(b'WAVEfmt ')
@@ -310,21 +275,14 @@ def gpu_tts_generator(text, lang, speaker_wav_path):
     try:
         yield bytes(create_wav_header(sample_rate=24000))
         for chunk in MODELS["tts"].synthesizer.tts_model.inference_stream(
-            text,
-            lang,
-            *MODELS["tts"].synthesizer.tts_model.get_conditioning_latents(audio_path=[speaker_wav_path]),
+            text, lang, *MODELS["tts"].synthesizer.tts_model.get_conditioning_latents(audio_path=[speaker_wav_path]),
             stream_chunk_size=20
         ):
             yield bytes((chunk * 32767).to(torch.int16).cpu().numpy().tobytes())
-        print("✨ [Generator Complete]")
-    except Exception as e:
-        print(f"❌ [Generator Error]: {e}")
     finally:
-        if speaker_wav_path and "default_speaker" not in speaker_wav_path:
-            if os.path.exists(speaker_wav_path): os.unlink(speaker_wav_path)
+        if speaker_wav_path and "default_speaker" not in speaker_wav_path and os.path.exists(speaker_wav_path): os.unlink(speaker_wav_path)
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
 
 app = FastAPI()
 
@@ -332,8 +290,7 @@ app = FastAPI()
 async def api_process(request: Request):
     try:
         data = await request.json()
-        result = core_process(data)
-        return result
+        return core_process(data)
     except Exception as e:
         traceback.print_exc()
         return {"error": str(e)}
@@ -343,37 +300,29 @@ async def api_tts_stream(request: Request):
     try:
         data = await request.json()
         speaker_wav_b64 = data.get("speaker_wav")
-        speaker_wav_path = None
         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
-        else:
-            speaker_wav_path = "default_speaker.wav"
+                f.write(sb); speaker_wav_path = f.name
+        else: speaker_wav_path = "default_speaker.wav"
         return StreamingResponse(gpu_tts_generator(data.get("text"), data.get("lang"), speaker_wav_path), media_type="audio/wav")
-    except Exception as e:
-        return {"error": str(e)}
+    except Exception as e: return {"error": str(e)}
 
 @app.get("/health")
-def health():
-    return {"status": "ok", "gpu": torch.cuda.is_available(), "time": time.ctime()}
+def health(): return {"status": "ok", "gpu": torch.cuda.is_available(), "time": time.ctime()}
 
 @app.post("/api/v1/clear_cache")
 async def clear_cache():
     try:
-        t0 = time.time()
-        gc.collect()
+        t0 = time.time(); gc.collect()
         if torch.cuda.is_available(): torch.cuda.empty_cache()
-        temp_dir = tempfile.gettempdir()
-        count = 0
+        temp_dir = tempfile.gettempdir(); count = 0
         for f in os.listdir(temp_dir):
             if f.endswith(".wav") or f.startswith("tm"):
                 try: os.unlink(os.path.join(temp_dir, f)); count += 1
                 except: pass
-        return {"status": "success", "cleaned_files": count, "duration": f"{time.time()-t0:.2f}s"}
-    except Exception as e:
-        return {"status": "error", "message": str(e)}
+        return {"status": "success", "cleaned_files": count}
+    except Exception as e: return {"status": "error", "message": str(e)}
 
 def gradio_fn(req_json):
     try: return json.dumps(core_process(json.loads(req_json)))