Update app.py

app.py

@@ -5,7 +5,7 @@ from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# 1. CONFIGURATION MATÉRIEL ET LISTE DES MODÈLES ROBOTSMALI
+# 1. CONFIGURATION ET MODÈLES
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 MODELS = {
@@ -23,24 +23,17 @@ MODELS = {
     "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "rnnt"),
 }
 
-# --- OPTIMISATION : DETECTION DE LA VIDEO DANS LE DOSSIER EXAMPLES ---
+# Détection de la vidéo d'exemple dans le dossier local
 def find_example_video():
-    # Liste des noms possibles basés sur ta capture d'écran
-    paths = [
-        "examples/MARALINKE_FIXED.mp4", 
-        "examples/MARALINKE.mp4",
-        "MARALINKE.mp4"
-    ]
+    paths = ["examples/MARALINKE_FIXED.mp4", "examples/MARALINKE.mp4", "MARALINKE.mp4"]
     for p in paths:
-        if os.path.exists(p):
-            return p
+        if os.path.exists(p): return p
     return None
 
 EXAMPLE_PATH = find_example_video()
-
 _cache = {}
 
-# 2. GESTION DE LA MÉMOIRE ET CHARGEMENT DU MODÈLE
+# 2. OPTIMISATION CHARGEMENT ET MÉMOIRE
 def clear_memory():
     _cache.clear()
     gc.collect()
@@ -52,13 +45,9 @@ def get_model(name):
     clear_memory()
     repo, _ = MODELS[name]
     
-    print(f"📥 Téléchargement de {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)
     
-    if not nemo_file: raise FileNotFoundError("Fichier .nemo introuvable.")
-
-    # Optimisation RobotsMali : Connecteur flexible pour éviter l'erreur state_dict
     from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
     model = nemo_asr.models.ASRModel.restore_from(
         nemo_file, 
@@ -73,48 +62,44 @@ def get_model(name):
     _cache[name] = model
     return model
 
-# 3. UTILITAIRES
+# 3. TRANSCRIPTION ET RENDU
 def format_srt_time(sec):
     td = time.gmtime(sec)
     ms = int((sec - int(sec)) * 1000)
     return f"{time.strftime('%H:%M:%S', td)},{ms:03}"
 
-# 4. PIPELINE DE TRANSCRIPTION
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: return "❌ Veuillez sélectionner une vidéo.", None
+        if not video_in: return "❌ Vidéo absente", None
         
-        # A. Extraction Audio
-        yield "⏳ Phase 1/4 : Extraction audio...", None
+        # FFmpeg optimisé (-threads 0)
+        yield "🚀 Phase 1 : Extraction Audio...", None
         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)
+        subprocess.run(f"ffmpeg -y -threads 0 -i {shlex.quote(video_in)} -vn -ac 1 -ar 16000 {full_wav}", shell=True, check=True)
         
-        # B. Segmentation
-        yield "⏳ Phase 2/4 : Segmentation (20s)...", None
+        yield "🚀 Phase 2 : Segmentation...", None
         subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time 20 -c copy {os.path.join(tmp_dir, 'seg_%03d.wav')}", shell=True, check=True)
         audio_segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
         
-        # C. Transcription IA
-        yield f"⏳ Phase 3/4 : Chargement de {model_name}...", None
+        # Inférence par BATCH (beaucoup plus rapide)
+        yield f"🎙️ Phase 3 : Transcription Batch avec {model_name}...", None
         model = get_model(model_name)
         
+        # 
+        batch_hypotheses = model.transcribe(audio_segments, batch_size=4, return_hypotheses=True)
+        
         all_words_ts = []
-        for idx, seg_path in enumerate(audio_segments):
+        for idx, hyp in enumerate(batch_hypotheses):
             base_time = idx * 20
-            yield f"🎙️ Transcription segment {idx+1}/{len(audio_segments)}...", None
-            hyp = model.transcribe([seg_path], return_hypotheses=True)[0]
             if isinstance(hyp, list): hyp = hyp[0]
             text = hyp.text if hasattr(hyp, 'text') else str(hyp)
             words = text.split()
-            
-            # Répartition temporelle
             gap = 20.0 / max(len(words), 1)
             for i, w in enumerate(words):
                 all_words_ts.append({"word": w, "start": base_time + (i * gap), "end": base_time + ((i+1) * gap)})
 
-        # D. Génération SRT et Rendu Vidéo
-        yield "⏳ Phase 4/4 : Incrustation sous-titres...", None
+        # Génération SRT
         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):
@@ -122,40 +107,36 @@ def pipeline(video_in, model_name):
                 f.write(f"{(i//6)+1}\n{format_srt_time(chunk[0]['start'])} --> {format_srt_time(chunk[-1]['end'])}\n")
                 f.write(" ".join([c['word'] for c in chunk]) + "\n\n")
 
-        out_path = os.path.abspath(f"robotsmali_result_{int(time.time())}.mp4")
+        # Rendu Final FFmpeg accéléré
+        yield "🚀 Phase 4 : Encodage Vidéo...", None
+        out_path = os.path.abspath(f"robotsmali_final_{int(time.time())}.mp4")
         safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
         
-        # Style : Jaune, Taille 18, Centré en bas
-        cmd = f"ffmpeg -y -i {shlex.quote(video_in)} -vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18,PrimaryColour=&H00FFFF'\" -c:v libx264 -preset ultrafast -c:a aac {out_path}"
+        cmd = f"ffmpeg -y -threads 0 -i {shlex.quote(video_in)} -vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18,PrimaryColour=&H00FFFF'\" -c:v libx264 -preset superfast -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:
         yield f"❌ Erreur : {str(e)}", None
     finally:
         if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir)
 
-# 5. INTERFACE GRADIO
+# 4. INTERFACE GRADIO
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.HTML("<div style='text-align:center;'><h1>🤖 RobotsMali Speech Laboratory</h1><p>Testez nos modèles de transcription et traduction</p></div>")
+    gr.HTML("<div style='text-align:center;'><h1>🤖 RobotsMali Speech Lab</h1></div>")
     
     with gr.Row():
         with gr.Column():
-            v_input = gr.Video(label="Vidéo")
+            v_input = gr.Video(label="Source")
             m_input = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Modèle")
-            run_btn = gr.Button("🚀 GÉNÉRER", variant="primary")
+            run_btn = gr.Button("🚀 LANCER LA TRANSCRIPTION", variant="primary")
             
-            # --- AFFICHAGE DE L'EXEMPLE SI TROUVÉ ---
             if EXAMPLE_PATH:
-                gr.Examples(
-                    examples=[[EXAMPLE_PATH, "Soloba V3 (CTC)"]],
-                    inputs=[v_input, m_input],
-                    label="Vidéo d'exemple"
-                )
+                gr.Examples(examples=[[EXAMPLE_PATH, "Soloba V3 (CTC)"]], inputs=[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])