Update app.py

app.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-import os, shlex, subprocess, tempfile, traceback, time, glob, gc
+import os, shlex, subprocess, tempfile, traceback, time, glob, gc, shutil
 import torch
 from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
@@ -30,7 +30,7 @@ def get_absolute_example():
 EXAMPLE_PATH = get_absolute_example()
 _cache = {}
 
-# 3. GESTION DE LA MÉMOIRE ET CHARGEMENT
+# 3. GESTION DE LA MÉMOIRE ET CHARGEMENT (CORRIGÉ)
 def clear_memory():
     """Libère proprement la RAM et la VRAM."""
     _cache.clear()
@@ -39,22 +39,26 @@ def clear_memory():
         torch.cuda.empty_cache()
 
 def get_model(name):
-    """Charge le modèle avec optimisation FP16 pour la vitesse."""
+    """Charge le modèle en utilisant ASRModel pour éviter les erreurs de state_dict."""
     if name in _cache: return _cache[name]
     
     clear_memory()
-    repo, mode = MODELS[name]
+    repo, _ = MODELS[name]
+    
+    print(f"📥 Téléchargement depuis {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 mode == "rnnt":
-        model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(nemo_file)
-    else:
-        model = nemo_asr.models.EncDecCTCModelBPE.restore_from(nemo_file)
+    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)
     
     model.to(DEVICE).eval()
     
-    # OPTIMISATION : Inférence en demi-précision (FP16) sur GPU
+    # Optimisation FP16
     if DEVICE == "cuda":
         model.half()
         
@@ -67,18 +71,18 @@ def format_srt_time(sec):
     ms = int((sec - int(sec)) * 1000)
     return f"{time.strftime('%H:%M:%S', td)},{ms:03}"
 
-# 5. PIPELINE DE TRANSCRIPTION OPTIMISÉ
+# 5. PIPELINE DE TRANSCRIPTION
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
         if not video_in: return "❌ Source vide", None
         
-        # A. Extraction Audio Rapide
+        # 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 -threads 0 {full_wav}", shell=True, check=True)
+        subprocess.run(f"ffmpeg -y -i {shlex.quote(video_in)} -vn -ac 1 -ar 16000 {full_wav}", shell=True, check=True)
         
-        # B. Segmentation (20s pour réduire le nombre d'appels IA)
+        # B. Segmentation
         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)
@@ -88,20 +92,20 @@ def pipeline(video_in, model_name):
         yield f"⏳ IA : Chargement de {model_name}...", None
         model = get_model(model_name)
         
-        # Détermination du stride (standard RobotsMali 0.02)
         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 Séquentielle
+        # 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
             
-            # Inférence
             hyp = model.transcribe([seg_path], return_hypotheses=True)[0]
-            words = (hyp.text if hasattr(hyp, 'text') else str(hyp)).split()
+            # Gestion des différents formats de retour NeMo
+            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):
@@ -109,7 +113,6 @@ def pipeline(video_in, model_name):
                     t_start = base_time + (offsets[i] * stride)
                     all_words_ts.append({"word": word, "start": t_start, "end": t_start + 0.45})
             else:
-                # 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)})
@@ -122,18 +125,15 @@ 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 (Ultra-rapide)
-        yield "⏳ Rendu vidéo (Ultra-rapide)...", None
+        # F. Encodage Vidéo Final
+        yield "⏳ Rendu vidéo final...", 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,BorderStyle=4'\" "
-            f"-c:v libx264 -preset ultrafast -pix_fmt yuv420p -movflags +faststart -c:a aac {out_path}"
+            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)
         
@@ -142,19 +142,22 @@ def pipeline(video_in, model_name):
     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)
 
 # 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.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(label="Source Vidéo", sources=["upload", "webcam"])
+            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.Markdown("### 💡 Exemple Rapide")
                 gr.Examples(examples=[[EXAMPLE_PATH, "Soloba V3 (CTC)"]], inputs=[v_in, m_sel])
 
         with gr.Column():
@@ -164,4 +167,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css="body {background-color: #0f172a;}")
     btn_run.click(pipeline, [v_in, m_sel], [status, v_out])
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch()