Update app.py

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 # 1. CONFIGURATION MATÉRIEL
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 2. DICTIONNAIRE DES MODÈLES (Mis à jour selon votre capture d'écran)
+# 2. DICTIONNAIRE DES MODÈLES
 MODELS = {
     "Soloba V3 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
     "Soloba V1.5 (TDT)":         ("RobotsMali/soloba-tdt-0.6b-v1.5", "rnnt"),
@@ -21,19 +21,18 @@ MODELS = {
 _cache = {}
 
 def clear_memory():
-    """Nettoie la VRAM et la RAM pour éviter les débordements."""
+    """Libère la VRAM et la RAM."""
     _cache.clear()
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-def load_model(name):
-    """Charge le modèle avec optimisation FP16 si possible."""
-    if name in _cache: return _cache[name]
+def get_model(name):
+    """Charge le modèle et le retourne directement."""
+    if name in _cache: 
+        return _cache[name]
     
-    yield f"⏳ Chargement du modèle {name}..."
     clear_memory()
-    
     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)
@@ -45,7 +44,7 @@ def load_model(name):
     
     model.to(DEVICE).eval()
     if DEVICE == "cuda":
-        model.half() # Utilisation de la demi-précision pour gagner 50% de VRAM
+        model.half()
     
     _cache[name] = model
     return model
@@ -58,37 +57,35 @@ def format_srt_time(sec):
 def pipeline(video_in, model_name):
     tmp_dir = tempfile.mkdtemp()
     try:
-        if not video_in: return "❌ Source vide", None, None
+        if not video_in: 
+            return "❌ Source vide", None
         
         # A. Extraction Audio
-        yield "⏳ Extraction de l'audio...", None, None
+        yield "⏳ Extraction de l'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 {full_wav}", shell=True, check=True)
         
-        # B. Segmentation Temporelle (10s)
+        # B. Segmentation (10s)
         segment_pattern = os.path.join(tmp_dir, "seg_%03d.wav")
         subprocess.run(f"ffmpeg -i {full_wav} -f segment -segment_time 10 -c copy {segment_pattern}", shell=True, check=True)
         audio_segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
         
-        # C. Chargement et Calcul de la Précision Temporelle
-        model_gen = load_model(model_name)
-        model = None
-        for update in model_gen:
-            if isinstance(update, str): yield update, None, None
-            else: model = update
+        # C. Chargement du modèle
+        yield f"⏳ Chargement du modèle {model_name}...", None
+        model = get_model(model_name)
             
-        # Calcul dynamique du stride (facteur de conversion frames -> secondes)
+        # Facteur de temps dynamique
         stride = 0.02 
         if hasattr(model, 'preprocessor') and hasattr(model.preprocessor, 'featurizer'):
             hop = model.preprocessor.featurizer.hop_length
             sr = model.preprocessor.featurizer.sample_rate
             stride = hop / sr
 
-        # D. Transcription par segments
+        # D. Transcription
         all_words_ts = []
         for idx, seg_path in enumerate(audio_segments):
             base_time = idx * 10.0
-            yield f"⏳ IA : Transcription segment {idx+1}/{len(audio_segments)}...", None, None
+            yield f"⏳ IA : Transcription {idx+1}/{len(audio_segments)}...", None
             
             hyp = model.transcribe([seg_path], return_hypotheses=True)[0]
             offsets = getattr(hyp, 'word_offsets', None)
@@ -99,12 +96,11 @@ def pipeline(video_in, model_name):
                     start_t = base_time + (offsets[i] * stride)
                     all_words_ts.append({"word": word, "start": start_t, "end": start_t + 0.45})
             else:
-                # Fallback : Répartition linéaire si les offsets manquent
                 gap = 10.0 / 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)})
 
-        # E. Génération du fichier SRT
+        # E. Génération SRT (Fichier temporaire interne)
         srt_path = os.path.join(tmp_dir, "final.srt")
         words_per_line = 6
         with open(srt_path, "w", encoding="utf-8") as f:
@@ -114,11 +110,10 @@ def pipeline(video_in, model_name):
                 f.write(f"{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 avec Incrustation (Burn-in)
-        yield "⏳ Rendu vidéo final ...", None, srt_path
+        # F. Encodage Vidéo
+        yield "⏳ Rendu vidéo final...", None
         out_path = os.path.abspath(f"robotsmali_final_{int(time.time())}.mp4")
         
-        # Protection des chemins pour FFmpeg (indispensable pour Windows/Linux)
         safe_srt = srt_path.replace("\\", "/").replace(":", "\\:")
         cmd_ffmpeg = (
             f"ffmpeg -y -i {shlex.quote(video_in)} "
@@ -127,28 +122,27 @@ def pipeline(video_in, model_name):
         )
         subprocess.run(cmd_ffmpeg, shell=True, check=True)
         
-        yield "✅ Transcription et Incrustation Terminées !", out_path, srt_path
+        yield "✅ Terminé !", out_path
 
     except Exception as e:
         traceback.print_exc()
-        yield f"❌ Erreur : {str(e)}", None, None
+        yield f"❌ Erreur : {str(e)}", None
 
-# 6. INTERFACE UTILISATEUR GRADIO
-with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
-    gr.HTML("<h1 style='text-align:center; color:#EAB308;'>🤖 ROBOTSMALI TRANSCRIPTION PRO</h1>")
+# 6. INTERFACE
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.HTML("<h1 style='text-align:center; color:#EAB308;'>🤖 ROBOTSMALI TRANSCRIPTION</h1>")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            v_in = gr.Video(label="Vidéo Source (Upload ou Webcam)", sources=["upload", "webcam"])
-            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Choisir le Modèle IA")
-            btn_run = gr.Button("🚀 GÉNÉRER LES SOUS-TITRES", variant="primary")
+        with gr.Column():
+            v_in = gr.Video(label="Source")
+            m_sel = gr.Dropdown(choices=list(MODELS.keys()), value="Soloba V3 (CTC)", label="Modèle IA")
+            btn_run = gr.Button("🚀 GÉNÉRER", variant="primary")
             
-        with gr.Column(scale=1):
-            status = gr.Markdown("### État\nPrêt à l'emploi.")
-            v_out = gr.Video(label="Vidéo Finale Incrustée")
-            f_srt = gr.File(label="Fichier Sous-titres (.SRT)")
+        with gr.Column():
+            status = gr.Markdown("### État\nPrêt.")
+            v_out = gr.Video(label="Vidéo Finale")
 
-    btn_run.click(pipeline, [v_in, m_sel], [status, v_out, f_srt])
+    btn_run.click(pipeline, [v_in, m_sel], [status, v_out])
 
 if __name__ == "__main__":
-    demo.launch(debug=True, show_error=True)
+    demo.launch(debug=True)