Update app.py

app.py

@@ -1,91 +1,65 @@
-import gradio as gr
-import whisper
-from whisper.utils import write_vtt
 
-from pytube import YouTube
+import gradio as gr 
 import os
 import sys
 import subprocess
+#from moviepy.editor import VideoFileClip
 
-loaded_model = whisper.load_model("base")
-current_size = 'base'
-
-def inference(link):
-
-  yt = YouTube(link)
-  
-  audio_path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp4")
-  print(f'audio path :  {audio_path}')
-  video_path = yt.streams.filter(file_extension='mp4')
-  
-  #options = whisper.DecodingOptions(without_timestamps=True)
-  options = dict(beam_size=5, best_of=5, fp16 = False)
-  translate_options = dict(task="inference", **options)
-  results = loaded_model.transcribe(audio_path,**translate_options)
-  
-  output_dir = ''
-  path = audio_path.split(".")[0]
-  
-  with open(os.path.join(output_dir, path + ".vtt"), "w") as vtt:
-    write_vtt(results["segments"], file=vtt)  
-
-  subtitle = path + ".vtt"
-  output_video = path + "_subtitled.mp4"
-
-  os.system(f"ffmpeg -i {video_path} -vf subtitles={subtitle} {output_video}")
-  
-  return output_video
-  
-def change_model(size):
-  if size == current_size:
-    return
-  loaded_model = whisper.load_model(size)
-  current_size = size
-
-def populate_metadata(link):
-  yt = YouTube(link)
-  return yt.thumbnail_url, yt.title
-
-title="Youtube Caption Generator"
-description="Generate captions of Youtube videos using OpenAI's Whisper"
-block = gr.Blocks()
+import whisper
+from whisper.utils import write_vtt
+
+model = whisper.load_model("medium")
+
+title = "Add Text/Caption to your YouTube Shorts - MultiLingual"
+
+def video2mp3(video_file, output_ext="mp3"):
+    filename, ext = os.path.splitext(video_file)
+    subprocess.call(["ffmpeg", "-y", "-i", video_file, f"{filename}.{output_ext}"], 
+                    stdout=subprocess.DEVNULL,
+                    stderr=subprocess.STDOUT)
+    return f"{filename}.{output_ext}"
+
+
+def translate(input_video):
+
+    audio_file = video2mp3(input_video)
+    
+    options = dict(beam_size=5, best_of=5, fp16 = False)
+    translate_options = dict(task="translate", **options)
+    result = model.transcribe(audio_file,**translate_options)
+
+    output_dir = ''
+    audio_path = audio_file.split(".")[0]
 
+    with open(os.path.join(output_dir, audio_path + ".vtt"), "w") as vtt:
+      write_vtt(result["segments"], file=vtt)
+
+    subtitle = audio_path + ".vtt"
+    output_video = audio_path + "_subtitled.mp4"
+
+    os.system(f"ffmpeg -i {input_video} -vf subtitles={subtitle} {output_video}")
+
+    return output_video
+
+block = gr.Blocks()
 with block:
-    gr.HTML(
-        """
-            <div style="text-align: center; max-width: 500px; margin: 0 auto;">
-              <div>
-                <h1>Youtube Caption Generator</h1>
-              </div>
-              <p style="margin-bottom: 10px; font-size: 94%">
-                Generate captions of Youtube videos using OpenAI's Whisper
-              </p>
-            </div>
-        """
-    )
+
     with gr.Group():
         with gr.Box():
-          sz = gr.Dropdown(label="Model Size", choices=['base','small', 'medium', 'large'], value='base')
-          
-          link = gr.Textbox(label="YouTube Link")
-          
-          with gr.Row().style(mobile_collapse=False, equal_height=True):
-            title = gr.Label(label="Video Title", placeholder="Title")
-            img = gr.Image(label="Thumbnail")
-            
-        #   text = gr.Textbox(
-        #       label="Transcription", 
-        #       placeholder="Transcription Output",
-        #       lines=5)
-          
-          op_video = gr.Video()
+            with gr.Row().style():
+                inp_video = gr.Video(
+                    label="Input Video",
+                    type="filepath",
+                    mirror_webcam = False
+                )
+                op_video = gr.Video()
+            btn = gr.Button("Generate Subtitle Video")
         
-          with gr.Row().style(mobile_collapse=False, equal_height=True): 
-              btn = gr.Button("Generate Captions")       
-          
-          # Events
-          btn.click(inference, inputs=[link], outputs=[op_video])
-          link.change(populate_metadata, inputs=[link], outputs=[img, title])
-          sz.change(change_model, inputs=[sz], outputs=[])
-
-block.launch(debug=True,enable_queue=True)
+        
+        
+
+
+        
+        btn.click(translate, inputs=[inp_video], outputs=[op_video])
+
+block.launch(enable_queue = True)