import gradio as gr
import os

import whisper
from whisper.utils import write_srt

model = whisper.load_model("base")

def transcribe(file):
    result = model.transcribe(file)

    #transcript
    with open(os.path.basename(file) + "-transcript.txt", 'w', encoding="utf-8") as f:
        f.write(result['text'])

    #subtitles
    with open(os.path.basename(file) + "-subs.srt", 'w', encoding="utf-8") as srt:
        write_srt(result["segments"], file=srt)

    download = []
    download.append(os.path.basename(file) + "-subs.srt");
    download.append(os.path.basename(file) + "-transcript.txt");

    return download

iface = gr.Interface(
    title = 'Whisper transcription and subtitles from file.', 
    fn=transcribe, 
    inputs=[
        gr.inputs.Audio(source="upload", type="filepath", label="Upload Audio")
    ],
    outputs=[
        gr.File(label="Download"),
    ],
    live=True)

iface.launch()