import whisper
from pytube import YouTube
from transformers import pipeline
import gradio as gr
import os

model = whisper.load_model("base")
summarizer = pipeline("summarization")

def get_audio(url):
    yt = YouTube(url)
    video = yt.streams.filter(only_audio=True).first()
    out_file = video.download(output_path=".")
    base, ext = os.path.splitext(out_file)
    new_file = base + '.mp3'
    os.rename(out_file, new_file)
    return new_file

def get_text_from_url(url):
    result = model.transcribe(get_audio(url))
    return result['text']

def get_text_from_file(file):
    # Assuming the uploaded file is already in MP3 format
    result = model.transcribe(file.name)
    return result['text']

def get_summary_from_url(url):
    article = get_text_from_url(url)
    b = summarizer(article)
    return b[0]['summary_text']

def get_summary_from_file(file):
    article = get_text_from_file(file)
    b = summarizer(article)
    return b[0]['summary_text']

def process_url(url):
    transcription = get_text_from_url(url)
    summary = get_summary_from_url(url)
    return summary, transcription

def process_file(file):
    transcription = get_text_from_file(file)
    summary = get_summary_from_file(file)
    return summary, transcription

with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Youtube video transcription with OpenAI's Whisper</center></h1>")
    gr.Markdown("<center>Enter the link of any youtube video or upload an MP4 file to get the transcription and a summary in text form.</center>")
    
    with gr.Tab('Youtube Video'):
        with gr.Row():
            input_text = gr.Textbox(placeholder='Enter the Youtube video URL', label='URL')
            output_summary = gr.Textbox(placeholder='Summary text of the Youtube Video', label='Summary')
            output_transcription = gr.Textbox(placeholder='Transcription of the video', label='Transcription')
        result_button = gr.Button('Process Youtube Video')
        
    with gr.Tab('Uploaded MP4'):
        with gr.Row():
            input_file = gr.File(label='Upload MP4')
            output_file_summary = gr.Textbox(placeholder='Summary text of the video', label='Summary')
            output_file_transcription = gr.Textbox(placeholder='Transcription of the video', label='Transcription')
        result_button_file = gr.Button('Process Uploaded MP4')

    result_button.click(process_url, inputs=input_text, outputs=[output_summary, output_transcription])
    result_button_file.click(process_file, inputs=input_file, outputs=[output_file_summary, output_file_transcription])
    
demo.launch(debug=True)