Update app.py

app.py

@@ -8,15 +8,15 @@ from transformers.utils import is_flash_attn_2_available
 from languages import get_language_names
 from subtitle_manager import Subtitle
 
-
 logging.basicConfig(level=logging.INFO)
 last_model = None
 pipe = None
 
-def write_file(output_file,subtitle):
+def write_file(output_file, subtitle):
     with open(output_file, 'w', encoding='utf-8') as f:
         f.write(subtitle)
 
+@spaces.GPU
 def create_pipe(model, flash):
     if torch.cuda.is_available():
         device = "cuda:0"
@@ -55,8 +55,9 @@ def create_pipe(model, flash):
     )
     return pipe
 
+@spaces.GPU
 def transcribe_webui_simple_progress(modelName, languageName, urlData, multipleFiles, microphoneData, task, flash,
-                                    chunk_length_s, batch_size, progress=gr.Progress()):
+                                     chunk_length_s, batch_size, progress=gr.Progress()):
     global last_model
     global pipe
 
@@ -69,7 +70,7 @@ def transcribe_webui_simple_progress(modelName, languageName, urlData, multipleF
     logging.info(f"chunk_length_s: {chunk_length_s}")
     logging.info(f"batch_size: {batch_size}")
 
-    if last_model == None:
+    if last_model is None:
         logging.info("first model")
         progress(0.1, desc="Loading Model..")
         pipe = create_pipe(modelName, flash)
@@ -88,7 +89,7 @@ def transcribe_webui_simple_progress(modelName, languageName, urlData, multipleF
 
     files = []
     if multipleFiles:
-        files+=multipleFiles
+        files += multipleFiles
     if urlData:
         files.append(urlData)
     if microphoneData:
@@ -107,8 +108,8 @@ def transcribe_webui_simple_progress(modelName, languageName, urlData, multipleF
         logging.info(file)
         outputs = pipe(
             file,
-            chunk_length_s=chunk_length_s,#30
-            batch_size=batch_size,#24
+            chunk_length_s=chunk_length_s,  # 30
+            batch_size=batch_size,  # 24
             generate_kwargs=generate_kwargs,
             return_timestamps=True,
         )
@@ -119,13 +120,13 @@ def transcribe_webui_simple_progress(modelName, languageName, urlData, multipleF
         srt = srt_sub.get_subtitle(outputs["chunks"])
         vtt = vtt_sub.get_subtitle(outputs["chunks"])
         txt = txt_sub.get_subtitle(outputs["chunks"])
-        write_file(file_out+".srt",srt)
-        write_file(file_out+".vtt",vtt)
-        write_file(file_out+".txt",txt)
-        files_out += [file_out+".srt", file_out+".vtt", file_out+".txt"]
+        write_file(file_out + ".srt", srt)
+        write_file(file_out + ".vtt", vtt)
+        write_file(file_out + ".txt", txt)
+        files_out += [file_out + ".srt", file_out + ".vtt", file_out + ".txt"]
 
     progress(1, desc="Completed!")
-    
+
     return files_out, vtt, txt
 
 
@@ -142,7 +143,7 @@ with gr.Blocks(title="Insanely Fast Whisper") as demo:
         "openai/whisper-large-v2", "distil-whisper/distil-large-v2",
         "openai/whisper-large-v3", "distil-whisper/distil-large-v3", "xaviviro/whisper-large-v3-catalan-finetuned-v2",
     ]
-    waveform_options=gr.WaveformOptions(
+    waveform_options = gr.WaveformOptions(
         waveform_color="#01C6FF",
         waveform_progress_color="#0066B4",
         skip_length=2,
@@ -150,25 +151,29 @@ with gr.Blocks(title="Insanely Fast Whisper") as demo:
     )
 
     simple_transcribe = gr.Interface(fn=transcribe_webui_simple_progress,
-        description=description,
-        article=article,
-        inputs=[
-            gr.Dropdown(choices=whisper_models, value="distil-whisper/distil-large-v2", label="Model", info="Select whisper model", interactive = True,),
-            gr.Dropdown(choices=["Automatic Detection"] + sorted(get_language_names()), value="Automatic Detection", label="Language", info="Select audio voice language", interactive = True,),
-            gr.Text(label="URL", info="(YouTube, etc.)", interactive = True),
-            gr.File(label="Upload Files", file_count="multiple"),
-            gr.Audio(sources=["upload", "microphone",], type="filepath", label="Input", waveform_options = waveform_options),
-            gr.Dropdown(choices=["transcribe", "translate"], label="Task", value="transcribe", interactive = True),
-            gr.Checkbox(label='Flash',info='Use Flash Attention 2'),
-            gr.Number(label='chunk_length_s',value=30, interactive = True),
-            gr.Number(label='batch_size',value=24, interactive = True)
-        ], outputs=[
-            gr.File(label="Download"),
-            gr.Text(label="Transcription"), 
-            gr.Text(label="Segments")
-        ]
-    )
+                                     description=description,
+                                     article=article,
+                                     inputs=[
+                                         gr.Dropdown(choices=whisper_models, value="distil-whisper/distil-large-v2",
+                                                     label="Model", info="Select whisper model", interactive=True, ),
+                                         gr.Dropdown(choices=["Automatic Detection"] + sorted(get_language_names()),
+                                                     value="Automatic Detection", label="Language",
+                                                     info="Select audio voice language", interactive=True, ),
+                                         gr.Text(label="URL", info="(YouTube, etc.)", interactive=True),
+                                         gr.File(label="Upload Files", file_count="multiple"),
+                                         gr.Audio(sources=["upload", "microphone", ], type="filepath", label="Input",
+                                                  waveform_options=waveform_options),
+                                         gr.Dropdown(choices=["transcribe", "translate"], label="Task",
+                                                     value="transcribe", interactive=True),
+                                         gr.Checkbox(label='Flash', info='Use Flash Attention 2'),
+                                         gr.Number(label='chunk_length_s', value=30, interactive=True),
+                                         gr.Number(label='batch_size', value=24, interactive=True)
+                                     ], outputs=[
+                                         gr.File(label="Download"),
+                                         gr.Text(label="Transcription"),
+                                         gr.Text(label="Segments")
+                                     ]
+                                     )
 
 if __name__ == "__main__":
-    demo.launch()
-
+    demo.launch()