import gradio as gr
import whisper
import os

model_size = list(whisper._MODELS.keys())

# main processing
def main(URL, model_size):
    video_path = youtube_dl(URL)
    text = transcript(video_path, model_size)
    return text

# Transcribe text using Whisper
def transcript(video_path, model_size):
    model = whisper.load_model(model_size)
    result = model.transcribe(video_path)
    return result["text"]

# Download the video
def youtube_dl(URL):
    # Download the video
    os.system(f"yt-dlp -f best -v {URL} -o target.mp4")
    # Path of downloaded video
    video_path = os.path.join(os.path.dirname(__file__), "target.mp4")
    return video_path

with gr.Blocks() as demo:

    with gr.Row():

        with gr.Column():

            with gr.Row():
                url = gr.Textbox(placeholder = 'Youtube video URL', label = 'URL')

            with gr.Row():
                model_size = gr.Dropdown(choices = model_size, value = 'tiny', label = "Model")

            with gr.Row():
                transcribe_btn = gr.Button('Transcribe')

        with gr.Column():
            outputs = gr.Textbox(placeholder = 'Transcription of the video', label = 'Transcription')

    transcribe_btn.click(fn = main,
                        inputs = [url, model_size],
                        outputs = outputs
                        )

if __name__ == "__main__":
    demo.launch()