Update app.py

app.py

@@ -1,27 +1,24 @@
-import os, warnings, tempfile, logging
+import os, warnings, logging, tempfile
 
-# --- 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 (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
 from moviepy.editor import VideoFileClip, CompositeVideoClip, TextClip
 from nemo.collections import asr as nemo_asr
 
-# --- Configuration ---
+# ---------------- CONFIG ---------------- #
+
 SR = 16000
-MAX_VIDEO_BYTES = 200_000_000  # ≈200MB limite stable HF Space
+MAX_VIDEO_BYTES = 200_000_000
 
 ASR_MODELS = {
     "Soloba CTC 0.6B V0": "RobotsMali/soloba-ctc-0.6b-v0",
@@ -46,44 +43,36 @@ def load_model(name):
     _CACHE[name] = (model, device)
     return model, device
 
-# ---------------- AUDIO EXTRACTION ---------------- #
+# ---------------- EXTRACT AUDIO (FORCE MONO) ---------------- #
 
 def extract_audio(video_path, wav_path):
-    # 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
+    if os.path.getsize(video_path) > MAX_VIDEO_BYTES:
+        raise RuntimeError("⚠️ Vidéo trop lourde (>200MB). Compressez puis réessayez.")
 
-    # 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 -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 ---------------- #
+    # Force audio mono + 16k (100% fiable)
+    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)
+    if sr == 0 or len(audio) == 0:
+        raise RuntimeError("⚠️ Audio introuvable ou illisible.")
+    return len(audio) / sr
+
+# ---------------- TRANSCRIBE ---------------- #
 
 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 if sr else 0
 
+    # Force mono propre + normalisation
+    if audio.ndim == 2:
+        audio = np.mean(audio, axis=1).astype(np.float32)
+    if np.max(np.abs(audio)) > 1:
+        audio = audio / np.max(np.abs(audio))
+
+    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)
 
-    # --- Soloni → timestamps réels via decode_and_align ---
+    # ---- Soloni : timestamps réels ---- #
     if "Soloni" in model_key and hasattr(model, "decode_and_align"):
         with torch.no_grad():
             proc, plen = model.preprocessor(
@@ -96,26 +85,26 @@ def transcribe(model, device, wav_path, model_key):
             )
 
         hyp = hyps[0][0] if isinstance(hyps[0], list) else hyps[0]
-
         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]
 
-    # --- Soloba + QuartzNet → alignement fluide stable ---
+    # ---- Soloba & QuartzNet fallback alignement fluide ---- #
     text = model.transcribe([wav_path])[0]
     words = text.split()
     if not words or total_s <= 0:
         return []
 
-    wps = max(2.0, len(words)/total_s)
+    wps = max(2.0, len(words) / total_s)
     subs, t = [], 0
     for w in words:
-        d = 1/wps
-        subs.append((t, min(total_s, t+d), w))
+        d = 1 / wps
+        subs.append((t, min(total_s, t + d), w))
         t += d
-        if t >= total_s: break
+        if t >= total_s:
+            break
     return subs
 
-# ---------------- HARD SUBTITLE ---------------- #
+# ---------------- BURN SUBTITLES ---------------- #
 
 def burn(video_path, subs):
     clip = None
@@ -123,16 +112,20 @@ def burn(video_path, subs):
     try:
         clip = VideoFileClip(video_path)
         W, H = clip.size
-
         layers = []
-        for s,e,w in subs:
+
+        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)
+            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(tc)
+            layers.append(txt)
 
         final = CompositeVideoClip([clip] + layers)
         out = "RobotsMali_Subtitled.mp4"
@@ -148,7 +141,7 @@ def burn(video_path, subs):
 # ---------------- PIPELINE ---------------- #
 
 def pipeline(video, model_name, progress=gr.Progress()):
-    progress(0.25, "Chargement du modèle…")
+    progress(0.3, "Chargement du modèle…")
     model, device = load_model(model_name)
 
     with tempfile.TemporaryDirectory() as td:
@@ -156,10 +149,8 @@ def pipeline(video, model_name, progress=gr.Progress()):
 
         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…")
+        progress(0.75, "Transcription…")
         subs = transcribe(model, device, wav, model_name)
         if not subs:
             return "⚠️ Aucun mot détecté.", None
@@ -168,20 +159,20 @@ def pipeline(video, model_name, progress=gr.Progress()):
     out = burn(video, subs)
 
     progress(1.0, "✅ Terminé")
-    return f"✅ Sous-titré avec **{model_name}**", out
+    return f"✅ Sous-titrage terminé avec **{model_name}**", out
 
 # ---------------- UI ---------------- #
 
 CSS = """
-body { background:#F6F9FF; font-family:Inter, sans-serif; }
-h1 { text-align:center; font-weight:800; color:#006CFF; }
-.gr-button { background:#007BFF !important; color:white !important; border-radius:8px; font-weight:700; }
+body { background:#F7FAFF; font-family:Inter, sans-serif; }
+h1 { text-align:center; font-weight:800; color:#005BFF; }
+.gr-button { background:#005BFF !important; color:white !important; border-radius:8px; font-weight:700; }
 """
 
 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>")
+    gr.Markdown("<h1>RobotsMali Caption Studio</h1><p>Transcription & Sous-titres Automatiques en Bambara</p>")
     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")
+    model = gr.Dropdown(list(ASR_MODELS.keys()), value="Soloni 114M TDT CTC V1", label="🧠 Modèle ASR")
     run = gr.Button("🚀 Générer")
     status = gr.Markdown()
     output = gr.Video()