import gradio as gr
import spaces
from gradio_imageslider import ImageSlider
from image_gen_aux import UpscaleWithModel
from image_gen_aux.utils import load_image


@spaces.GPU
def upscale_image(image):
    original = load_image(image)

    upscaler = UpscaleWithModel.from_pretrained("OzzyGT/UltraSharp").to("cuda")
    image = upscaler(original)

    return original, image


title = """<h1 align="center">Image Upscaler</h1>"""

with gr.Blocks() as demo:
    gr.HTML(title)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            run_button = gr.Button("Upscale")
        with gr.Column():
            result = ImageSlider(
                interactive=False,
                label="Generated Image",
                elem_id="result-image",
                position=0.1,
            )

    run_button.click(
        fn=upscale_image,
        inputs=[input_image],
        outputs=result,
    )


demo.launch(share=False)