Upload app.py with huggingface_hub

app.py

@@ -5,19 +5,25 @@ 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
-import scipy.io.wavfile
 
-# --- [v147] 🚀 HYBRID ENGINE: GPU-STT + CPU-TTS ---
-print(f"--- [v147] 📡 BOOTING HYBRID ENGINE ---")
+# --- [v148] 🚀 FULL FEATURE H200 ENGINE (STT + TRANS + TTS + FALLBACK) ---
+print(f"--- [v148] 📡 BOOTING FULL ENGINE ---")
 
-from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
-from datasets import load_dataset
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+from TTS.api import TTS
 from deep_translator import GoogleTranslator
 
+try:
+    import chatterbox_utils
+    HAS_CHATTERBOX = True
+except ImportError:
+    HAS_CHATTERBOX = False
+
 try:
     import spaces
     HAS_SPACES = True
@@ -30,110 +36,140 @@ except ImportError:
                 return func
             return decorator
 
+# --- System Config & Stability Flags ---
 os.environ["COQUI_TOS_AGREED"] = "1"
 os.environ["PYTHONWARNINGS"] = "ignore"
+# Disable libraries that might conflict with MIG partitions
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-MODELS = {"stt": None, "tts_processor": None, "tts_model": None, "tts_vocoder": None, "tts_speaker": None}
-
-# Load TTS at startup (CPU)
-print("--- [v147] 📥 LOADING TTS (SpeechT5) ON CPU ---")
-MODELS["tts_processor"] = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-MODELS["tts_model"] = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-MODELS["tts_vocoder"] = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
-embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-MODELS["tts_speaker"] = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
-print("--- [v147] ✅ TTS READY (CPU) ---")
+MODELS = {"stt": None, "tts": None}
 
-@spaces.GPU(duration=60)
-def process_stt(audio_b64, lang):
-    """Speech-to-Text on GPU (Whisper)"""
+def load_stt_gpu():
+    """Load Whisper on GPU with stability flags."""
     global MODELS
-    
     if MODELS.get("stt") is None:
-        print("--- [v147] 📥 LOADING WHISPER ON GPU ---")
+        print("--- [v148] 📥 LOADING WHISPER (Large-v3-Turbo) ---")
+        model_id = "openai/whisper-large-v3-turbo"
+        # Force FP32 and Eager attention to combat CUBLAS errors on H200 drivers
+        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="openai/whisper-base",
+            model=model,
+            tokenizer=processor.tokenizer,
+            feature_extractor=processor.feature_extractor,
+            torch_dtype=torch.float32,
             device="cuda",
-            torch_dtype=torch.float32
+            model_kwargs={"attn_implementation": "eager"}
         )
-        print("--- [v147] ✅ WHISPER READY (GPU) ---")
-    
+        print("--- [v148] ✅ WHISPER LOADED (FP32/EAGER) ---")
+
+def load_tts_gpu():
+    """Load XTTS on GPU."""
+    global MODELS
+    if MODELS.get("tts") is None:
+        print("--- [v148] 📥 LOADING XTTS v2 ---")
+        MODELS["tts"] = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
+        # Explicit stability cast
+        MODELS["tts"].to(torch.float32) 
+        print("--- [v148] ✅ XTTS LOADED (FP32) ---")
+
+@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
+        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})
+        # Use batch_size=1 for kernel stability
+        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)
-
-def process_tts(text):
-    """Text-to-Speech on CPU (SpeechT5)"""
-    inputs = MODELS["tts_processor"](text=text, return_tensors="pt")
-    speech = MODELS["tts_model"].generate_speech(inputs["input_ids"], MODELS["tts_speaker"], vocoder=MODELS["tts_vocoder"])
+        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)
     
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as out_f:
-        out_p = out_f.name
-    scipy.io.wavfile.write(out_p, rate=16000, data=speech.numpy())
+    if not mapped_lang:
+        if HAS_CHATTERBOX:
+            print(f"--- [v148] 📦 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."}
     
-    with open(out_p, "rb") as f:
-        audio_b64 = base64.b64encode(f.read()).decode()
-    os.unlink(out_p)
-    return audio_b64
+    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
+    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
+        # XTTS inference
+        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": "147", "mode": "HYBRID_GPU_STT_CPU_TTS"}
+        if action == "health": return {"status": "awake", "v": "148", "mode": "FULL_FEATURE"}
         
-        print(f"--- [v147] 🛠️ PROCESSING: {action} ---")
+        print(f"--- [v148] 🛠️ ENGINE ACTION: {action} ---")
         t1 = time.time()
         
+        # 🎙️ STT
         stt_text = None
         if action in ["stt", "s2st"]:
-            stt_text = process_stt(data.get("file"), data.get("lang"))
-            if action == "stt":
-                return {"text": stt_text}
+            stt_text = core_stt(data.get("file"), data.get("lang"))
+            if action == "stt": return {"text": stt_text}
 
+        # 🔊 TTS
         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": ""}
+            if len(text) < 2: return {"audio": ""} if action == "tts" else {"text": stt_text, "translated": "", "audio": ""}
             
-            audio_b64 = process_tts(text)
+            audio_res = core_tts(text, (data.get("lang") if action == "tts" else target), data.get("speaker_wav"))
             
-            if action == "tts":
-                return {"audio": audio_b64}
-            return {"text": stt_text, "translated": trans_text, "audio": audio_b64}
+            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}
 
     except Exception as e:
-        print(f"❌ [v147] ERROR: {traceback.format_exc()}")
+        print(f"❌ [v148] ERROR: {traceback.format_exc()}")
         return {"error": str(e)}
     finally:
-        print(f"--- [v147] ✨ DONE ({time.time()-t1:.1f}s) ---")
+        print(f"--- [v148] ✨ DONE ({time.time()-t1:.1f}s) ---")
+        torch.cuda.empty_cache()
 
 @app.get("/health")
-def health():
-    return {"status": "ok", "v": "147", "mode": "HYBRID_GPU_STT_CPU_TTS", "spaces": HAS_SPACES}
+def health(): return {"status": "ok", "v": "148", "mode": "FULL_H200", "gpu": HAS_SPACES}
 
 @app.get("/", response_class=HTMLResponse)
-def root():
-    return f"<html><body><h1>🚀 AI Engine v147 (HYBRID)</h1><p>STT: GPU Whisper | TTS: CPU SpeechT5</p></body></html>"
+def root(): return "<html><body><h1>🚀 AI Engine v148 (FULL FEATURE H200)</h1></body></html>"
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)