Update app.py

app.py

@@ -5,10 +5,9 @@ from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# 1. CONFIGURATION MATÉRIEL
+# 1. CONFIGURATION MATÉRIEL ET MODÈLES
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 2. DICTIONNAIRE DES MODÈLES
 MODELS = {
     "Soloba V3 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
     "Soloba V1.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
@@ -18,19 +17,30 @@ MODELS = {
     "Soloba V0.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v0.5", "rnnt"),
 }
 
+# 2. LOCALISATION DE LA VIDÉO D'EXEMPLE
+def get_absolute_example():
+    names = ["MARALINKE.mp4", "maralinke.mp4", "example.mp4"]
+    dirs = [".", "examples", "/home/user/app", "/home/user/app/examples"]
+    for d in dirs:
+        for n in names:
+            p = os.path.join(d, n)
+            if os.path.exists(p): return os.path.abspath(p)
+    return None
+
+EXAMPLE_PATH = get_absolute_example()
 _cache = {}
 
+# 3. GESTION DE LA MÉMOIRE ET CHARGEMENT
 def clear_memory():
-    """Libère la VRAM et la RAM."""
+    """Libère proprement la RAM et la VRAM."""
     _cache.clear()
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
 def get_model(name):
-    """Charge le modèle et le retourne directement."""
-    if name in _cache: 
-        return _cache[name]
+    """Charge le modèle avec optimisation FP16 pour la vitesse."""
+    if name in _cache: return _cache[name]
     
     clear_memory()
     repo, mode = MODELS[name]
@@ -43,104 +53,113 @@ def get_model(name):
         model = nemo_asr.models.EncDecCTCModelBPE.restore_from(nemo_file)
     
     model.to(DEVICE).eval()
+    
+    # OPTIMISATION : Inférence en demi-précision (FP16) sur GPU
     if DEVICE == "cuda":
         model.half()
-    
+        
     _cache[name] = model
     return model
 
+# 4. UTILITAIRE TEMPOREL
 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}"
 
+# 5. PIPELINE DE TRANSCRIPTION OPTIMISÉ
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: 
-            return "❌ Source vide", None
+        if not video_in: return "❌ Source vide", None
         
-        # A. Extraction Audio
-        yield "⏳ Extraction de l'audio...", None
+        # A. Extraction Audio Rapide
+        yield "⏳ 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 -i {shlex.quote(video_in)} -vn -ac 1 -ar 16000 -threads 0 {full_wav}", shell=True, check=True)
         
-        # B. Segmentation (10s)
+        # B. Segmentation (20s pour réduire le nombre d'appels IA)
+        seg_time = 20
         segment_pattern = os.path.join(tmp_dir, "seg_%03d.wav")
-        subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time 10 -c copy {segment_pattern}", shell=True, check=True)
+        subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time {seg_time} -c copy {segment_pattern}", shell=True, check=True)
         audio_segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
         
-        # C. Chargement du modèle
-        yield f"⏳ Chargement du modèle {model_name}...", None
+        # C. Initialisation Modèle
+        yield f"⏳ IA : Chargement de {model_name}...", None
         model = get_model(model_name)
-            
-        # Facteur de temps dynamique
-        stride = 0.02 
+        
+        # Détermination du stride (standard RobotsMali 0.02)
+        stride = 0.02
         if hasattr(model, 'preprocessor') and hasattr(model.preprocessor, 'featurizer'):
-            hop = model.preprocessor.featurizer.hop_length
-            sr = model.preprocessor.featurizer.sample_rate
-            stride = hop / sr
+            stride = model.preprocessor.featurizer.hop_length / model.preprocessor.featurizer.sample_rate
 
-        # D. Transcription
+        # D. Transcription Séquentielle
         all_words_ts = []
         for idx, seg_path in enumerate(audio_segments):
-            base_time = idx * 10.0
+            base_time = idx * seg_time
             yield f"⏳ IA : Transcription {idx+1}/{len(audio_segments)}...", None
             
+            # Inférence
             hyp = model.transcribe([seg_path], return_hypotheses=True)[0]
+            words = (hyp.text if hasattr(hyp, 'text') else str(hyp)).split()
             offsets = getattr(hyp, 'word_offsets', None)
-            words = hyp.text.split() if hasattr(hyp, 'text') else str(hyp).split()
 
             if offsets and len(offsets) == len(words):
                 for i, word in enumerate(words):
-                    start_t = base_time + (offsets[i] * stride)
-                    all_words_ts.append({"word": word, "start": start_t, "end": start_t + 0.45})
+                    t_start = base_time + (offsets[i] * stride)
+                    all_words_ts.append({"word": word, "start": t_start, "end": t_start + 0.45})
             else:
-                gap = 10.0 / max(len(words), 1)
+                # Fallback linéaire si les offsets sont indisponibles
+                gap = float(seg_time) / 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)})
 
-        # E. Génération SRT (Fichier temporaire interne)
+        # E. Génération du SRT
         srt_path = os.path.join(tmp_dir, "final.srt")
-        words_per_line = 6
         with open(srt_path, "w", encoding="utf-8") as f:
-            for i in range(0, len(all_words_ts), words_per_line):
-                chunk = all_words_ts[i:i+words_per_line]
-                f.write(f"{(i//words_per_line)+1}\n")
-                f.write(f"{format_srt_time(chunk[0]['start'])} --> {format_srt_time(chunk[-1]['end'])}\n")
+            for i in range(0, len(all_words_ts), 6):
+                chunk = all_words_ts[i:i+6]
+                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")
 
-        # F. Encodage Vidéo
-        yield "⏳ Rendu vidéo final...", None
+        # F. Encodage Vidéo Final (Ultra-rapide)
+        yield "⏳ Rendu vidéo (Ultra-rapide)...", None
         out_path = os.path.abspath(f"robotsmali_final_{int(time.time())}.mp4")
         
+        # Protection des chemins pour FFmpeg
         safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
+        
+        # OPTIMISATION : -preset ultrafast pour minimiser le temps de rendu
         cmd_ffmpeg = (
             f"ffmpeg -y -i {shlex.quote(video_in)} "
-            f"-vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18,OutlineColour=&H80000000,BorderStyle=4'\" "
-            f"-c:v libx264 -preset fast -pix_fmt yuv420p -movflags +faststart -c:a aac {out_path}"
+            f"-vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18,BorderStyle=4'\" "
+            f"-c:v libx264 -preset ultrafast -pix_fmt yuv420p -movflags +faststart -c:a aac {out_path}"
         )
         subprocess.run(cmd_ffmpeg, shell=True, check=True)
         
-        yield "✅ Terminé !", out_path
+        yield "✅ Terminé avec succès !", out_path
 
     except Exception as e:
         traceback.print_exc()
         yield f"❌ Erreur : {str(e)}", None
 
-# 6. INTERFACE
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.HTML("<h1 style='text-align:center; color:#EAB308;'>🤖 ROBOTSMALI TRANSCRIPTION</h1>")
+# 6. INTERFACE GRADIO
+with gr.Blocks(theme=gr.themes.Soft(), css="body {background-color: #0f172a;}") as demo:
+    gr.HTML("<h1 style='text-align:center; color:#facc15;'>🤖 ROBOTSMALI TRANSCRIPTION</h1>")
     
     with gr.Row():
         with gr.Column():
-            v_in = gr.Video(label="Source")
+            v_in = gr.Video(label="Source Vidéo", sources=["upload", "webcam"])
             m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Modèle IA")
-            btn_run = gr.Button("🚀 GÉNÉRER", variant="primary")
+            btn_run = gr.Button("🚀 GÉNÉRER SOUS-TITRES", variant="primary")
             
+            if EXAMPLE_PATH:
+                gr.Markdown("### 💡 Exemple Rapide")
+                gr.Examples(examples=[[EXAMPLE_PATH, "Soloba V3 (CTC)"]], inputs=[v_in, m_sel])
+
         with gr.Column():
             status = gr.Markdown("### État\nPrêt.")
-            v_out = gr.Video(label="Vidéo Finale")
+            v_out = gr.Video(label="Résultat Final")
 
     btn_run.click(pipeline, [v_in, m_sel], [status, v_out])