import whisper
import gradio as gr

# Preload default model
model = whisper.load_model("small")

# Cache models to avoid re-downloading
models = {
    "tiny": None,
    "base": None,
    "small": model,
    "medium": None,
    "large": None
}

def load_model(selected_model):
    if models[selected_model] is None:
        models[selected_model] = whisper.load_model(selected_model)
    return models[selected_model]

def transcribe(audio, model_name):
    model = load_model(model_name)
    result = model.transcribe(audio)
    return result["text"]

gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath", label="Upload Audio"),
        gr.Dropdown(["tiny", "base", "small", "medium", "large"], value="small", label="Choose Whisper Model")
    ],
    outputs="text",
    title="Whisper Multilingual STT (All Models)",
    description="Upload an audio file and choose which Whisper model to use for transcription. Supports multiple languages including Korean, Urdu, Arabic, etc."
).launch()