Update app.py

app.py

@@ -1,40 +1,38 @@
 # -*- coding: utf-8 -*-
-import os, shlex, subprocess, tempfile, traceback, textwrap, time, glob
+import os, shlex, subprocess, tempfile, traceback, time, glob, gc
 import torch
 from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# 1. CONFIGURATION DU MATÉRIEL ET DES MODÈLES
+# 1. CONFIGURATION MATÉRIEL
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# 2. DICTIONNAIRE DES MODÈLES (Mis à jour selon votre capture d'écran)
 MODELS = {
-    "Soloba V1 (CTC)":         ("RobotsMali/soloba-ctc-0.6b-v1", "ctc"),           
-    "Soloni V1 (RNNT)":        ("RobotsMali/soloni-114m-tdt-ctc-v1", "rnnt"),
-    "QuartzNet V1 (CTC-char)": ("RobotsMali/stt-bm-quartznet15x5-v1", "ctc_char"),       
-    "Soloba V0 (CTC)":         ("RobotsMali/soloba-ctc-0.6b-v0", "ctc"),           
-    "Soloni V0 (RNNT)":        ("RobotsMali/soloni-114m-tdt-ctc-v0", "rnnt"),
-    "QuartzNet V0 (CTC-char)": ("RobotsMali/stt-bm-quartznet15x5-v0", "ctc_char"),
+    "Soloba V3 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
+    "Soloba V1.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
+    "Soloni V3 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v3", "rnnt"),
+    "Soloni V2 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v2", "rnnt"),
+    "Soloni MSE (Experimental)": ("RobotsMali/lau-soloni-114m-mse-k1", "ctc"),
+    "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. CHARGEMENT DES MODÈLES IA
+def clear_memory():
+    """Nettoie la VRAM et la RAM pour éviter les débordements."""
+    _cache.clear()
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
 def load_model(name):
+    """Charge le modèle avec optimisation FP16 si possible."""
     if name in _cache: return _cache[name]
-    _cache.clear()
-    if torch.cuda.is_available(): torch.cuda.empty_cache()
+    
+    yield f"⏳ Chargement du modèle {name}..."
+    clear_memory()
     
     repo, mode = MODELS[name]
     folder = snapshot_download(repo, local_dir_use_symlinks=False)
@@ -42,71 +40,71 @@ def load_model(name):
     
     if mode == "rnnt":
         model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(nemo_file)
-    elif mode == "ctc_char":
-        model = nemo_asr.models.EncDecCTCModel.restore_from(nemo_file)
     else:
         model = nemo_asr.models.EncDecCTCModelBPE.restore_from(nemo_file)
     
     model.to(DEVICE).eval()
+    if DEVICE == "cuda":
+        model.half() # Utilisation de la demi-précision pour gagner 50% de VRAM
+    
     _cache[name] = model
     return model
 
-# 4. UTILITAIRE DE FORMATAGE SRT
 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 TRAITEMENT (SEGMENTATION 10S + OFFSETS)
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: return "❌ Erreur : Source vide", None
+        if not video_in: return "❌ Source vide", None, None
         
-        # Étape A : Extraction et Segmentation Audio
-        yield "⏳ Découpage de l'audio en segments de 10s...", None
+        # A. Extraction Audio
+        yield "⏳ Extraction de l'audio...", None, 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)
         
+        # B. Segmentation Temporelle (10s)
         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)
-        
         audio_segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
-        model = load_model(model_name)
         
-        # Étape B : Transcription segmentée avec Offsets natifs
+        # C. Chargement et Calcul de la Précision Temporelle
+        model_gen = load_model(model_name)
+        model = None
+        for update in model_gen:
+            if isinstance(update, str): yield update, None, None
+            else: model = update
+            
+        # Calcul dynamique du stride (facteur de conversion frames -> secondes)
+        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
+
+        # D. Transcription par segments
         all_words_ts = []
         for idx, seg_path in enumerate(audio_segments):
             base_time = idx * 10.0
-            yield f"⏳ IA : Transcription segment {idx+1}/{len(audio_segments)}...", None
+            yield f"⏳ IA : Transcription segment {idx+1}/{len(audio_segments)}...", None, None
             
-            # Utilisation de return_hypotheses pour les timestamps
             hyp = model.transcribe([seg_path], return_hypotheses=True)[0]
-            
-            if hasattr(hyp, 'word_offsets') and hyp.word_offsets:
-                words = hyp.text.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):
-                    # Facteur de conversion frame->seconde (standard NeMo 0.02)
-                    rel_start = hyp.word_offsets[i] * 0.02
-                    all_words_ts.append({
-                        "word": word, 
-                        "start": base_time + rel_start, 
-                        "end": base_time + rel_start + 0.45
-                    })
+                    start_t = base_time + (offsets[i] * stride)
+                    all_words_ts.append({"word": word, "start": start_t, "end": start_t + 0.45})
             else:
-                # Fallback temporel si offsets non dispos
-                words = (hyp.text if hasattr(hyp, 'text') else str(hyp)).split()
-                if words:
-                    gap = 10.0 / 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)
-                        })
+                # Fallback : Répartition linéaire si les offsets manquent
+                gap = 10.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)})
 
-        # Étape C : Génération du SRT optimisé
-        yield "⏳ Création du fichier de sous-titres...", None
+        # E. Génération du fichier SRT
         srt_path = os.path.join(tmp_dir, "final.srt")
         words_per_line = 6
         with open(srt_path, "w", encoding="utf-8") as f:
@@ -116,43 +114,41 @@ def pipeline(video_in, model_name):
                 f.write(f"{format_srt_time(chunk[0]['start'])} --> {format_srt_time(chunk[-1]['end'])}\n")
                 f.write(" ".join([c['word'] for c in chunk]) + "\n\n")
 
-        # Étape D : Incrustation Finale (Burn-in)
-        yield "⏳ Rendu vidéo final...", None
+        # F. Encodage Vidéo avec Incrustation (Burn-in)
+        yield "⏳ Rendu vidéo final ...", None, srt_path
         out_path = os.path.abspath(f"robotsmali_final_{int(time.time())}.mp4")
+        
+        # Protection des chemins pour FFmpeg (indispensable pour Windows/Linux)
+        safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
         cmd_ffmpeg = (
             f"ffmpeg -y -i {shlex.quote(video_in)} "
-            f"-vf \"subtitles={shlex.quote(srt_path)}:force_style='Alignment=2,FontSize=20,PrimaryColour=&H00FFFF&'\" "
-            f"-c:v libx264 -pix_fmt yuv420p -movflags +faststart -c:a aac {out_path}"
+            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}"
         )
         subprocess.run(cmd_ffmpeg, shell=True, check=True)
         
-        yield "✅ Synchronisation parfaite terminée !", out_path
+        yield "✅ Transcription et Incrustation Terminées !", out_path, srt_path
 
     except Exception as e:
         traceback.print_exc()
-        yield f"❌ Erreur : {str(e)}", None
+        yield f"❌ Erreur : {str(e)}", None, None
 
 # 6. INTERFACE UTILISATEUR GRADIO
-with gr.Blocks(theme=gr.themes.Soft(), css="body {background-color: #0b1120;}") as demo:
-    gr.HTML("<h1 style='text-align:center; color:#facc15;'>🤖 ROBOTSMALI </h1>")
-    gr.Markdown("<p style='text-align:center; color:white;'>Segmentation NeMo</p>")
+with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
+    gr.HTML("<h1 style='text-align:center; color:#EAB308;'>🤖 ROBOTSMALI TRANSCRIPTION PRO</h1>")
     
     with gr.Row():
-        with gr.Column():
-            v_in = gr.Video(label="Source (Webcam ou Fichier)", sources=["upload", "webcam"], interactive=True)
-            
-            m_sel = gr.Dropdown(list(MODELS.keys()), value="Soloba V1 (CTC)", label="Modèle IA")
-            btn_run = gr.Button("🚀 GÉNÉRER SOUS-TITRES", variant="primary")
+        with gr.Column(scale=1):
+            v_in = gr.Video(label="Vidéo Source (Upload ou Webcam)", sources=["upload", "webcam"])
+            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Choisir le Modèle IA")
+            btn_run = gr.Button("🚀 GÉNÉRER LES SOUS-TITRES", variant="primary")
             
-            if EXAMPLE_PATH:
-                gr.Markdown("### 💡 Exemple")
-                gr.Examples(examples=[[EXAMPLE_PATH, "Soloba V1 (CTC)"]], inputs=[v_in, m_sel])
-
-        with gr.Column():
-            status = gr.Markdown("### État\nPrêt")
-            v_out = gr.Video(label="Résultat Final")
+        with gr.Column(scale=1):
+            status = gr.Markdown("### État\nPrêt à l'emploi.")
+            v_out = gr.Video(label="Vidéo Finale Incrustée")
+            f_srt = gr.File(label="Fichier Sous-titres (.SRT)")
 
-    btn_run.click(pipeline, [v_in, m_sel], [status, v_out])
+    btn_run.click(pipeline, [v_in, m_sel], [status, v_out, f_srt])
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True, show_error=True)