Update app.py

app.py

@@ -1,6 +1,6 @@
 # =====================================================
-# HF SPACES – ROBOTSMALI ASR (SAFE MULTI-MODELS)
-# Default: Soloni (lightweight)
+# HF SPACES – ROBOTSMALI ASR (STABLE MULTI-MODELS)
+# Basé sur le pipeline vidéo fonctionnel
 # =====================================================
 
 import os
@@ -8,22 +8,36 @@ os.environ["GRADIO_DISABLE_API"] = "1"
 
 import gradio as gr
 import torch
-import nemo.collections.asr as nemo_asr
 import warnings
 import gc
 import time
+from huggingface_hub import snapshot_download
+from nemo.collections import asr as nemo_asr
 
 warnings.filterwarnings("ignore")
 
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
 # =====================================================
-# MODELS REGISTRY
+# MODELS REGISTRY (TOUS LES MODÈLES VIDÉO)
 # =====================================================
 MODELS = {
+    # Soloba CTC
+    "Soloba V3 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
+    "Soloba V2 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v2", "ctc"),
+    "Soloba V1 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v1", "ctc"),
+
+    # Soloba TDT
+    "Soloba V1.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
+    "Soloba V0.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v0.5", "rnnt"),
+
+    # Soloni
     "Soloni V3 (TDT-CTC) – Rapide (défaut)": ("RobotsMali/soloni-114m-tdt-ctc-v3", "rnnt"),
-    "Soloni V2 (TDT-CTC)":                  ("RobotsMali/soloni-114m-tdt-ctc-v2", "rnnt"),
-    "Soloni V1 (TDT-CTC)":                  ("RobotsMali/soloni-114m-tdt-ctc-v1", "rnnt"),
+    "Soloni V2 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v2", "rnnt"),
+    "Soloni V1 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v1", "rnnt"),
 
-    "Soloba V3 (CTC) – Haute précision":    ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
+    # Traduction
+    "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "rnnt"),
 }
 
 DEFAULT_MODEL = "Soloni V3 (TDT-CTC) – Rapide (défaut)"
@@ -32,16 +46,15 @@ current_model = None
 current_model_name = None
 
 # =====================================================
-# MODEL LOADER (ONE MODEL AT A TIME)
+# MODEL LOADER (ROBUSTE)
 # =====================================================
 def load_model(model_name):
     global current_model, current_model_name
 
-    # Already loaded
     if current_model_name == model_name:
         return current_model
 
-    # Clean previous model from memory
+    # Nettoyage mémoire
     if current_model is not None:
         del current_model
         current_model = None
@@ -49,22 +62,30 @@ def load_model(model_name):
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-    model_id, model_type = MODELS[model_name]
+    repo, _ = MODELS[model_name]
 
     start = time.time()
 
-    if model_type == "ctc":
-        model = nemo_asr.models.EncDecCTCModel.from_pretrained(model_id)
-    else:
-        model = nemo_asr.models.EncDecRNNTModel.from_pretrained(model_id)
+    # Téléchargement manuel du repo
+    folder = snapshot_download(repo, local_dir_use_symlinks=False)
 
-    model.eval()
+    # Recherche du fichier .nemo
+    nemo_file = next(
+        os.path.join(folder, f)
+        for f in os.listdir(folder)
+        if f.endswith(".nemo")
+    )
 
-    if torch.cuda.is_available():
-        model = model.cuda()
+    # Chargement stable
+    model = nemo_asr.models.ASRModel.restore_from(
+        nemo_file,
+        map_location=torch.device(DEVICE)
+    )
 
-    if hasattr(model, "decoding"):
-        model.decoding.strategy = "greedy"
+    model.eval()
+
+    if DEVICE == "cuda":
+        model = model.half()
 
     current_model = model
     current_model_name = model_name
@@ -82,14 +103,21 @@ def transcribe(model_name, audio_path):
         return "❌ Aucun fichier audio fourni."
 
     try:
-        asr_model = load_model(model_name)
-        with torch.no_grad():
-            result = asr_model.transcribe([audio_path])
+        model = load_model(model_name)
+
+        with torch.inference_mode():
+            result = model.transcribe(
+                [audio_path],
+                batch_size=1,
+                num_workers=0
+            )
 
-        return f"✅ Transcription :\n\n{result[0]}"
+        text = result[0] if isinstance(result, list) else str(result)
+
+        return f"✅ Transcription :\n\n{text}"
 
     except Exception as e:
-        return f"❌ Erreur pendant la transcription :\n{e}"
+        return f"❌ Erreur pendant la transcription :\n{str(e)}"
 
 # =====================================================
 # UI
@@ -99,9 +127,9 @@ with gr.Blocks() as demo:
         """
         ## 🤖 RobotsMali ASR – Démo officielle
 
-        🟢 **Soloni** (léger) est utilisé par défaut pour un démarrage rapide  
-        ⚠️ **Soloba** est plus précis mais peut prendre plusieurs minutes à charger  
-        ℹ️ Un seul modèle est chargé à la fois pour éviter les crashs mémoire
+        🟢 **Soloni** démarre plus rapidement  
+        ⚠️ **Soloba 0.6B** peut prendre du temps à charger  
+        ℹ️ Un seul modèle est chargé à la fois pour stabilité maximale
         """
     )
 
@@ -122,6 +150,7 @@ with gr.Blocks() as demo:
     )
 
     btn = gr.Button("🎙️ Transcrire")
+
     btn.click(
         fn=transcribe,
         inputs=[model_selector, audio],
@@ -129,6 +158,6 @@ with gr.Blocks() as demo:
     )
 
 # =====================================================
-# HF SPACES LAUNCH (MANDATORY)
+# HF SPACES LAUNCH
 # =====================================================
-demo.launch()
+demo.launch(server_name="0.0.0.0", server_port=7860)