Update app.py

app.py

@@ -5,66 +5,43 @@ import soundfile as sf
 from moviepy.editor import VideoFileClip, CompositeVideoClip, ImageClip
 from PIL import Image, ImageDraw, ImageFont
 from nemo.collections import asr as nemo_asr
+from huggingface_hub import hf_hub_download
 from ctc_segmentation import ctc_segmentation, CtcSegmentationParameters, prepare_text
 
 
-# =============================
-# LISTE OFFICIELLE DES MODELES ROBOTSMALI
-# =============================
-
 MODELS = {
-    "Soloni V0": "RobotsMali/soloni-114m-tdt-ctc-V0",
-    "Soloni V1": "RobotsMali/soloni-114m-tdt-ctc-V1",
+    "Soloni V0": ("RobotsMali/soloni-114m-tdt-ctc-V0", "soloni-114m-tdt-ctc-V0.nemo", "rnnt"),
+    "Soloni V1": ("RobotsMali/soloni-114m-tdt-ctc-V1", "soloni-114m-tdt-ctc-V1.nemo", "rnnt"),
 
-    "Soloba V0": "RobotsMali/soloba-ctc-0.6b-V0",
-    "Soloba V1": "RobotsMali/soloba-ctc-0.6b-V1",
+    "Soloba V0": ("RobotsMali/soloba-ctc-0.6b-V0", None, "ctc"),
+    "Soloba V1": ("RobotsMali/soloba-ctc-0.6b-V1", None, "ctc"),
 
-    "QuartzNet V0": "RobotsMali/stt-bm-quartznet15x5-V0",
-    "QuartzNet V1": "RobotsMali/stt-bm-quartznet15x5-V1"
+    "QuartzNet V0": ("RobotsMali/stt-bm-quartznet15x5-V0", None, "ctc"),
+    "QuartzNet V1": ("RobotsMali/stt-bm-quartznet15x5-V1", None, "ctc"),
 }
 
 
-# =============================
-# EXTRACTION AUDIO (SOLIDE & COMPATIBLE HF)
-# =============================
-
 def extract_audio(video_path, wav_path):
-    (
-        VideoFileClip(video_path)
-        .audio
-        .write_audiofile(
-            wav_path,
-            fps=16000,
-            codec="pcm_s16le",
-            verbose=False,
-            logger=None
-        )
-    )
-
-
-# =============================
-# TRANSCRIPTION + ALIGNEMENT
-# =============================
+    (VideoFileClip(video_path).audio.write_audiofile(
+        wav_path, fps=16000, codec="pcm_s16le", verbose=False, logger=None
+    ))
+
 
 def transcribe(model, device, wav, model_name):
     audio, sr = sf.read(wav)
     if audio.ndim == 2:
         audio = np.mean(audio, axis=1)
-
     x = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
     ln = torch.tensor([x.shape[1]]).to(device)
     total_s = len(audio) / sr
 
-    # === Soloni → timestamps natifs ===
     if "Soloni" in model_name:
         with torch.no_grad():
             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]
         return [(w.start_offset_ms/1000, w.end_offset_ms/1000, w.word) for w in hyp.words]
 
-    # === Soloba & QuartzNet → CTC Forced Alignment ===
     text = model.transcribe([wav])[0].strip()
     if not text:
         return []
@@ -84,90 +61,69 @@ def transcribe(model, device, wav, model_name):
                 timings[i+1] * tps if i+1 < len(timings) else total_s,
                 words[i]) for i in range(len(words))]
 
-    # Groupage lisible (max 4 mots par sous-titre)
     grouped, temp = [], []
     for w in aligned:
         temp.append(w)
         if len(temp) >= 4:
-            grouped.append(temp)
-            temp = []
-    if temp:
-        grouped.append(temp)
+            grouped.append(temp); temp = []
+    if temp: grouped.append(temp)
 
     return [(g[0][0], g[-1][1], " ".join([w[2] for w in g])) for g in grouped]
 
 
-# =============================
-# INCRUSTATION SOUS-TITRES (SANS IMAGEMAGICK)
-# =============================
-
 def burn(video, subs):
     clip = VideoFileClip(video)
     W, H = clip.size
-
     try:
         font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", int(H/20))
     except:
         font = ImageFont.load_default()
 
     layers = []
-    for s, e, text in subs:
-        img = Image.new("RGBA", (W, int(H*0.12)), (0,0,0,140))
+    for s,e,text in subs:
+        img = Image.new("RGBA",(W,int(H*0.12)),(0,0,0,140))
         draw = ImageDraw.Draw(img)
         bbox = draw.textbbox((0,0), text, font=font)
         tw, th = bbox[2]-bbox[0], bbox[3]-bbox[1]
-        draw.text(((W-tw)//2, (int(H*0.12)-th)//2), text, font=font, fill="white")
-        layers.append(ImageClip(np.array(img)).set_start(s).set_duration(e-s).set_position(("center", int(H*0.85))))
+        draw.text(((W-tw)//2,(int(H*0.12)-th)//2), text, font=font, fill="white")
+        layers.append(ImageClip(np.array(img)).set_start(s).set_duration(e-s).set_position(("center",int(H*0.85))))
 
     final = CompositeVideoClip([clip] + layers)
     out = "RobotsMali_Subtitled.mp4"
     final.write_videofile(out, codec="libx264", audio_codec="aac", fps=clip.fps, verbose=False, logger=None)
-
-    clip.close()
-    final.close()
+    clip.close(); final.close()
     return out
 
 
-# =============================
-# PIPELINE PRINCIPAL
-# =============================
-
 def pipeline(video_file, model_name):
     if video_file is None:
         return "Veuillez importer une vidéo.", None
 
+    repo, nemo_file, mode = MODELS[model_name]
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Chargement correct selon le modèle
-    if "Soloni" in model_name:
-        model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.from_pretrained(model_name=MODELS[model_name])
+    if mode == "rnnt":
+        nemo_path = hf_hub_download(repo, filename=nemo_file)
+        model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(nemo_path)
     else:
-        model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=MODELS[model_name])
+        model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=repo)
 
-    model = model.to(device)
-    model.eval()
+    model = model.to(device); model.eval()
 
     wav = "audio.wav"
     extract_audio(video_file, wav)
     subs = transcribe(model, device, wav, model_name)
     final = burn(video_file, subs)
-
     return "✅ Sous-titres générés.", final
 
 
-# =============================
-# INTERFACE (DESIGN CONSERVÉ)
-# =============================
-
 with gr.Blocks() as demo:
     gr.Markdown("# 🎙️ **RobotsMali — Sous-titrage automatique Bambara**")
-
     video = gr.Video(label="Vidéo")
     model = gr.Dropdown(list(MODELS.keys()), value="Soloni V1", label="Modèle")
     btn = gr.Button("⚡ Générer les sous-titres")
     status = gr.Markdown()
-    out = gr.Video(label="Résultat (avec sous-titres)")
-
+    out = gr.Video(label="Résultat")
     btn.click(pipeline, inputs=[video, model], outputs=[status, out])
 
 demo.launch()