Update app.py

app.py

@@ -1,134 +1,136 @@
 # -*- coding: utf-8 -*-
-import os
-import tempfile
-import glob
-import shutil
+import os, shlex, subprocess, tempfile, traceback, time, glob, gc, shutil
 import torch
-import gradio as gr
 from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
+import gradio as gr
 
-# Configuration
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-SEGMENT_DURATION = 10.0
+SEGMENT_DURATION = 10.0  # Augmenté à 10s pour l'audio pur (plus efficace)
 
-# Modèle unique pour CPU (CTC est plus léger)
+# --- CONFIGURATION DES MODÈLES (Identique au script vidéo) ---
 MODELS = {
-    "Soloni V3 (Recommandé CPU)": "RobotsMali/soloni-114m-tdt-ctc-v3"
+    "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 V1.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
+    "Soloba V0.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v0.5", "rnnt"),
+    "Soloni V3 (TDT-CTC)":       ("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"),
+    "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "rnnt"),
 }
 
 _cache = {}
 
-# ==========================
-# CHARGEMENT DU MODÈLE
-# ==========================
-def get_model(model_name):
-    if model_name in _cache:
-        return _cache[model_name]
-
-    repo_id = MODELS[model_name]
-    print(f"⏳ Chargement du modèle : {repo_id}")
+def get_model(name):
+    """Charge le modèle avec gestion de la mémoire."""
+    if name in _cache:
+        return _cache[name]
     
-    # Téléchargement local
-    folder = snapshot_download(repo_id, local_dir_use_symlinks=False)
-
-    # Trouve le fichier .nemo
-    try:
-        nemo_file = next(
-            os.path.join(folder, f)
-            for f in os.listdir(folder)
-            if f.endswith(".nemo")
-        )
-    except StopIteration:
-        raise FileNotFoundError("Fichier .nemo introuvable.")
-
-    # Restauration générique (Auto-détecte CTC/RNNT)
-    model = nemo_asr.models.ASRModel.restore_from(
-        nemo_file,
-        map_location=torch.device(DEVICE)
-    )
+    # Libérer la mémoire des anciens modèles si nécessaire
+    if len(_cache) > 0:
+        _cache.clear()
+        gc.collect()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+
+    repo, _ = MODELS[name]
+    print(f"⏳ Chargement de {name} depuis {repo}...")
+    
+    folder = snapshot_download(repo, local_dir_use_symlinks=False)
+    nemo_file = next((os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".nemo")), None)
 
+    # Chargement via ASRModel (Factory)
+    model = nemo_asr.models.ASRModel.restore_from(nemo_file, map_location=torch.device(DEVICE))
     model.eval()
-    _cache[model_name] = model
+    
+    if DEVICE == "cuda":
+        model = model.half()
+        
+    _cache[name] = model
     return model
 
-# ==========================
-# LOGIQUE DE TRANSCRIPTION
-# ==========================
-def pipeline(audio_path, model_name):
-    if not audio_path:
-        yield "❌ Erreur", "Veuillez fournir un fichier audio."
-        return
-
+# --- PIPELINE ASR ---
+def pipeline(audio_in, model_name):
     tmp_dir = tempfile.mkdtemp()
-    
     try:
-        yield "⏳ Traitement audio...", ""
+        if not audio_in:
+            yield "❌ Fichier audio manquant", ""
+            return
 
-        # Segmentation FFmpeg (16kHz Mono requis)
-        # On utilise une syntaxe simple pour éviter les erreurs de shell
-        os.system(
-            f"ffmpeg -y -i '{audio_path}' -f segment -segment_time {SEGMENT_DURATION} "
-            f"-ac 1 -ar 16000 {tmp_dir}/seg_%03d.wav > /dev/null 2>&1"
-        )
+        yield "⏳ Traitement audio & Segmentation...", ""
+        
+        # Normalisation : Mono, 16kHz
+        wav_path = os.path.join(tmp_dir, "clean.wav")
+        subprocess.run(f"ffmpeg -y -i {shlex.quote(audio_in)} -ac 1 -ar 16000 {wav_path}", shell=True, check=True)
+        
+        # Découpage en segments
+        subprocess.run(f"ffmpeg -i {wav_path} -f segment -segment_time {SEGMENT_DURATION} -c copy {os.path.join(tmp_dir, 'seg_%03d.wav')}", shell=True)
 
-        segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
-        segments = [s for s in segments if os.path.getsize(s) > 1000]
+        valid_segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
+        valid_segments = [f for f in valid_segments if os.path.getsize(f) > 1000]
 
-        if not segments:
-            yield "❌ Erreur", "Impossible de segmenter l'audio."
+        if not valid_segments:
+            yield "❌ Erreur de segmentation", ""
             return
 
-        yield f"🎙️ Transcription ({len(segments)} segments)...", ""
-
+        yield f"🎙️ Transcription de {len(valid_segments)} segments...", ""
+        
         model = get_model(model_name)
 
         with torch.inference_mode():
-            # Transcription par batch pour le CPU
-            results = model.transcribe(segments, batch_size=2, num_workers=0)
+            # Utilisation de la méthode stable sans Lhotse
+            batch_hyp = model.transcribe(
+                valid_segments, 
+                batch_size=4, 
+                return_hypotheses=True,
+                num_workers=0
+            )
 
-        # Gestion des différents formats de sortie de NeMo
-        if isinstance(results, tuple):
-            text_list = results[0]
-        else:
-            text_list = results
+        # Extraction du texte
+        results = []
+        for hyp in batch_hyp:
+            text = hyp.text if hasattr(hyp, 'text') else str(hyp)
+            if text:
+                results.append(text)
 
-        final_text = " ".join(text_list)
-        yield "✅ Terminé", final_text
+        final_text = " ".join(results)
+        
+        yield "✅ Transcription terminée", final_text
 
     except Exception as e:
-        yield "❌ Erreur Système", str(e)
-
+        print(traceback.format_exc())
+        yield f"❌ Erreur : {str(e)}", ""
     finally:
         if os.path.exists(tmp_dir):
             shutil.rmtree(tmp_dir)
 
-# ==========================
-# INTERFACE UTILISATEUR
-# ==========================
-with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 RobotsMali ASR")
-    gr.Markdown("Outil de transcription automatique optimisé pour le CPU.")
-
+# --- INTERFACE GRADIO ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🤖 RobotsMali Speech-to-Text")
+    gr.Markdown("Transcription multi-modèles pour les langues du Mali.")
+    
     with gr.Row():
         with gr.Column():
-            audio_input = gr.Audio(type="filepath", label="Audio")
+            audio_input = gr.Audio(label="Audio (Fichier ou Micro)", type="filepath")
             model_input = gr.Dropdown(
-                choices=list(MODELS.keys()),
-                value="Soloni V3 (Recommandé CPU)",
-                label="Modèle"
+                choices=list(MODELS.keys()), 
+                value="Soloni V3 (TDT-CTC)", 
+                label="Choisir un modèle"
             )
-            btn = gr.Button("🚀 Transcrire", variant="primary")
-
+            run_btn = gr.Button("🚀 TRANSCRIRE", variant="primary")
+            
         with gr.Column():
-            status_output = gr.Textbox(label="Statut", interactive=False)
-            text_output = gr.Textbox(label="Texte Transcrit", lines=10)
+            status = gr.Markdown("### État : Prêt")
+            text_output = gr.Textbox(label="Résultat", lines=15, show_copy_button=True)
 
-    btn.click(
-        fn=pipeline,
-        inputs=[audio_input, model_input],
-        outputs=[status_output, text_output]
+    run_btn.click(
+        fn=pipeline, 
+        inputs=[audio_input, model_input], 
+        outputs=[status, text_output]
     )
 
+    gr.HTML("<center><p style='color: gray;'>Optimisé pour CPU & GPU Hugging Face</p></center>")
+
 if __name__ == "__main__":
     demo.launch()