Update app.py

app.py

@@ -1,11 +1,23 @@
 import gradio as gr
 from moviepy.editor import VideoFileClip
 import os
-from googletrans import Translator
-
 import whisper
-model = whisper.load_model("base")
+import torch
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+
+# Load Whisper model (use 'small' for faster transcription)
+model = whisper.load_model("small")
+
+# Load M2M-100 model & tokenizer
+m2m_model_name = "facebook/m2m100_418M"  # Use "facebook/m2m100_1.2B" for better accuracy
+tokenizer = M2M100Tokenizer.from_pretrained(m2m_model_name)
+translator_model = M2M100ForConditionalGeneration.from_pretrained(m2m_model_name)
 
+# Move model to GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+translator_model.to(device)
+
+# Supported languages for translation (must match M2M-100 language codes)
 TRANSLATION_LANGUAGES = {
     "Urdu": "ur",
     "French": "fr",
@@ -16,7 +28,17 @@ TRANSLATION_LANGUAGES = {
     "Hindi": "hi"
 }
 
-translator = Translator()
+def translate_text_m2m(text_list, target_lang):
+    """ Translates a list of English texts into the target language using M2M-100. """
+    tokenizer.src_lang = "en"
+    
+    # Tokenize and translate in batches
+    inputs = tokenizer(text_list, return_tensors="pt", padding=True, truncation=True).to(device)
+    outputs = translator_model.generate(**inputs, forced_bos_token_id=tokenizer.get_lang_id(target_lang))
+    
+    # Decode output
+    translated_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    return translated_texts
 
 def generate_translated_subtitles(video_path, target_language):
     # Extract audio from video
@@ -24,7 +46,7 @@ def generate_translated_subtitles(video_path, target_language):
     audio_path = "temp_audio.wav"
     video.audio.write_audiofile(audio_path)
     
-    # Transcribe (but do NOT translate) using Whisper
+    # Transcribe (without translation) using Whisper
     result = model.transcribe(audio_path, language="en")
     
     # Clean up temporary audio file
@@ -33,15 +55,14 @@ def generate_translated_subtitles(video_path, target_language):
     # Extract all subtitle texts for batch translation
     texts = [segment['text'] for segment in result['segments']]
     
-    # Perform batch translation (all at once)
-    translations = translator.translate(texts, dest=TRANSLATION_LANGUAGES[target_language])
+    # Translate using M2M-100
+    translated_texts = translate_text_m2m(texts, TRANSLATION_LANGUAGES[target_language])
     
     # Format subtitles in .srt format
     srt_content = ""
-    for index, (segment, translation) in enumerate(zip(result['segments'], translations)):
+    for index, (segment, translated_text) in enumerate(zip(result['segments'], translated_texts)):
         start_time = segment['start']
         end_time = segment['end']
-        translated_text = translation.text
 
         # Convert seconds to SRT time format (HH:MM:SS,mmm)
         def format_time(seconds):
@@ -77,8 +98,8 @@ iface = gr.Interface(
         gr.Textbox(label="Translated Subtitles", lines=10),
         gr.File(label="Download Translated .srt File")  # Correct file download
     ],
-    title="Video to Translated Subtitles",
-    description="Upload an English video, select a language, and get translated subtitles."
+    title="Video to Translated Subtitles (Offline M2M-100)",
+    description="Upload an English video, select a language, and get translated subtitles offline."
 )
 
-iface.launch(share=True, debug=True)
+iface.launch(share=True)