import gradio as gr
from animesr.inference_animesr_video import main


def AnimeSr_Video(
    source: str,
    model_id: str,
):

    main(source=source, model_id=model_id)    
    return 'output.mp4'
       

inputs = [
    gr.Video(),
    gr.inputs.Dropdown(
        label="Model",
        choices=[
            "kadirnar/AnimeSR_Paper_Model",
            "kadirnar/AnimeSR_v2",
        ],
        default="kadirnar/AnimeSR_Paper_Model",
    ),
]

outputs = gr.Video()
title = "AnimeSR: Learning Real-World Super-Resolution Models for Animation Videos"

examples = [
    [
        "/home/kadir/AI/AnimeSR/videos/panda.mp4",
        "kadirnar/AnimeSR_v2",
    ],
    [
        "/home/kadir/AI/AnimeSR/videos/tom.mp4",
        "kadirnar/AnimeSR_Paper_Model",
    ],
    [
        "/home/kadir/AI/AnimeSR/videos/ironman.mp4",
        "kadirnar/AnimeSR_v2",
    ]
]

demo_app = gr.Interface(
    fn=AnimeSr_Video,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    cache_examples=True,
    theme='huggingface',
)
demo_app.launch(debug=True, enable_queue=True)