import gradio as gr
from animesr.inference_animesr_video import main
from utils import url_download
       
MODEL_LIST = [
    "kadirnar/AnimeSR_Paper_Model",
    "kadirnar/AnimeSR_v2"
]

examples = [
    ["https://www.youtube.com/watch?v=KJeD4ErwO04", "kadirnar/AnimeSR_v2"]
    #["https://www.youtube.com/watch?v=icPHcK_cCF4", "kadirnar/AnimeSR_v2"]
    
]

def AnimeSr_Video(
    source: str,
    model_id: str,
):
    if source.startswith('http'):
        source = url_download(source)
    else:
        source = source
    main(source=source, model_id=model_id)    
    save_path = 'output.mp4'
    return source, save_path


app = gr.Blocks()
with app:
    gr.Markdown("# **<p align='center'>AnimeSR: Learning Real-World Super-Resolution Models for Animation Videos</p>**")
    gr.Markdown(
        """
        <p style='text-align: center'>
        Follow me for more! 
        <br> <a href='https://twitter.com/kadirnar_ai' target='_blank'>twitter</a> | <a href='https://github.com/kadirnar' target='_blank'>github</a> | <a href='https://www.linkedin.com/in/kadir-nar/' target='_blank'>linkedin</a> | 
        </p>
        """
    )    
    with gr.Row():
        with gr.Column():
            with gr.Tab("URL"):
                input_text = gr.inputs.Textbox(lines=1, label="Video URL")
                input_text_button = gr.Button(value="Predict")
            
            with gr.Tab("Local"):
                gr.Markdown("### **Video**")
                input_video = gr.inputs.Video(label="Video")  
                input_video_button = gr.Button(value="Predict")    
            
            input_dropdown =gr.inputs.Dropdown(choices=MODEL_LIST)
        with gr.Column():
            output = gr.outputs.Video(label="Input Video")
        
        with gr.Column():
            video_gif = gr.outputs.Video(label="Output Video")
    
    gr.Examples(examples, inputs=[input_text, input_dropdown], outputs=[output, video_gif], fn=AnimeSr_Video, cache_examples=False)
    input_text_button.click(AnimeSr_Video, inputs=[input_text, input_dropdown], outputs=[output, video_gif])
    input_video_button.click(AnimeSr_Video, inputs=[input_video, input_dropdown], outputs=[output, video_gif])
    
app.launch()