import gradio as gr
from PIL import Image
import numpy as np
from aura_sr import AuraSR
import torch
import os


USE_TORCH_COMPILE = False
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

def process_image(input_image):
    if input_image is None:
        raise gr.Error("Please provide an image to upscale.")
    print("get input image: ", input_image)
    # Upscale the image using AuraSR
    upscaled_image = process_image_on_gpu(input_image)
    print("upscaled_image: ", upscaled_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(type="pil", label="输入图片")
            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)