Update app.py

app.py

@@ -5,119 +5,92 @@ from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# 1. CONFIGURATION MATÉRIEL ET MODÈLES
+# 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"),
+    "Soloba V1 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v1", "ctc"),
+    "Soloba V0 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v0", "ctc"),
     "Soloba V1.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
+    "Soloba V0.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v0.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 V1 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v1", "rnnt"),
+    "Soloni V0 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v0", "rnnt"),
     "Soloni MSE (Experimental)": ("RobotsMali/lau-soloni-114m-mse-k1", "ctc"),
-    "Soloba V0.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v0.5", "rnnt"),
+    "Traduction Soloni (ST)":    ("RobotsMali/st-soloni-114m-tdt-ctc", "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 (CORRIGÉ)
 def clear_memory():
-    """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 en utilisant ASRModel pour éviter les erreurs de state_dict."""
     if name in _cache: return _cache[name]
-    
     clear_memory()
     repo, _ = MODELS[name]
     
-    print(f"📥 Téléchargement depuis {repo}...")
+    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(f"Aucun fichier .nemo trouvé dans {folder}")
-
-    # SOLUTION : Utilisation de ASRModel.restore_from pour la détection automatique
-    # Cela évite l'erreur 'Unexpected key(s) in state_dict'
-    model = nemo_asr.models.ASRModel.restore_from(nemo_file)
+    # 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()
-    
-    # Optimisation FP16
     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
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: return "❌ Source vide", None
+        if not video_in: return "❌ Vidéo manquante", None
         
-        # A. Extraction Audio
         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)
         
-        # B. Segmentation
+        yield "⏳ Segmentation...", None
         seg_time = 20
-        segment_pattern = os.path.join(tmp_dir, "seg_%03d.wav")
-        subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time {seg_time} -c copy {segment_pattern}", shell=True, check=True)
+        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)
         audio_segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
         
-        # C. Initialisation Modèle
-        yield f"⏳ IA : Chargement de {model_name}...", None
+        yield f"⏳ IA : Chargement {model_name}...", None
         model = get_model(model_name)
         
-        stride = 0.02
-        if hasattr(model, 'preprocessor') and hasattr(model.preprocessor, 'featurizer'):
-            stride = model.preprocessor.featurizer.hop_length / model.preprocessor.featurizer.sample_rate
-
-        # D. Transcription
         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
-            
             hyp = model.transcribe([seg_path], return_hypotheses=True)[0]
-            # Gestion des différents formats de retour NeMo
+            if isinstance(hyp, list): hyp = hyp[0]
+            
             text = hyp.text if hasattr(hyp, 'text') else str(hyp)
             words = text.split()
-            offsets = getattr(hyp, 'word_offsets', None)
-
-            if offsets and len(offsets) == len(words):
-                for i, word in enumerate(words):
-                    t_start = base_time + (offsets[i] * stride)
-                    all_words_ts.append({"word": word, "start": t_start, "end": t_start + 0.45})
-            else:
-                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)})
+            # On simule un timing si les offsets manquent
+            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 du 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):
@@ -125,46 +98,28 @@ 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")
 
-        # F. Encodage Vidéo Final
-        yield "⏳ Rendu vidéo final...", None
-        out_path = os.path.abspath(f"robotsmali_final_{int(time.time())}.mp4")
+        yield "⏳ Rendu vidéo...", None
+        out_path = os.path.abspath(f"resultat_{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)
         
-        cmd_ffmpeg = (
-            f"ffmpeg -y -i {shlex.quote(video_in)} "
-            f"-vf \"subtitles='{safe_srt}':force_style='Alignment=2,FontSize=18'\" "
-            f"-c:v libx264 -preset ultrafast -pix_fmt yuv420p -c:a aac {out_path}"
-        )
-        subprocess.run(cmd_ffmpeg, shell=True, check=True)
-        
-        yield "✅ Terminé avec succès !", out_path
-
+        yield "✅ Terminé !", out_path
     except Exception as e:
-        traceback.print_exc()
         yield f"❌ Erreur : {str(e)}", None
     finally:
-        # Nettoyage des fichiers temporaires
-        if os.path.exists(tmp_dir):
-            shutil.rmtree(tmp_dir)
+        if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir)
 
-# 6. INTERFACE GRADIO
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.HTML("<h1 style='text-align:center;'>🤖 ROBOTSMALI TRANSCRIPTION</h1>")
-    
+# INTERFACE
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 RobotsMali ASR Test")
     with gr.Row():
         with gr.Column():
-            v_in = gr.Video(label="Source Vidéo")
-            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Modèle IA")
-            btn_run = gr.Button("🚀 GÉNÉRER SOUS-TITRES", variant="primary")
-            
-            if EXAMPLE_PATH:
-                gr.Examples(examples=[[EXAMPLE_PATH, "Soloba V3 (CTC)"]], inputs=[v_in, m_sel])
-
+            v_in = gr.Video()
+            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)")
+            btn = gr.Button("Lancer", variant="primary")
         with gr.Column():
-            status = gr.Markdown("### État\nPrêt.")
-            v_out = gr.Video(label="Résultat Final")
-
-    btn_run.click(pipeline, [v_in, m_sel], [status, v_out])
+            status = gr.Markdown("Prêt.")
+            v_out = gr.Video()
+    btn.click(pipeline, [v_in, m_sel], [status, v_out])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(debug=True; share=True)