Update app.py

app.py

@@ -1,15 +1,11 @@
 # -*- coding: utf-8 -*-
-# !apt-get install -y ffmpeg
-# !pip install gradio huggingface_hub torch
-# !pip install git+https://github.com/NVIDIA/NeMo.git@main#egg=nemo_toolkit[all]
-
 import os, shlex, subprocess, tempfile, traceback, time, glob, gc, shutil
 import torch
 from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# 1. CONFIGURATION ET MODÈLES (INTÉGRALITÉ CONSERVÉE)
+# 1. CONFIGURATION ET MODÈLES
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 SEGMENT_DURATION = 5.0 
 
@@ -25,55 +21,42 @@ MODELS = {
     "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "rnnt"),
 }
 
-# --- SECTION EXEMPLE (RÉTABLIE) ---
+# --- SECTION EXEMPLE ---
 def find_example_video():
     paths = ["examples/MARALINKE_FIXED.mp4", "examples/MARALINKE.mp4", "MARALINKE.mp4"]
     for p in paths:
         if os.path.exists(p): return p
-    
-    print("⬇️ Téléchargement de la vidéo d'exemple...")
-    example_url = "https://huggingface.co/spaces/RobotsMali/Soloni-Demo/resolve/main/examples/MARALINKE.mp4"
-    target_path = "examples/MARALINKE.mp4"
-    os.makedirs("examples", exist_ok=True)
-    try:
-        subprocess.run(f"wget {example_url} -O {target_path}", shell=True, check=True)
-        return target_path
-    except Exception:
-        return None
+    return None
 
 EXAMPLE_PATH = find_example_video()
 _cache = {}
 
-# 2. GESTION MÉMOIRE ET CHARGEMENT SÉCURISÉ
-def clear_memory():
-    _cache.clear()
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
+# 2. GESTION MÉMOIRE ET CHARGEMENT (NOUVELLE MÉTHODE ANTI-BUG)
 def get_model(name):
     if name in _cache: return _cache[name]
-    clear_memory()
-    repo, _ = MODELS[name]
+    gc.collect()
+    if torch.cuda.is_available(): torch.cuda.empty_cache()
     
+    repo, model_type = MODELS[name]
     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)
-    if not nemo_file: raise FileNotFoundError("Fichier .nemo introuvable.")
-
-    # FIX CRITIQUE : Importation et instanciation explicite du connecteur
-    from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
     
-    # On force NeMo à utiliser cette instance précise pour éviter le bug __init__()
-    model = nemo_asr.models.ASRModel.restore_from(
-        nemo_file, 
-        map_location=torch.device(DEVICE),
-        save_restore_connector=SaveRestoreConnector()
-    )
+    print(f"📦 Restauration forcée du modèle : {name}")
+    
+    # MÉTHODE ALTERNATIVE : On bypass le connecteur automatique
+    try:
+        # Tentative avec la méthode classique mais simplifiée au maximum
+        model = nemo_asr.models.ASRModel.restore_from(nemo_file, map_location=torch.device(DEVICE))
+    except TypeError:
+        # Si l'erreur object.init() survient, on utilise le chargement par classe spécifique
+        print("⚠️ Détection du bug NeMo, basculement sur le chargement spécifique...")
+        if "ctc" in name.lower() or model_type == "ctc":
+            model = nemo_asr.models.EncDecCTCModel.restore_from(nemo_file, map_location=torch.device(DEVICE))
+        else:
+            model = nemo_asr.models.EncDecHybridRNNTCTCModel.restore_from(nemo_file, map_location=torch.device(DEVICE))
     
     model.eval()
-    if DEVICE == "cuda":
-        model = model.half()
-        
+    if DEVICE == "cuda": model = model.half()
     _cache[name] = model
     return model
 
@@ -83,53 +66,42 @@ def format_srt_time(sec):
     ms = int((sec - int(sec)) * 1000)
     return f"{time.strftime('%H:%M:%S', td)},{ms:03}"
 
-# 4. PIPELINE DE TRANSCRIPTION
+# 4. PIPELINE
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in:
-            yield "❌ Aucune vidéo sélectionnée.", None
-            return
+        if not video_in: return "❌ Vidéo manquante", None
         
-        yield "⏳ Extraction audio...", None
+        # Audio
         full_wav = os.path.join(tmp_dir, "full.wav")
         subprocess.run(f"ffmpeg -y -i {shlex.quote(video_in)} -vn -ac 1 -ar 16000 {full_wav}", shell=True, check=True)
         
-        yield f"⏳ Segmentation ({SEGMENT_DURATION}s)...", None
+        # Segmentation
+        yield "⏳ Segmentation...", None
         subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time {SEGMENT_DURATION} -c copy {os.path.join(tmp_dir, 'seg_%03d.wav')}", shell=True, check=True)
         
         files = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
         valid_segments = [f for f in files if os.path.getsize(f) > 1500]
         
-        if not valid_segments:
-            yield "❌ Erreur : Audio trop court ou invalide.", None
-            return
-
-        yield f"⏳ Chargement de {model_name}...", None
+        # Transcription
+        yield f"🎙️ Chargement & Transcription ({model_name})...", None
         model = get_model(model_name)
         
-        yield f"🎙️ Transcription de {len(valid_segments)} segments...", None
-        b_size = 16 if DEVICE == "cuda" else 2
-        
         with torch.inference_mode():
-            batch_hypotheses = model.transcribe(valid_segments, batch_size=b_size, return_hypotheses=True)
+            batch_hypotheses = model.transcribe(valid_segments, batch_size=16 if DEVICE=="cuda" else 2, return_hypotheses=True)
         
+        # SRT Generation
         all_words_ts = []
         for idx, hyp in enumerate(batch_hypotheses):
             base_time = idx * SEGMENT_DURATION
             text = hyp.text if hasattr(hyp, 'text') else str(hyp)
             words = text.split()
             if not words: continue
-            
             gap = SEGMENT_DURATION / len(words)
             for i, w in enumerate(words):
-                all_words_ts.append({
-                    "word": w, 
-                    "start": base_time + (i * gap), 
-                    "end": base_time + ((i+1) * gap)
-                })
+                all_words_ts.append({"word": w, "start": base_time + (i * gap), "end": base_time + ((i+1) * gap)})
 
-        yield "⏳ Encodage vidéo...", None
+        # Encodage
         srt_path = os.path.join(tmp_dir, "final.srt")
         with open(srt_path, "w", encoding="utf-8") as f:
             for i in range(0, len(all_words_ts), 6):
@@ -139,35 +111,29 @@ def pipeline(video_in, model_name):
 
         out_path = os.path.abspath(f"resultat_{int(time.time())}.mp4")
         safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
-        
-        cmd = f"ffmpeg -y -i {shlex.quote(video_in)} -vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18,PrimaryColour=&H00FFFF,Outline=1'\" -c:v libx264 -preset ultrafast -c:a copy {out_path}"
+        cmd = f"ffmpeg -y -i {shlex.quote(video_in)} -vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18'\" -c:v libx264 -preset ultrafast -c:a copy {out_path}"
         subprocess.run(cmd, shell=True, check=True)
         
-        yield "✅ Terminé !", out_path
+        yield "✅ Succès !", out_path
 
     except Exception as e:
-        traceback.print_exc()
         yield f"❌ Erreur : {str(e)}", None
     finally:
         if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir)
 
-# 5. INTERFACE GRADIO
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.HTML("<div style='text-align:center;'><h1>🤖 RobotsMali Speech Lab</h1></div>")
-    
+# 5. INTERFACE
+with gr.Blocks() as demo:
+    gr.HTML("<h1 style='text-align:center;'>RobotsMali Speech Lab</h1>")
     with gr.Row():
         with gr.Column():
-            v_input = gr.Video(label="Vidéo Source")
+            v_input = gr.Video(label="Source")
             m_input = gr.Dropdown(choices=list(MODELS.keys()), value="Soloni V3 (TDT-CTC)", label="Modèle")
             run_btn = gr.Button("🚀 GÉNÉRER", variant="primary")
-            
-            if EXAMPLE_PATH:
-                gr.Examples(examples=[[EXAMPLE_PATH, "Soloni V3 (TDT-CTC)"]], inputs=[v_input, m_input])
-
+            if EXAMPLE_PATH: gr.Examples([[EXAMPLE_PATH, "Soloni V3 (TDT-CTC)"]], [v_input, m_input])
         with gr.Column():
-            status = gr.Markdown("### État\nPrêt.")
+            status = gr.Markdown("Prêt.")
             v_output = gr.Video(label="Résultat")
 
     run_btn.click(pipeline, [v_input, m_input], [status, v_output])
 
-demo.queue().launch()
+demo.launch()