Update app.py

app.py

@@ -3,32 +3,63 @@ import whisper
 from pytube import YouTube
 import os
 
-loaded_model = whisper.load_model("base")
-current_size = 'base'
+class GradioInference():
+  def __init__(self):
+    self.sizes = list(whisper._MODELS.keys())
+    self.file = "Upload audio/video"
+    self.current_size = "base"
+    self.loaded_model = whisper.load_model(self.current_size)
+    self.yt = None
+  
+  def __call__(self, link, file, size, subs):
+    if self.yt is None:
+      self.yt = YouTube(link)
+      path = self.yt.streams.filter(only_audio=True)[0].download(filename="tmp.mp4")
 
-def inference(link):
-  if link.startswith('https://'):
-    yt = YouTube(link)
-    path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp4")
-  else:
-    path = link
-  options = whisper.DecodingOptions(without_timestamps=True)
-  results = loaded_model.transcribe(path)
-  return results['text']
+    if self.file == None:
+      self.file = path
 
-def change_model(size):
-  if size == current_size:
-    return
-  loaded_model = whisper.load_model(size)
-  current_size = size
+    if size != self.current_size:
+      self.loaded_model = whisper.load_model(size)
+      self.current_size = size
+    results = self.loaded_model.transcribe(path)
 
-def populate_metadata(link):
-  if link.startswith('https://'):
-    yt = YouTube(link)
-    return yt.thumbnail_url, yt.title
+    if subs == "None":
+      return results["text"]
+    elif subs == ".srt":
+      return self.srt(results["segments"])
+    elif ".csv" == ".csv":
+      return self.csv(results["segments"])
+   
+  def srt(self, segments):
+    output = ""
+    for i, segment in enumerate(segments):
+      output += f"{i+1}\n"
+      output += f"{self.format_time(segment['start'])} --> {self.format_time(segment['end'])}\n"
+      output += f"{segment['text']}\n\n"
+    return output
+  
+  def csv(self, segments):
+    output = ""
+    for segment in segments:
+      output += f"{segment['start']},{segment['end']},{segment['text']}\n"
+    return output
 
+  def format_time(self, time):
+    hours = time//3600
+    minutes = (time - hours*3600)//60
+    seconds = time - hours*3600 - minutes*60
+    milliseconds = (time - int(time))*1000
+    return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d},{int(milliseconds):03d}"
+    
+  def populate_metadata(self, link):
+    self.yt = YouTube(link)
+    return self.yt.thumbnail_url, self.yt.title
+
+gio = GradioInference()
 title="Youtube Whisperer"
 description="Speech to text transcription of Youtube videos using OpenAI's Whisper"
+
 block = gr.Blocks()
 
 with block:
@@ -46,23 +77,18 @@ with block:
     )
     with gr.Group():
         with gr.Box():
-          sz = gr.Dropdown(label="Model Size", choices=['base','small', 'medium', 'large'], value='base')
-          
-          link = gr.File(label="YouTube Link or Upload a Video")
-          
-          with gr.Row().style(mobile_collapse=False, equal_height=True):
-            title = gr.Label(label="Video Title", placeholder="Title")
+          with gr.Row().style(equal_height=True):
+            sz = gr.Dropdown(label="Model Size", choices=gio.sizes, value='base')
+            file = gr.File(label="Upload Video/Audio")
+          with gr.Row().style(equal_height=True):
+            wt = gr.Radio(["None", ".srt", ".csv"], label="With Timestamps?")
+          link = gr.Textbox(label="YouTube Link")
+          title = gr.Label(label="Video Title")
+          with gr.Row().style(equal_height=True):
             img = gr.Image(label="Thumbnail")
-          text = gr.Textbox(
-              label="Transcription", 
-              placeholder="Transcription Output",
-              lines=5)
-          with gr.Row().style(mobile_collapse=False, equal_height=True): 
+            text = gr.Textbox(label="Transcription", placeholder="Transcription Output", lines=10)
+          with gr.Row().style(equal_height=True): 
               btn = gr.Button("Transcribe")       
-          
-         # Events
-          btn.click(inference, inputs=[link], outputs=[text])
-          link.change(populate_metadata, inputs=[link], outputs=[img, title])
-          sz.change(change_model, inputs=[sz], outputs=[])
-
-block.launch(debug=True)
+          btn.click(gio, inputs=[link, lang, sz, wt], outputs=[text])
+          link.change(gio.populate_metadata, inputs=[link], outputs=[img, title])
+block.launch()