import gradio as gr
from infer import infer_image, infer_video

theme = gr.themes.Base(
    primary_hue=gr.themes.colors.purple,
    neutral_hue=gr.themes.colors.gray,
    font=gr.themes.GoogleFont("DM Sans"),
).set(
    body_background_fill="#0e0e0e",
    body_background_fill_dark="#0e0e0e",
    block_background_fill="#181818",
    block_background_fill_dark="#181818",
    block_border_color="#2a2a2a",
    block_border_color_dark="#2a2a2a",
    block_label_text_color="#888888",
    block_label_text_color_dark="#888888",
    block_title_text_color="#f0f0f0",
    block_title_text_color_dark="#f0f0f0",
    border_color_accent="#a855f7",
    border_color_accent_dark="#a855f7",
    button_primary_background_fill="#a855f7",
    button_primary_background_fill_dark="#a855f7",
    button_primary_background_fill_hover="#b87ef7",
    button_primary_background_fill_hover_dark="#b87ef7",
    button_primary_text_color="#ffffff",
    button_primary_text_color_dark="#ffffff",
    button_secondary_background_fill="#222222",
    button_secondary_background_fill_dark="#222222",
    button_secondary_border_color="#333333",
    button_secondary_border_color_dark="#333333",
    button_secondary_text_color="#f0f0f0",
    button_secondary_text_color_dark="#f0f0f0",
    input_background_fill="#181818",
    input_background_fill_dark="#181818",
    input_border_color="#333333",
    input_border_color_dark="#333333",
    input_placeholder_color="#555555",
    input_placeholder_color_dark="#555555",
    slider_color="#a855f7",
    slider_color_dark="#a855f7",
    checkbox_background_color="#181818",
    checkbox_background_color_dark="#181818",
    checkbox_background_color_selected="#a855f7",
    checkbox_background_color_selected_dark="#a855f7",
    body_text_color="#f0f0f0",
    body_text_color_dark="#f0f0f0",
    body_text_color_subdued="#888888",
    body_text_color_subdued_dark="#888888",
)

with gr.Blocks(theme=theme, title="Image Upscaler") as demo:
    gr.Markdown("""
    # ✦ Image Upscaler
    Upscale images and videos using Real-ESRGAN AI
    """)

    with gr.Tabs():
        with gr.TabItem("🖼 Image"):
            with gr.Row():
                with gr.Column():
                    img_input = gr.Image(type='pil', label='Select image')
                    img_scale = gr.Radio([2, 4, 8], value=4, label="Scale")
                    img_btn = gr.Button("✦ Upscale", variant="primary")
                with gr.Column():
                    img_output = gr.Image(type="pil", label="Output Image")
            img_btn.click(fn=infer_image, inputs=[img_input, img_scale], outputs=img_output)

        with gr.TabItem("🎬 Video"):
            with gr.Row():
                with gr.Column():
                    vid_input = gr.Video(label='Select video')
                    vid_scale = gr.Radio([2, 4, 8], value=2, label="Scale")
                    vid_btn = gr.Button("✦ Upscale", variant="primary")
                with gr.Column():
                    vid_output = gr.Video(label="Output Video")
            vid_btn.click(fn=infer_video, inputs=[vid_input, vid_scale], outputs=vid_output)

demo.launch()