import gradio as gr
from PIL import Image
from aura_sr import AuraSR

# # Force CPU usage
# torch.set_default_type(torch.FloatTensor)
# torch.set_default_device('cpu')
#
# # Override torch.load to always use CPU
# original_load = torch.load
# torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, map_location=torch.device('cpu'))

# Initialize the AuraSR model
aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2/model.safetensors")


# # Restore original torch.load
# torch.load = original_load

def process_image(input_image):
    if input_image is None:
        raise gr.Error("Please provide an image to upscale.")

    # Convert to PIL Image for resizing
    pil_image = Image.fromarray(input_image)

    # Upscale the image using AuraSR
    upscaled_image = process_image_on_gpu(pil_image)

    return upscaled_image


def process_image_on_gpu(pil_image):
    return aura_sr.upscale_4x(pil_image)


title = """<h1 align="center">AuraSR-v2：一款基于GAN图像修复工具，可从低分辨率图片生成高分辨率图片</h1>"""

with gr.Blocks() as demo:
    gr.HTML(title)

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="输入图片", type="numpy")
            process_btn = gr.Button("生成")
        with gr.Column(scale=1):
            gallery = gr.Image(label="生成图片")

    process_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=gallery
    )

    # Add examples
    gr.Examples(
        examples=[
            "image1.png",
            "image3.png"
        ],
        inputs=input_image,
        outputs=gallery,
        fn=process_image,
        cache_examples=True
    )

demo.launch(server_name='0.0.0.0', share=True)