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


# 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 star_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> STAR </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="STAR"):
            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 = star_demo(result_dir)
    venhancer_iface.queue(max_size=12)
    venhancer_iface.launch(max_threads=1)