Update app.py

app.py

@@ -1,88 +1,174 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import srt
 import torch
 import os
+import math
+from datetime import timedelta
 
-# --- Configuration ---
-MODEL_CHECKPOINT = "facebook/nllb-200-distilled-1.3B"
-SRC_LANG = "eng_Latn"
-TGT_LANG = "arb_Arab"
+# --- إعدادات الموديلات ---
+# 1. موديل الترجمة (NLLB)
+TRANSLATION_MODEL = "facebook/nllb-200-distilled-1.3B" 
+# 2. موديل تفريغ الصوت (Whisper المطور)
+WHISPER_MODEL = "distil-whisper/distil-large-v3"
 
-# --- Load Model Directly ---
-print("Loading model...")
-# We use the tokenizer to convert text to numbers
-tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_CHECKPOINT)
-print("Model loaded!")
+print("Jari Tahmeel Al-Models... (Loading Models...)")
 
-def batch_translate(texts, batch_size=8):
+# --- تحميل موديل الترجمة (NLLB) ---
+tokenizer_nllb = AutoTokenizer.from_pretrained(TRANSLATION_MODEL)
+model_nllb = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATION_MODEL)
+
+# --- تحميل موديل الصوت (Whisper) ---
+# نستخدم chunk_length_s لتقسيم الصوت الطويل
+whisper_pipe = pipeline(
+    "automatic-speech-recognition",
+    model=WHISPER_MODEL,
+    torch_dtype=torch.float32,
+    device="cpu",
+    chunk_length_s=30,
+    stride_length_s=5,
+)
+
+print("Tam Tahmeel Al-Models Binajah! (Models Loaded!)")
+
+# ---------------------------------------------------------
+# الجزء الأول: دوال الترجمة (NLLB Logic)
+# ---------------------------------------------------------
+def batch_translate(texts, src_lang, tgt_lang, batch_size=8, progress=gr.Progress()):
     results = []
+    tokenizer_nllb.src_lang = src_lang
     
-    # Set the source language
-    tokenizer.src_lang = SRC_LANG
-
-    for i in range(0, len(texts), batch_size):
-        batch = texts[i : i + batch_size]
-        
-        # 1. Tokenize the batch
-        inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    total_batches = (len(texts) + batch_size - 1) // batch_size
+    
+    for i, start_idx in enumerate(range(0, len(texts), batch_size)):
+        # progress(i / total_batches, desc=f"Translating batch {i+1}/{total_batches}")
+        batch = texts[start_idx : start_idx + batch_size]
         
-        # 2. Get the Target Language ID properly (The Fix)
-        # We use convert_tokens_to_ids() instead of accessing the internal dictionary
-        forced_bos_token_id = tokenizer.convert_tokens_to_ids(TGT_LANG)
+        inputs = tokenizer_nllb(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        forced_bos_token_id = tokenizer_nllb.convert_tokens_to_ids(tgt_lang)
         
-        # 3. Generate translation
         with torch.no_grad():
-            generated_tokens = model.generate(
+            generated_tokens = model_nllb.generate(
                 **inputs,
                 forced_bos_token_id=forced_bos_token_id,
                 max_length=512
             )
             
-        # 4. Decode the results
-        batch_results = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+        batch_results = tokenizer_nllb.batch_decode(generated_tokens, skip_special_tokens=True)
         results.extend(batch_results)
-        
     return results
 
-def process_srt(filepath):
-    if filepath is None:
-        return None
-    
+def process_translation(filepath, src_lang_code, tgt_lang_code):
+    if filepath is None: return None
     try:
         with open(filepath, 'r', encoding='utf-8') as f:
             content = f.read()
-        subtitle_generator = srt.parse(content)
-        subtitles = list(subtitle_generator)
+        subtitles = list(srt.parse(content))
     except Exception as e:
-        return f"Error parsing SRT: {str(e)}"
+        return f"Error: {str(e)}"
 
-    # Translate content
-    texts_to_translate = [sub.content for sub in subtitles]
-    translated_texts = batch_translate(texts_to_translate)
+    texts = [sub.content for sub in subtitles]
+    translated = batch_translate(texts, src_lang_code, tgt_lang_code)
     
-    # Update subtitles
-    for sub, trans_text in zip(subtitles, translated_texts):
-        sub.content = trans_text
+    for sub, trans in zip(subtitles, translated):
+        sub.content = trans
 
-    # Save output
-    output_path = "translated_subtitles.srt"
-    with open(output_path, 'w', encoding='utf-8') as f:
+    out_path = "translated_subtitles.srt"
+    with open(out_path, 'w', encoding='utf-8') as f:
         f.write(srt.compose(subtitles))
+    return out_path
+
+# ---------------------------------------------------------
+# الجزء الثاني: دوال استخراج الصوت (Whisper Logic)
+# ---------------------------------------------------------
+def format_timestamp(seconds):
+    td = timedelta(seconds=seconds)
+    # تنسيق SRT يتطلب ساعات:دقائق:ثواني,مللي
+    total_seconds = int(td.total_seconds())
+    hours = total_seconds // 3600
+    minutes = (total_seconds % 3600) // 60
+    secs = total_seconds % 60
+    millis = int(td.microseconds / 1000)
+    return f"{hours:02}:{minutes:02}:{secs:02},{millis:03}"
+
+def video_to_srt(video_path, progress=gr.Progress()):
+    if video_path is None: return None
+    
+    progress(0.1, desc="Extracting Audio & Transcribing...")
     
-    return output_path
+    # تشغيل الـ Whisper Pipeline
+    # نطلب منه إرجاع الطوابع الزمنية (timestamps)
+    outputs = whisper_pipe(video_path, return_timestamps=True, generate_kwargs={"language": "english"})
+    
+    chunks = outputs.get("chunks", [])
+    if not chunks:
+        # أحيانًا يكون المخرج نصًا كاملاً إذا كان الفيديو قصيرًا جدًا
+        chunks = [{"text": outputs.get("text", ""), "timestamp": (0.0, 5.0)}]
 
-# --- Gradio Interface ---
-with gr.Blocks(title="NLLB SRT Translator") as demo:
-    gr.Markdown("# 🇬🇧 English to 🇸🇦 Arabic SRT Translator")
+    progress(0.8, desc="Formatting SRT...")
+    
+    # تحويل مخرجات ويسبر إلى صيغة SRT
+    srt_subtitles = []
+    for i, chunk in enumerate(chunks):
+        text = chunk['text'].strip()
+        start, end = chunk['timestamp']
+        
+        # حماية في حال كان الـ end غير موجود (None)
+        if end is None: end = start + 5.0 
+        
+        srt_subtitles.append(
+            srt.Subtitle(index=i+1, start=timedelta(seconds=start), end=timedelta(seconds=end), content=text)
+        )
     
-    with gr.Row():
-        input_file = gr.File(label="Upload English SRT", file_types=[".srt"])
-        output_file = gr.File(label="Download Arabic SRT")
+    out_path = "generated_captions.srt"
+    with open(out_path, 'w', encoding='utf-8') as f:
+        f.write(srt.compose(srt_subtitles))
+        
+    return out_path
+
+# ---------------------------------------------------------
+# واجهة المستخدم (Gradio Tabs)
+# ---------------------------------------------------------
+with gr.Blocks(title="The Ultimate Subtitler") as demo:
+    gr.Markdown("# 🎥 The Ultimate Subtitle Tool")
     
-    btn = gr.Button("Translate", variant="primary")
-    btn.click(fn=process_srt, inputs=input_file, outputs=output_file)
+    with gr.Tabs():
+        # --- التبويب الأول: من فيديو إلى SRT ---
+        with gr.TabItem("Step 1: Video to SRT (Whisper)"):
+            gr.Markdown("### استخرج ملف الترجمة الإنجليزية من أي فيديو")
+            with gr.Row():
+                video_input = gr.Video(label="Upload Video")
+                srt_output_gen = gr.File(label="Generated English SRT")
+            
+            gen_btn = gr.Button("Generate SRT from Video", variant="primary")
+            gen_btn.click(video_to_srt, inputs=video_input, outputs=srt_output_gen)
+
+        # --- التبويب الثاني: ترجمة الـ SRT ---
+        with gr.TabItem("Step 2: Translate SRT (NLLB)"):
+            gr.Markdown("### ترجم ملف الـ SRT إلى العربية (أو لغات أخرى)")
+            
+            with gr.Row():
+                srt_input = gr.File(label="Upload SRT File (English)")
+                
+                with gr.Column():
+                    # خيارات اللغات لتكون الأداة شاملة
+                    src_lang = gr.Dropdown(
+                        ["eng_Latn", "spa_Latn", "fra_Latn", "deu_Latn"], 
+                        label="Source Language", value="eng_Latn"
+                    )
+                    tgt_lang = gr.Dropdown(
+                        ["arb_Arab", "arz_Arab (Egyptian)", "eng_Latn", "fra_Latn"], 
+                        label="Target Language", value="arb_Arab"
+                    )
+                
+                srt_output_trans = gr.File(label="Translated SRT")
+            
+            trans_btn = gr.Button("Translate Subtitles", variant="primary")
+            trans_btn.click(
+                process_translation, 
+                inputs=[srt_input, src_lang, tgt_lang], 
+                outputs=srt_output_trans
+            )
 
 if __name__ == "__main__":
     demo.launch()