Zero-gpu

app.py

@@ -2,6 +2,8 @@ import os
 import time
 
 import gradio as gr
+import spaces
+import torch
 from pydub import AudioSegment
 from transformers import pipeline
 
@@ -28,44 +30,52 @@ model_dict = {
     ),
 }
 
-# Download models
-for model in model_dict.values():
-    pipeline("automatic-speech-recognition", model=model)
+print("Initializing pipelines...")
+pipe_dict = {
+    k: pipeline(
+        "automatic-speech-recognition",
+        model=v,
+        device="cuda" if torch.cuda.is_available() or is_hf else "cpu",
+    )
+    for k, v in model_dict.items()
+}
 
 
+@spaces.GPU
 def transcribe_common(audio: str, model: str) -> tuple[str, float]:
+    print(f"Transcribing {audio} with {model}")
     # Get duration of audio
     duration = AudioSegment.from_file(audio).duration_seconds
     if duration > 15:
         return "Audio too long, limit is 15 seconds", 0
     start_time = time.time()
-    pipe = pipeline("automatic-speech-recognition", model=model)
+    result = pipe_dict[model](audio, generate_kwargs=generate_kwargs)["text"]
     end_time = time.time()
-    return pipe(audio, generate_kwargs=generate_kwargs)["text"], end_time - start_time
+    return result, end_time - start_time
 
 
 def transcribe_large_v2(audio) -> tuple[str, float]:
-    return transcribe_common(audio, model_dict["whisper-large-v2"])
+    return transcribe_common(audio, "whisper-large-v2")
 
 
 def transcribe_large_v3(audio) -> tuple[str, float]:
-    return transcribe_common(audio, model_dict["whisper-large-v3"])
+    return transcribe_common(audio, "whisper-large-v3")
 
 
 def transcribe_large_v3_turbo(audio) -> tuple[str, float]:
-    return transcribe_common(audio, model_dict["whisper-large-v3-turbo"])
+    return transcribe_common(audio, "whisper-large-v3-turbo")
 
 
 def transcribe_kotoba_v1(audio) -> tuple[str, float]:
-    return transcribe_common(audio, model_dict["kotoba-whisper-v1.0"])
+    return transcribe_common(audio, "kotoba-whisper-v1.0")
 
 
 def transcribe_kotoba_v2(audio) -> tuple[str, float]:
-    return transcribe_common(audio, model_dict["kotoba-whisper-v2.0"])
+    return transcribe_common(audio, "kotoba-whisper-v2.0")
 
 
 def transcribe_galgame_whisper(audio) -> tuple[str, float]:
-    return transcribe_common(audio, model_dict["galgame-whisper-wip"])
+    return transcribe_common(audio, "galgame-whisper-wip")
 
 
 initial_md = """
@@ -91,39 +101,39 @@ generate_kwargs = {
 with gr.Blocks() as app:
     gr.Markdown(initial_md)
     audio = gr.Audio(type="filepath")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Galgame-Whisper (WIP)")
+            button_galgame = gr.Button("Transcribe with Galgame-Whisper (WIP)")
+            time_galgame = gr.Textbox(label="Time taken")
+            output_galgame = gr.Textbox(label="Result")
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Whisper-Large-V2")
             button_v2 = gr.Button("Transcribe with Whisper-Large-V2")
-            output_v2 = gr.Textbox()
-            time_v2 = gr.Textbox("Time taken")
+            time_v2 = gr.Textbox(label="Time taken")
+            output_v2 = gr.Textbox(label="Result")
         with gr.Column():
             gr.Markdown("### Whisper-Large-V3")
             button_v3 = gr.Button("Transcribe with Whisper-Large-V3")
-            output_v3 = gr.Textbox()
-            time_v3 = gr.Textbox("Time taken")
+            time_v3 = gr.Textbox(label="Time taken")
+            output_v3 = gr.Textbox(label="Result")
         with gr.Column():
             gr.Markdown("### Whisper-Large-V3-Turbo")
             button_v3_turbo = gr.Button("Transcribe with Whisper-Large-V3-Turbo")
-            output_v3_turbo = gr.Textbox()
-            time_v3_turbo = gr.Textbox()
+            time_v3_turbo = gr.Textbox(label="Time taken")
+            output_v3_turbo = gr.Textbox(label="Result")
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Kotoba-Whisper-V1.0")
             button_kotoba_v1 = gr.Button("Transcribe with Kotoba-Whisper-V1.0")
-            output_kotoba_v1 = gr.Textbox()
-            time_kotoba_v1 = gr.Textbox("Time taken")
+            time_kotoba_v1 = gr.Textbox(label="Time taken")
+            output_kotoba_v1 = gr.Textbox(label="Result")
         with gr.Column():
             gr.Markdown("### Kotoba-Whisper-V2.0")
             button_kotoba_v2 = gr.Button("Transcribe with Kotoba-Whisper-V2.0")
-            output_kotoba_v2 = gr.Textbox()
-            time_kotoba_v2 = gr.Textbox("Time taken")
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### Galgame-Whisper (WIP)")
-            button_galgame = gr.Button("Transcribe with Galgame-Whisper (WIP)")
-            output_galgame = gr.Textbox()
-            time_galgame = gr.Textbox("Time taken")
+            time_kotoba_v2 = gr.Textbox(label="Time taken")
+            output_kotoba_v2 = gr.Textbox(label="Result")
 
     button_v2.click(transcribe_large_v2, inputs=audio, outputs=[output_v2, time_v2])
     button_v3.click(transcribe_large_v3, inputs=audio, outputs=[output_v3, time_v3])

requirements.txt

@@ -1,4 +1,5 @@
 gradio
 numpy<2
+spaces
 torch
 transformers