import gradio as gr
import whisperx
import whisper
import torch
import spaces


@spaces.GPU
def transcribe(audio_file):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Transcribe with original Whisper
    model = whisper.load_model("medium", device)
    result = model.transcribe(audio_file)
    return result

inputs = gr.Audio(sources="upload", type="filepath")
outputs = gr.JSON()

gr.Interface(fn=transcribe, inputs=inputs, outputs=outputs).launch()