Update app.py

app.py

@@ -1,19 +1,27 @@
-import os, warnings, tempfile
+import os, warnings, tempfile, logging
+
+# --- Sécurité & compatibilité environnement ---
 warnings.filterwarnings("ignore")
+logging.getLogger("nemo_logger").setLevel(logging.ERROR)
+os.environ["NEMO_FORCE_CPU"] = "1"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+import torch
+torch.set_grad_enabled(False)
 
-# Débloquer ImageMagick (compatible tous environnements)
-for p in ["/etc/ImageMagick/policy.xml", "/etc/ImageMagick-6/policy.xml"]:
+# Débloquer ImageMagick (HF varie selon build)
+for p in ("/etc/ImageMagick/policy.xml", "/etc/ImageMagick-6/policy.xml"):
     if os.path.exists(p):
         os.system(f'sed -i "s/rights=\\"none\\"/rights=\\"read|write\\"/g" "{p}"')
 
 import gradio as gr
 import numpy as np
 import soundfile as sf
-import torch
 from moviepy.editor import VideoFileClip, CompositeVideoClip, TextClip
 from nemo.collections import asr as nemo_asr
 
+# --- Configuration ---
 SR = 16000
+MAX_VIDEO_BYTES = 200_000_000  # ≈200MB limite stable HF Space
 
 ASR_MODELS = {
     "Soloba CTC 0.6B V0": "RobotsMali/soloba-ctc-0.6b-v0",
@@ -26,50 +34,79 @@ ASR_MODELS = {
 
 _CACHE = {}
 
+# ---------------- LOAD MODEL ---------------- #
+
 def load_model(name):
     if name in _CACHE:
         return _CACHE[name]
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = nemo_asr.models.ASRModel.from_pretrained(model_name=ASR_MODELS[name]).to(device).eval()
+    model = nemo_asr.models.ASRModel.from_pretrained(
+        model_name=ASR_MODELS[name]
+    ).to(device).eval()
     _CACHE[name] = (model, device)
     return model, device
 
+# ---------------- AUDIO EXTRACTION ---------------- #
+
 def extract_audio(video_path, wav_path):
-    # Extraction audio stable (pas to_soundarray)
+    # Protéger contre fichiers trop lourds
+    try:
+        if os.path.getsize(video_path) > MAX_VIDEO_BYTES:
+            raise RuntimeError("⚠️ Vidéo trop lourde (>200MB). Compressez puis réessayez.")
+    except:
+        pass
+
+    # Extraction fiable (ffmpeg / MoviePy)
     try:
         clip = VideoFileClip(video_path)
+        if clip.audio is None:
+            clip.close()
+            raise RuntimeError("⚠️ Aucun audio détecté dans la vidéo.")
         clip.audio.write_audiofile(
             wav_path, fps=SR, codec="pcm_s16le", verbose=False, logger=None
         )
+        duration = clip.duration or 0
         clip.close()
+        return duration
     except:
-        os.system(f"ffmpeg -i '{video_path}' -ac 1 -ar {SR} -vn -y '{wav_path}' >/dev/null 2>&1")
-    audio, sr = sf.read(wav_path)
-    return len(audio)/sr
+        os.system(f"ffmpeg -y -i '{video_path}' -ac 1 -ar {SR} -vn '{wav_path}' >/dev/null 2>&1")
+        audio, sr = sf.read(wav_path)
+        return len(audio)/sr if sr else 0.0
+
+# ---------------- TRANSCRIPTION ---------------- #
 
 def transcribe(model, device, wav_path, model_key):
     audio, sr = sf.read(wav_path)
     if audio.ndim > 1: audio = audio.mean(axis=1)
-    total_s = len(audio)/sr
+    total_s = len(audio)/sr if sr else 0
 
     x = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
     ln = torch.tensor([x.shape[1]]).to(device)
 
-    # Timestamps Soloni (decode_and_align → hyp.words)
-    if "Soloni" in model_key:
+    # --- Soloni → timestamps réels via decode_and_align ---
+    if "Soloni" in model_key and hasattr(model, "decode_and_align"):
         with torch.no_grad():
-            proc, plen = model.preprocessor(x, ln)
-            hyps = model.decode_and_align(proc, plen)
+            proc, plen = model.preprocessor(
+                input_signal=x,
+                input_signal_length=ln
+            )
+            hyps = model.decode_and_align(
+                encoder_output=proc,
+                encoded_lengths=plen
+            )
+
         hyp = hyps[0][0] if isinstance(hyps[0], list) else hyps[0]
-        if hasattr(hyp, "words"):
+
+        if hasattr(hyp, "words") and hyp.words:
             return [(w.start_offset_ms/1000, w.end_offset_ms/1000, w.word) for w in hyp.words]
 
-    # CTC models → alignement dynamique
+    # --- Soloba + QuartzNet → alignement fluide stable ---
     text = model.transcribe([wav_path])[0]
     words = text.split()
-    if not words: return []
+    if not words or total_s <= 0:
+        return []
 
-    wps = max(2.2, len(words)/total_s)
+    wps = max(2.0, len(words)/total_s)
     subs, t = [], 0
     for w in words:
         d = 1/wps
@@ -78,41 +115,63 @@ def transcribe(model, device, wav_path, model_key):
         if t >= total_s: break
     return subs
 
+# ---------------- HARD SUBTITLE ---------------- #
+
 def burn(video_path, subs):
-    clip = VideoFileClip(video_path)
-    W, H = clip.size
-
-    layers = []
-    for s,e,w in subs:
-        txt = TextClip(w.upper(), fontsize=int(H/20), color="white",
-                       stroke_color="black", stroke_width=2,
-                       method="caption", size=(int(W*0.9), None)
-        ).set_start(s).set_duration(e-s).set_position(("center", int(H*0.88)))
-        layers.append(txt)
-
-    final = CompositeVideoClip([clip] + layers)
-    out = "RobotsMali_Subtitled.mp4"
-    final.write_videofile(out, codec="libx264", audio_codec="aac", verbose=False, logger=None)
-    return out
+    clip = None
+    final = None
+    try:
+        clip = VideoFileClip(video_path)
+        W, H = clip.size
+
+        layers = []
+        for s,e,w in subs:
+            if e <= s: continue
+            tc = TextClip(
+                w.upper(), fontsize=int(H/20), color="white",
+                stroke_color="black", stroke_width=2,
+                method="caption", size=(int(W*0.9), None)
+            ).set_start(s).set_duration(e - s).set_position(("center", int(H*0.88)))
+            layers.append(tc)
+
+        final = CompositeVideoClip([clip] + layers)
+        out = "RobotsMali_Subtitled.mp4"
+        final.write_videofile(out, codec="libx264", audio_codec="aac", verbose=False, logger=None)
+        return out
+
+    finally:
+        try: final.close()
+        except: pass
+        try: clip.close()
+        except: pass
+
+# ---------------- PIPELINE ---------------- #
 
 def pipeline(video, model_name, progress=gr.Progress()):
-    progress(0.2, "Chargement du modèle…")
+    progress(0.25, "Chargement du modèle…")
     model, device = load_model(model_name)
 
     with tempfile.TemporaryDirectory() as td:
         wav = f"{td}/audio.wav"
-        progress(0.45, "Extraction audio…")
-        extract_audio(video, wav)
 
-        progress(0.7, "Transcription en Bambara…")
+        progress(0.5, "Extraction audio…")
+        duration = extract_audio(video, wav)
+        if duration <= 0:
+            return "❌ Audio introuvable ou illisible.", None
+
+        progress(0.75, "Transcription en Bambara…")
         subs = transcribe(model, device, wav, model_name)
+        if not subs:
+            return "⚠️ Aucun mot détecté.", None
 
-    progress(0.9, "Incrustation des sous-titres…")
+    progress(0.95, "Incrustation des sous-titres…")
     out = burn(video, subs)
 
     progress(1.0, "✅ Terminé")
     return f"✅ Sous-titré avec **{model_name}**", out
 
+# ---------------- UI ---------------- #
+
 CSS = """
 body { background:#F6F9FF; font-family:Inter, sans-serif; }
 h1 { text-align:center; font-weight:800; color:#006CFF; }
@@ -121,11 +180,12 @@ h1 { text-align:center; font-weight:800; color:#006CFF; }
 
 with gr.Blocks(css=CSS, title="RobotsMali Caption Studio") as demo:
     gr.Markdown("<h1>RobotsMali Caption Studio</h1><p>Sous-titrage Automatique en Bambara</p>")
-    video = gr.File(label="🎥 Importer une vidéo", type="filepath")
-    model = gr.Dropdown(list(ASR_MODELS.keys()), value="Soloni 114M TDT CTC V1")
-    run = gr.Button("🚀 Générer les sous-titres")
+    video = gr.File(label="🎥 Importer une vidéo (max 200MB)", type="filepath")
+    model = gr.Dropdown(list(ASR_MODELS.keys()), value="Soloni 114M TDT CTC V1", label="🧠 Sélection du modèle ASR")
+    run = gr.Button("🚀 Générer")
     status = gr.Markdown()
     output = gr.Video()
+
     run.click(pipeline, inputs=[video, model], outputs=[status, output])
 
-demo.launch()
+demo.launch(server_name="0.0.0.0", server_port=7860, share=False, ssr_mode=False)