Update app.py

app.py

@@ -5,9 +5,8 @@ from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# 1. CONFIGURATION ET MODÈLES
+# 1. CONFIGURATION
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
 MODELS = {
     "Soloba V3 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
     "Soloba V2 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v2", "ctc"),
@@ -23,7 +22,6 @@ MODELS = {
     "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "rnnt"),
 }
 
-# Détection de la vidéo d'exemple dans le dossier local
 def find_example_video():
     paths = ["examples/MARALINKE_FIXED.mp4", "examples/MARALINKE.mp4", "MARALINKE.mp4"]
     for p in paths:
@@ -33,64 +31,65 @@ def find_example_video():
 EXAMPLE_PATH = find_example_video()
 _cache = {}
 
-# 2. OPTIMISATION CHARGEMENT ET MÉMOIRE
-def clear_memory():
-    _cache.clear()
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
+# 2. GESTION MÉMOIRE
 def get_model(name):
     if name in _cache: return _cache[name]
-    clear_memory()
-    repo, _ = MODELS[name]
+    _cache.clear()
+    gc.collect()
+    if torch.cuda.is_available(): torch.cuda.empty_cache()
     
+    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)
     
     from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
     model = nemo_asr.models.ASRModel.restore_from(
-        nemo_file, 
-        map_location=torch.device(DEVICE),
+        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
 
-# 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}"
 
+# 3. PIPELINE AVEC GÉNÉRATEUR (POUR ÉVITER LE TIMEOUT)
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: return "❌ Vidéo absente", None
-        
-        # FFmpeg optimisé (-threads 0)
-        yield "🚀 Phase 1 : Extraction Audio...", None
+        if not video_in:
+            yield "❌ Erreur : Aucune vidéo détectée.", None
+            return
+
+        # Phase 1 : Extraction Audio
+        yield "⏳ Phase 1/4 : Extraction de l'audio (FFmpeg)...", None
         full_wav = os.path.join(tmp_dir, "full.wav")
         subprocess.run(f"ffmpeg -y -threads 0 -i {shlex.quote(video_in)} -vn -ac 1 -ar 16000 {full_wav}", shell=True, check=True)
-        
-        yield "🚀 Phase 2 : Segmentation...", None
+
+        # Phase 2 : Découpage
+        yield "⏳ Phase 2/4 : Analyse et découpage des segments...", 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")))
-        
-        # Inférence par BATCH (beaucoup plus rapide)
-        yield f"🎙️ Phase 3 : Transcription Batch avec {model_name}...", None
+
+        # Phase 3 : Transcription (C'est ici que ça coupe d'habitude)
+        yield f"⏳ Phase 3/4 : Chargement de l'IA {model_name}...", None
         model = get_model(model_name)
         
-        # 
-        batch_hypotheses = model.transcribe(audio_segments, batch_size=4, return_hypotheses=True)
+        yield f"🎙️ Transcription de {len(audio_segments)} segments en cours...", None
+        
+        # Inférence Batch pour gagner du temps
+        # On réduit le batch_size si on est sur CPU pour éviter de saturer la RAM
+        b_size = 2 if DEVICE == "cpu" else 4
+        batch_hypotheses = model.transcribe(audio_segments, batch_size=b_size, return_hypotheses=True)
         
         all_words_ts = []
         for idx, hyp in enumerate(batch_hypotheses):
+            # On envoie un signal Gradio à chaque segment traité pour garder la connexion vivante
+            yield f"📝 Traitement du texte : {idx+1}/{len(audio_segments)} segments finis...", None
             base_time = idx * 20
             if isinstance(hyp, list): hyp = hyp[0]
             text = hyp.text if hasattr(hyp, 'text') else str(hyp)
@@ -99,7 +98,8 @@ def pipeline(video_in, model_name):
             for i, w in enumerate(words):
                 all_words_ts.append({"word": w, "start": base_time + (i * gap), "end": base_time + ((i+1) * gap)})
 
-        # Génération SRT
+        # Phase 4 : Rendu Final
+        yield "⏳ Phase 4/4 : Incrustation des sous-titres dans la 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):
@@ -107,38 +107,39 @@ 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")
 
-        # 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")
+        out_path = os.path.abspath(f"resultat_{int(time.time())}.mp4")
         safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
         
-        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}"
+        # Encodage ultra-rapide
+        cmd = f"ffmpeg -y -threads 0 -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 "✅ Succès !", out_path
+        yield "✅ Succès ! La vidéo est prête ci-dessous.", out_path
 
     except Exception as e:
-        yield f"❌ Erreur : {str(e)}", None
+        yield f"❌ Erreur critique : {str(e)}", None
     finally:
         if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir)
 
-# 4. INTERFACE GRADIO
+# 4. INTERFACE GRADIO AVEC TIMEOUT ÉTENDU
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.HTML("<div style='text-align:center;'><h1>🤖 RobotsMali Speech Lab</h1></div>")
+    gr.HTML("<h1 style='text-align:center;'>🤖 RobotsMali Transcription Pro</h1>")
     
     with gr.Row():
         with gr.Column():
-            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("🚀 LANCER LA TRANSCRIPTION", variant="primary")
+            v_input = gr.Video(label="Vidéo (MP4)")
+            m_input = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Modèle IA")
+            run_btn = gr.Button("🚀 LANCER", variant="primary")
             
             if EXAMPLE_PATH:
                 gr.Examples(examples=[[EXAMPLE_PATH, "Soloba V3 (CTC)"]], inputs=[v_input, m_input])
 
         with gr.Column():
-            status = gr.Markdown("Prêt.")
-            v_output = gr.Video(label="Résultat")
+            status = gr.Markdown("### État\nPrêt.")
+            v_output = gr.Video(label="Vidéo finale")
 
+    # Important : on utilise queue() pour gérer les longues attentes
     run_btn.click(pipeline, [v_input, m_input], [status, v_output])
 
-demo.launch()
+# Configuration de la file d'attente (Queue) pour éviter les coupures HTTP
+demo.queue(max_size=10).launch(show_error=True)