Update app.py

app.py

@@ -1,35 +1,139 @@
-# app.py
+# =====================================================
+# HF SPACES – ROBOTSMALI ASR (SAFE MULTI-MODELS)
+# Default: Soloni (lightweight)
+# =====================================================
+
+import os
+os.environ["GRADIO_DISABLE_API"] = "1"
+
 import gradio as gr
-import nemo.collections.asr as nemo_asr
 import torch
+import nemo.collections.asr as nemo_asr
 import warnings
+import gc
+import time
 
 warnings.filterwarnings("ignore")
 
-MODEL = None
+# =====================================================
+# MODELS REGISTRY
+# =====================================================
+MODELS = {
+    "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"),
+
+    "Soloba V3 (CTC) – Haute précision":    ("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 V1.5 (TDT)":                    ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
+    "Soloba V0.5 (TDT)":                    ("RobotsMali/soloba-tdt-0.6b-v0.5", "rnnt"),
+}
+
+DEFAULT_MODEL = "Soloni V3 (TDT-CTC) – Rapide (défaut)"
+
+current_model = None
+current_model_name = None
+
+# =====================================================
+# MODEL LOADER (ONE MODEL AT A TIME)
+# =====================================================
+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
+    if current_model is not None:
+        del current_model
+        current_model = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    model_id, model_type = MODELS[model_name]
+
+    start = time.time()
 
-def load_model():
-    global MODEL
-    if MODEL is None:
-        MODEL = nemo_asr.models.EncDecCTCModel.from_pretrained(
-            "nvidia/stt_en_conformer_ctc_small"
-        )
-        MODEL.eval()
-    return MODEL
+    if model_type == "ctc":
+        model = nemo_asr.models.EncDecCTCModel.from_pretrained(model_id)
+    else:
+        model = nemo_asr.models.EncDecRNNTModel.from_pretrained(model_id)
 
-def transcribe(audio_path):
+    model.eval()
+
+    if torch.cuda.is_available():
+        model = model.cuda()
+
+    if hasattr(model, "decoding"):
+        model.decoding.strategy = "greedy"
+
+    current_model = model
+    current_model_name = model_name
+
+    load_time = round(time.time() - start, 2)
+    print(f"[INFO] Model loaded: {model_name} in {load_time}s")
+
+    return model
+
+# =====================================================
+# TRANSCRIPTION
+# =====================================================
+def transcribe(model_name, audio_path):
     if audio_path is None:
-        return "Aucun audio fourni"
+        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()
-    with torch.no_grad():
-        text = model.transcribe([audio_path])[0]
-    return text
+        return f"✅ Transcription :\n\n{result[0]}"
 
+    except Exception as e:
+        return f"❌ Erreur pendant la transcription :\n{e}"
+
+# =====================================================
+# UI
+# =====================================================
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎙️ ASR Demo – Safe Mode")
-    audio = gr.Audio(type="filepath")
-    out = gr.Textbox()
-    gr.Button("Transcrire").click(transcribe, audio, out)
+    gr.Markdown(
+        """
+        ## 🤖 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
+        """
+    )
+
+    model_selector = gr.Dropdown(
+        choices=list(MODELS.keys()),
+        value=DEFAULT_MODEL,
+        label="Choisir le modèle ASR"
+    )
+
+    audio = gr.Audio(
+        type="filepath",
+        label="Téléverser un fichier audio (wav, mp3, flac)"
+    )
+
+    output = gr.Textbox(
+        label="Résultat de la transcription",
+        lines=8
+    )
+
+    btn = gr.Button("🎙️ Transcrire")
+    btn.click(
+        fn=transcribe,
+        inputs=[model_selector, audio],
+        outputs=output
+    )
 
+# =====================================================
+# HF SPACES LAUNCH (MANDATORY)
+# =====================================================
 demo.launch()