Update app.py

app.py

@@ -4,6 +4,7 @@ import random
 import numpy as np
 import json
 from datetime import timedelta
+import tempfile
 import gradio as gr
 from groq import Groq
 
@@ -280,7 +281,7 @@ def generate_subtitles(audio_file_path, prompt, language, auto_detect_language):
 
     # If there's an error during file check
     if error_message:
-        return error_message
+        return error_message, None, None
 
     with open(processed_path, "rb") as file:
         transcription_json = client.audio.transcriptions.create(
@@ -297,45 +298,46 @@ def generate_subtitles(audio_file_path, prompt, language, auto_detect_language):
     transcription_text = transcription_json['text']
 
     srt_content = create_srt_from_text(transcription_text)
-    return srt_content
-
-
-def add_subtitles_to_video(video_file_path, srt_content):
-    """Adds subtitles to a video using ffmpeg."""
-    output_file_path = os.path.splitext(video_file_path)[0] + "_with_subs.mp4"
-    
-    try:
-        with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_srt_file:
-            temp_srt_path = temp_srt_file.name
-            temp_srt_file.write(srt_content)
-
-        # Use temporary file with -i for subtitles 
-        subprocess.run(
-            [
-                "ffmpeg",
-                "-i",
-                video_file_path,
-                "-i",
-                temp_srt_path,
-                "-map",
-                "0:v",
-                "-map",
-                "0:a",
-                "-map",
-                "1:s?",
-                "-c:v",
-                "copy",
-                "-c:a",
-                "copy",
-                "-c:s",
-                "mov_text",
-                output_file_path,
-            ],
-            check=True,
-        )
-        return output_file_path
-    except subprocess.CalledProcessError as e:
-        return f"Error during subtitle addition: {e}"
+
+    # Generate subtitles and add to video if MP4
+    if audio_file_path.lower().endswith(".mp4"):
+        try:
+            with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_srt_file:
+                temp_srt_path = temp_srt_file.name
+                temp_srt_file.write(srt_content)
+
+                # Use temporary file with -i for subtitles 
+                subprocess.run(
+                    [
+                        "ffmpeg",
+                        "-i",
+                        audio_file_path,
+                        "-i",
+                        temp_srt_path,
+                        "-map",
+                        "0:v",
+                        "-map",
+                        "0:a",
+                        "-map",
+                        "1:s?",
+                        "-c:v",
+                        "copy",
+                        "-c:a",
+                        "copy",
+                        "-c:s",
+                        "mov_text",
+                        audio_file_path.replace(".mp4", "_with_subs.mp4"), # Create new file name
+                    ],
+                    check=True,
+                )
+                return srt_content, audio_file_path.replace(".mp4", "_with_subs.mp4"), None
+        except subprocess.CalledProcessError as e:
+            return srt_content, None, f"Error during subtitle addition: {e}"
+        finally:
+            if os.path.exists(temp_srt_path):
+                os.remove(temp_srt_path)
+
+    return srt_content, None, None 
         
 
 with gr.Blocks() as demo:
@@ -467,6 +469,7 @@ with gr.Blocks() as demo:
                             )
                     transcribe_button_subtitles = gr.Button("Generate Subtitles")
                     srt_output = gr.Textbox(label="SRT Output")
+                    srt_download_button = gr.Button("Download SRT")
                     video_output = gr.File(label="Output Video (with Subtitles)")
                     transcribe_button_subtitles.click(
                         generate_subtitles,
@@ -476,12 +479,13 @@ with gr.Blocks() as demo:
                             language_subtitles,
                             auto_detect_language_subtitles,
                         ],
-                        outputs=srt_output,
+                        outputs=[srt_output, video_output, gr.Error()],
                     )
-                    srt_output.change(
-                        add_subtitles_to_video,
-                        inputs=[audio_input_subtitles, srt_output],
-                        outputs=video_output,
+                    srt_download_button.click(
+                        lambda x: x,
+                        inputs=srt_output,
+                        outputs=srt_output,
+                        download=True,
                     )
 
 demo.launch()