Update app.py

app.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """ 
-ROBOTSMALI — Sous-titrage Bambara (V5.6 - Production Final)
-Logiciel de transcription et d'incrustation vidéo pour le Bambara.
+ROBOTSMALI — Sous-titrage Bambara (VERSION INTÉGRALE V6.0)
+Incrustation de sous-titres avec tous les modèles RobotsMali.
 """
 import os
 import shlex
@@ -10,6 +10,7 @@ import tempfile
 import traceback
 import random
 import textwrap
+import time
 from pathlib import Path
 
 import numpy as np
@@ -20,12 +21,13 @@ from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 import gradio as gr
 
-# ---------------------------- # CONFIGURATION IA # ----------------------------
+# ---------------------------- # CONFIGURATION # ----------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 random.seed(1234)
 np.random.seed(1234)
 torch.manual_seed(1234)
 
+# TOUS VOS MODÈLES SONT ICI
 MODELS = {
     "Soloni V1 (RNNT)":        ("RobotsMali/soloni-114m-tdt-ctc-v1", "rnnt"),
     "Soloni V0 (RNNT)":        ("RobotsMali/soloni-114m-tdt-ctc-v0", "rnnt"),
@@ -35,7 +37,7 @@ MODELS = {
     "QuartzNet V0 (CTC-char)": ("RobotsMali/stt-bm-quartznet15x5-v0", "ctc_char"),
 }
 
-# Chemin vers la vidéo d'exemple
+# EXEMPLE CONFIGURÉ
 VIDEO_EXAMPLES = [
     ["examples/MARALINKE.mp4", "Soloba V1 (CTC)"]
 ]
@@ -45,9 +47,10 @@ _cache = {}
 # ---------------------------- # FONCTIONS TECHNIQUES # ----------------------------
 
 def run_cmd(cmd):
+    """Exécute une commande système."""
     res = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
     if res.returncode != 0:
-        raise RuntimeError(f"Erreur système: {res.stdout}")
+        raise RuntimeError(f"Erreur FFmpeg: {res.stdout}")
     return res.stdout
 
 def ffprobe_duration(path):
@@ -57,7 +60,14 @@ def ffprobe_duration(path):
     except: return None
 
 def load_model(name):
+    """Charge le modèle sélectionné et nettoie le cache si nécessaire."""
     if name in _cache: return _cache[name]
+    
+    # Nettoyage pour économiser la RAM
+    if len(_cache) > 0:
+        _cache.clear()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+
     repo, mode = 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)
@@ -75,30 +85,17 @@ def load_model(name):
     return model
 
 def extract_audio(video_path, out_wav):
-    """Prépare la vidéo (H.264) et extrait l'audio 16kHz."""
+    """Stabilisation du codec (pour la webcam) et extraction audio."""
     tmp_fd, stabilized_mp4 = tempfile.mkstemp(suffix="_stabilized.mp4")
     os.close(tmp_fd)
-    # Réencodage pour supporter tous les formats (Webcam/WebM compris)
+    # On force le H.264 pour éviter les erreurs de lecture
     run_cmd(f'ffmpeg -hide_banner -loglevel error -y -i {shlex.quote(video_path)} -c:v libx264 -preset ultrafast -crf 23 -c:a aac {shlex.quote(stabilized_mp4)}')
     run_cmd(f'ffmpeg -hide_banner -loglevel error -y -i {shlex.quote(stabilized_mp4)} -vn -ac 1 -ar 16000 -f wav {shlex.quote(out_wav)}')
     if os.path.exists(stabilized_mp4): os.remove(stabilized_mp4)
 
-def clean_audio(wav_path):
-    audio, sr = sf.read(wav_path)
-    if audio.ndim == 2: audio = audio.mean(axis=1)
-    if sr != 16000:
-        audio = librosa.resample(audio.astype(float), orig_sr=sr, target_sr=16000)
-    max_val = np.max(np.abs(audio)) if audio.size > 0 else 0.0
-    if max_val > 1e-6: audio = audio / max_val * 0.9
-    clean_path = wav_path.replace(".wav", "_clean.wav")
-    sf.write(clean_path, audio, 16000)
-    return clean_path, audio, 16000
-
-# ---------------------------- # TRANSCRIPTION & BURNING # ----------------------------
-
 def burn_subtitles(video_path, words, duration):
-    """Crée un fichier SRT et l'incruste dans la vidéo."""
-    out_path = "RobotsMali_Subtitled.mp4"
+    """Génère le fichier SRT et l'incruste dans la vidéo finale."""
+    out_path = f"output_{int(time.time())}.mp4"
     chunk_size = 7
     with tempfile.NamedTemporaryFile(suffix=".srt", mode="w", encoding="utf-8", delete=False) as tf:
         for i, idx in enumerate(range(0, len(words), chunk_size)):
@@ -114,51 +111,55 @@ def burn_subtitles(video_path, words, duration):
             tf.write(f"{i+1}\n{t_srt(start)} --> {t_srt(end)}\n{txt}\n\n")
         srt_name = tf.name
 
+    # Encodage ultra-rapide pour éviter le timeout
     vf = f"subtitles={shlex.quote(srt_name)}:force_style='Fontsize=22,PrimaryColour=&HFFFFFF&,OutlineColour=&H000000&'"
-    run_cmd(f'ffmpeg -hide_banner -loglevel error -y -i {shlex.quote(video_path)} -vf {shlex.quote(vf)} -c:v libx264 -crf 23 -c:a aac {shlex.quote(out_path)}')
+    run_cmd(f'ffmpeg -hide_banner -loglevel error -y -i {shlex.quote(video_path)} -vf {shlex.quote(vf)} -c:v libx264 -preset ultrafast -crf 28 -c:a copy {shlex.quote(out_path)}')
     os.remove(srt_name)
     return out_path
 
+# ---------------------------- # PIPELINE # ----------------------------
+
 def pipeline(video_input, model_name):
     try:
-        if not video_input: return "❌ Aucune vidéo", None
+        if not video_input: return "❌ Veuillez charger une vidéo", None
         video_path = video_input
         
-        yield "⏳ Phase 1/3 : Analyse du fichier...", None
+        yield "⏳ Phase 1/3 : Stabilisation & Audio...", None
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tf:
             wav_path = tf.name
         
         extract_audio(video_path, wav_path)
-        clean_wav, audio, sr = clean_audio(wav_path)
-        duration = ffprobe_duration(video_path) or (len(audio)/sr)
+        duration = ffprobe_duration(video_path) or 10.0 # fallback
         
-        yield f"⏳ Phase 2/3 : Transcription IA ({model_name})...", None
+        yield f"⏳ Phase 2/3 : Analyse IA ({model_name})...", None
         model = load_model(model_name)
-        res = model.transcribe([clean_wav])[0]
+        res = model.transcribe([wav_path])[0]
         text = res.text if hasattr(res, 'text') else str(res)
         words = [w for w in text.split() if len(w) > 1]
         
-        if not words: return "⚠️ Pas de parole détectée", None
+        if not words: return "⚠️ Pas de parole détectée.", None
         
-        yield "⏳ Phase 3/3 : Incrustation des sous-titres...", None
-        final_video = burn_subtitles(video_path, words, duration)
+        yield "⏳ Phase 3/3 : Génération des sous-titres...", None
+        final_v = burn_subtitles(video_path, words, duration)
         
-        yield "✅ Succès !", final_video
+        if os.path.exists(wav_path): os.remove(wav_path)
+        yield "✅ Sous-titrage terminé !", final_v
+
     except Exception as e:
         traceback.print_exc()
-        yield f"❌ Erreur: {str(e)}", None
+        yield f"❌ Erreur critique : {str(e)}", None
 
-# ---------------------------- # INTERFACE GRADIO # ----------------------------
+# ---------------------------- # INTERFACE ARTISTIQUE # ----------------------------
 
 custom_css = """
 body { background-color: #0b0e14; }
 .gradio-container { background: rgba(17, 25, 40, 0.8) !important; backdrop-filter: blur(12px); border-radius: 20px; border: 1px solid rgba(255, 255, 255, 0.1); }
-#title-container { text-align: center; padding: 20px; }
+#title-header { text-align: center; padding: 20px; }
 .gr-button-primary { background: linear-gradient(135deg, #059669, #10b981) !important; border: none !important; }
 """
 
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    with gr.Column(elem_id="title-container"):
+    with gr.Column(elem_id="title-header"):
         gr.HTML("""
             <h1 style='color:#facc15; font-size: 2.5rem; margin:0;'>🤖 ROBOTSMALI</h1>
             <p style='color:#94a3b8; font-style:italic;'>Intelligence Artificielle pour le Bambara</p>
@@ -167,21 +168,21 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### 📥 Source")
+            gr.Markdown("### 📥 Source Vidéo")
             v_in = gr.Video(label=None, mirror_webcam=False)
             m_sel = gr.Dropdown(list(MODELS.keys()), value="Soloba V1 (CTC)", label="Modèle IA")
             btn = gr.Button("🚀 GÉNÉRER LES SOUS-TITRES", variant="primary")
             
         with gr.Column():
             gr.Markdown("### 📤 Résultat")
-            status = gr.Markdown("*En attente de traitement...*")
+            status = gr.Markdown("*Prêt pour le traitement...*")
             v_out = gr.Video(label=None)
 
-    # Section des exemples avec cache_examples=False pour débloquer le clic
+    # EXEMPLES : cache_examples=False est crucial pour que le clic fonctionne
     gr.Examples(
         examples=VIDEO_EXAMPLES,
         inputs=[v_in, m_sel],
-        label="📺 Exemples Disponibles",
+        label="📺 Vidéo d'exemple",
         cache_examples=False
     )