from pytube import YouTube
import ffmpeg
from gradio_client import Client
import subprocess
import gradio as gr

client = Client("abidlabs/music-separation")

placeholder_url = "https://www.youtube.com/watch?v=bfATETsvqt4&ab_channel=MiniMuslims"
placeholder_file = "example_acapella_video.mp4"

def acapellify(url):
    if url == placeholder_url:
        return placeholder_file
    
    yt = YouTube(url)
    try:
        video = yt.streams.get_highest_resolution()
        video.download(max_retries=5)
    except:
        raise gr.Error("YouTube video could not be downloaded. Try clicking Submit again and make sure that the URL is correct.")

    video_path = video.get_file_path()

    subprocess.run(['ffmpeg', '-i', video_path, '-vn', '-c:a', 'copy', 'audio.m4a', '-y'])

    result = client.predict(
        'audio.m4a',
        api_name="/predict"
    )

    subprocess.call(['ffmpeg', '-y', '-i', video_path, '-i', result[0], '-map', '0:v', '-map', '1:a', '-c:v', 'copy', '-c:a', 'aac', '-strict', 'experimental', 'acapella_video.mp4'])

    return "acapella_video.mp4"

io = gr.Interface(acapellify, gr.Textbox(label="URL to YouTube video", value=placeholder_url), gr.Video(label="Acapella output"))

with gr.Blocks() as demo:
    gr.Markdown("# Acapellify: Get a Vocals-only from any YouTube Video in 1 Click")
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Example Original Video**")
            gr.Video("old_video.mp4")
        with gr.Column():
            gr.Markdown("**Example Acapella Video**")
            gr.Video(placeholder_file)   

    gr.Markdown("**Try it yourself:** (for now, I recommend short videos <3 minutes)")
    io.render()
    
demo.launch()