import spaces
import os
import gradio as gr
from video_super_resolution.scripts.inference_sr import VEnhancer_sr

examples = [
    ["prompts/astronaut.mp4", "An astronaut flying in space, featuring a steady and smooth perspective", 4, 24, 250],
    ["prompts/cat.mp4", "A cat wearing sunglasses at a pool", 4, 24, 250],
    ["prompts/fish.mp4", "Clown fish swimming through the coral reef", 4, 24, 300],
    ["prompts/iron_man.mp4", "Iron Man flying in the sky", 4, 24, 250],
    ["prompts/raccoon.mp4", "A cute raccoon playing guitar in a boat on the ocean", 4, 24, 250],
    ["prompts/shanghai.mp4", "The bund Shanghai by Hokusai, in the style of Ukiyo", 4, 24, 250],
    ["prompts/gwen.mp4", "Gwen Stacy reading a book, black and white", 4, 24, 250],
]

# Define a GPU-decorated function for enhancement
@spaces.GPU(duration=1200)
def enhance_with_gpu(input_video, input_text):
    return venhancer.enhance_a_video(input_video, input_text)

def venhancer_demo(result_dir="./tmp/"):
    css = """#input_video {max-width: 1024px !important} #output_vid {max-width: 2048px; max-height:1280px}"""
    global venhancer
    venhancer = VEnhancer_sr(result_dir)
    
    with gr.Blocks(analytics_enabled=False, css=css) as venhancer_iface:
        gr.Markdown(
            "<div align='center'> <h1> VEnhancer </span> </h1> \
                     <a style='font-size:18px;color: #000000' href='https://arxiv.org/abs/2501.02976'> [ArXiv] </a>\
                     <a style='font-size:18px;color: #000000' href='https://nju-pcalab.github.io/projects/STAR'> [Project Page] </a> \
                     <a style='font-size:18px;color: #000000' href='https://github.com/NJU-PCALab/STAR'> [Github] </a> </div>"
        )
        with gr.Tab(label="VEnhancer"):
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            input_video = gr.Video(label="Input Video", elem_id="input_video")
                        with gr.Row():
                            input_text = gr.Text(label="Prompts")
                        end_btn = gr.Button("Generate")
                    with gr.Row():
                        output_video = gr.Video(
                            label="Generated Video", elem_id="output_vid", autoplay=True, show_share_button=True
                        )

                gr.Examples(
                    examples=examples,
                    inputs=[input_video, input_text],
                    outputs=[output_video],
                    fn=enhance_with_gpu,  # Use the GPU-decorated function
                    cache_examples=False,
                )
            end_btn.click(
                inputs=[input_video, input_text],
                outputs=[output_video],
                fn=enhance_with_gpu,  # Use the GPU-decorated function
            )

    return venhancer_iface


if __name__ == "__main__":
    result_dir = os.path.join("./", "results")
    venhancer_iface = venhancer_demo(result_dir)
    venhancer_iface.queue(max_size=12)
    venhancer_iface.launch(max_threads=1)