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
 
-# CONFIGURATION
+# 1. CONFIGURATION
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 MODELS = {
@@ -23,35 +23,28 @@ MODELS = {
     "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "rnnt"),
 }
 
-_cache = {}
+# Recherche de la vidéo d'exemple dans le dossier courant
+EXAMPLE_VIDEO = "MARALINKE.mp4" if os.path.exists("MARALINKE.mp4") else None
 
-def clear_memory():
-    _cache.clear()
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+_cache = {}
 
 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()
     
-    print(f"📥 Téléchargement de {repo}...")
+    repo, _ = 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)
     
-    # Chargement robuste pour éviter l'erreur Unexpected Key
     from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
     model = nemo_asr.models.ASRModel.restore_from(
         nemo_file, 
         map_location=torch.device(DEVICE),
         save_restore_connector=SaveRestoreConnector()
     )
-    
     model.to(DEVICE).eval()
-    if DEVICE == "cuda":
-        model.half()
-        
+    if DEVICE == "cuda": model.half()
     _cache[name] = model
     return model
 
@@ -63,34 +56,32 @@ def format_srt_time(sec):
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: return "❌ Vidéo manquante", None
+        if not video_in: return "❌ Source vide", None
         
-        yield "⏳ Extraction audio...", None
+        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)
         
-        yield "⏳ Segmentation...", None
-        seg_time = 20
-        subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time {seg_time} -c copy {os.path.join(tmp_dir, 'seg_%03d.wav')}", shell=True, check=True)
+        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")))
         
-        yield f"⏳ IA : Chargement {model_name}...", None
+        yield f"⏳ Phase 3 : IA ({model_name})...", None
         model = get_model(model_name)
         
         all_words_ts = []
         for idx, seg_path in enumerate(audio_segments):
-            base_time = idx * seg_time
-            yield f"⏳ IA : Transcription {idx+1}/{len(audio_segments)}...", None
+            base_time = idx * 20
+            yield f"🎙️ Transcription {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()
-            # On simule un timing si les offsets manquent
-            gap = float(seg_time) / max(len(words), 1)
+            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)})
 
+        yield "⏳ Phase 4 : Rendu Vidéo...", None
         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):
@@ -98,10 +89,9 @@ 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")
 
-        yield "⏳ Rendu vidéo...", None
-        out_path = os.path.abspath(f"resultat_{int(time.time())}.mp4")
+        out_path = os.path.abspath(f"output_{int(time.time())}.mp4")
         safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
-        subprocess.run(f"ffmpeg -y -i {shlex.quote(video_in)} -vf \"subtitles='{safe_srt}'\" -c:v libx264 -preset ultrafast -c:a aac {out_path}", shell=True, check=True)
+        subprocess.run(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 aac {out_path}", shell=True, check=True)
         
         yield "✅ Terminé !", out_path
     except Exception as e:
@@ -110,16 +100,27 @@ def pipeline(video_in, model_name):
         if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir)
 
 # INTERFACE
-with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 RobotsMali ASR Test")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.HTML("<h1 style='text-align:center;'>🤖 RobotsMali Transcription</h1>")
+    
     with gr.Row():
         with gr.Column():
-            v_in = gr.Video()
-            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)")
-            btn = gr.Button("Lancer", variant="primary")
+            v_in = gr.Video(label="Vidéo Source")
+            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Modèle")
+            btn = gr.Button("GÉNÉRER", variant="primary")
+            
+            # REVOICI L'EXEMPLE ICI
+            if EXAMPLE_VIDEO:
+                gr.Examples(
+                    examples=[[EXAMPLE_VIDEO, "Soloba V3 (CTC)"]],
+                    inputs=[v_in, m_sel],
+                    label="Exemple disponible"
+                )
+
         with gr.Column():
             status = gr.Markdown("Prêt.")
-            v_out = gr.Video()
+            v_out = gr.Video(label="Résultat")
+
     btn.click(pipeline, [v_in, m_sel], [status, v_out])
 
-demo.launch(debug=True, share=True)
+demo.launch()