Update app.py

app.py

@@ -4,17 +4,35 @@ import os
 from faster_whisper import WhisperModel
 from moviepy.video.io.VideoFileClip import VideoFileClip
 import logging
+import google.generativeai as genai
 
 # Suppress moviepy logs
 logging.getLogger("moviepy").setLevel(logging.ERROR)
 
-# Define the model and device
+# Configure Gemini API
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+
+# Create the Gemini model
+generation_config = {
+    "temperature": 1,
+    "top_p": 0.95,
+    "top_k": 40,
+    "max_output_tokens": 8192,
+    "response_mime_type": "text/plain",
+}
+
+model = genai.GenerativeModel(
+    model_name="gemini-2.0-flash-exp",
+    generation_config=generation_config,
+)
+
+# Define the Whisper model and device
 MODEL_NAME = "Systran/faster-whisper-large-v3"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 compute_type = "float32" if device == "cuda" else "int8"
 
 # Load the Whisper model
-model = WhisperModel(MODEL_NAME, device=device, compute_type=compute_type)
+whisper_model = WhisperModel(MODEL_NAME, device=device, compute_type=compute_type)
 
 # List of all supported languages in Whisper
 SUPPORTED_LANGUAGES = [
@@ -46,7 +64,7 @@ def extract_audio_from_video(video_file):
 def generate_subtitles(audio_file, language="Auto Detect"):
     """Generate subtitles from an audio file using Whisper."""
     # Transcribe the audio
-    segments, info = model.transcribe(
+    segments, info = whisper_model.transcribe(
         audio_file,
         task="transcribe",
         language=None if language == "Auto Detect" else language.lower(),
@@ -67,7 +85,7 @@ def generate_subtitles(audio_file, language="Auto Detect"):
         # Add to SRT
         srt_subtitles += f"{i}\n{start_time_srt} --> {end_time_srt}\n{text}\n\n"
     
-    return srt_subtitles
+    return srt_subtitles, info.language
 
 def format_timestamp(seconds):
     """Convert seconds to SRT timestamp format (HH:MM:SS,mmm)."""
@@ -77,23 +95,51 @@ def format_timestamp(seconds):
     milliseconds = int((seconds - int(seconds)) * 1000)
     return f"{hours:02}:{minutes:02}:{int(seconds):02},{milliseconds:03}"
 
-def process_video(video_file, language="Auto Detect"):
-    """Process a video file to generate subtitles."""
+def translate_text(text, target_language):
+    """Translate text using Gemini."""
+    prompt = f"Translate the following text into {target_language}:\n\n{text}"
+    response = model.generate_content(prompt)
+    return response.text
+
+def translate_srt(srt_text, target_language):
+    """Translate an SRT file while preserving timestamps."""
+    translated_srt = ""
+    for line in srt_text.split("\n"):
+        if " --> " in line:  # Timestamp line
+            translated_srt += line + "\n"
+        elif line.strip().isdigit():  # Subtitle index line
+            translated_srt += line + "\n"
+        elif line.strip():  # Text line
+            translated_srt += translate_text(line, target_language) + "\n"
+        else:  # Empty line
+            translated_srt += "\n"
+    return translated_srt
+
+def process_video(video_file, language="Auto Detect", translate_to=None):
+    """Process a video file to generate and translate subtitles."""
     # Extract audio from the video
     audio_file = extract_audio_from_video(video_file)
     
     # Generate subtitles
-    subtitles = generate_subtitles(audio_file, language)
+    subtitles, detected_language = generate_subtitles(audio_file, language)
     
-    # Save subtitles to an SRT file
-    srt_file = "subtitles.srt"
-    with open(srt_file, "w", encoding="utf-8") as f:
+    # Save original subtitles to an SRT file
+    original_srt_file = "original_subtitles.srt"
+    with open(original_srt_file, "w", encoding="utf-8") as f:
         f.write(subtitles)
     
+    # Translate subtitles if a target language is provided
+    translated_srt_file = None
+    if translate_to and translate_to != "None":
+        translated_subtitles = translate_srt(subtitles, translate_to)
+        translated_srt_file = "translated_subtitles.srt"
+        with open(translated_srt_file, "w", encoding="utf-8") as f:
+            f.write(translated_subtitles)
+    
     # Clean up extracted audio file
     os.remove(audio_file)
     
-    return srt_file
+    return original_srt_file, translated_srt_file, detected_language
 
 # Define the Gradio interface
 with gr.Blocks(title="AutoSubGen - AI Video Subtitle Generator") as demo:
@@ -101,7 +147,7 @@ with gr.Blocks(title="AutoSubGen - AI Video Subtitle Generator") as demo:
     with gr.Column():
         gr.Markdown("# 🎥 AutoSubGen")
         gr.Markdown("### AI-Powered Video Subtitle Generator")
-        gr.Markdown("Automatically generate subtitles for your videos in **SRT format**. Supports **100+ languages** and **auto-detection**.")
+        gr.Markdown("Automatically generate and translate subtitles for your videos in **SRT format**. Supports **100+ languages** and **auto-detection**.")
     
     # Main content
     with gr.Tab("Generate Subtitles"):
@@ -114,14 +160,23 @@ with gr.Blocks(title="AutoSubGen - AI Video Subtitle Generator") as demo:
                 value="Auto Detect",
                 scale=1
             )
+            translate_to_dropdown = gr.Dropdown(
+                choices=["None"] + SUPPORTED_LANGUAGES[1:],  # Exclude "Auto Detect"
+                label="Translate To",
+                value="None",
+                scale=1
+            )
         generate_button = gr.Button("Generate Subtitles", variant="primary")
-        subtitle_output = gr.File(label="Download Subtitles (SRT)")
+        with gr.Row():
+            original_subtitle_output = gr.File(label="Download Original Subtitles (SRT)")
+            translated_subtitle_output = gr.File(label="Download Translated Subtitles (SRT)")
+        detected_language_output = gr.Textbox(label="Detected Language")
     
     # Link button to function
     generate_button.click(
         process_video,
-        inputs=[video_input, language_dropdown],
-        outputs=subtitle_output
+        inputs=[video_input, language_dropdown, translate_to_dropdown],
+        outputs=[original_subtitle_output, translated_subtitle_output, detected_language_output]
     )
 
 # Launch the Gradio interface with a public link