import spaces
import gradio as gr
from gradio_imageslider import ImageSlider
from PIL import Image
import numpy as np
from aura_sr import AuraSR
import torch

torch.set_default_tensor_type(torch.FloatTensor)
original_load = torch.load
torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, map_location=torch.device('cpu'))

aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")

torch.load = original_load

@spaces.GPU
def process_image(input_image, scale_factor):
    if input_image is None:
        raise gr.Error("请提供一张图片。")

    pil_image = Image.fromarray(input_image)
    original_width, original_height = pil_image.size

    upscaled_image = aura_sr.upscale_4x(pil_image)

    target_scale = float(scale_factor)
    
    if target_scale != 4.0:
        # 计算目标尺寸
        new_width = int(original_width * target_scale)
        new_height = int(original_height * target_scale)
        # 使用高质量重采样算法调整尺寸
        upscaled_image = upscaled_image.resize((new_width, new_height), Image.LANCZOS)

    result_array = np.array(upscaled_image)

    return (input_image, result_array), upscaled_image

title = """<h1 align="center">AuraSR-v2 动态放大版</h1>"""

with gr.Blocks() as demo:
    gr.HTML(title)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="输入图片", type="numpy")

            scale_slider = gr.Slider(
                minimum=1.0, 
                maximum=4.0, 
                value=4.0, 
                step=0.5, 
                label="放大倍数 (Scale Factor)"
            )
            
            process_btn = gr.Button(value="开始放大", variant="primary")
            
        with gr.Column(scale=1):
            # 图片对比滑块
            output_slider = ImageSlider(label="对比效果", type="numpy")

            download_output = gr.Image(
                label="下载结果 (PNG)", 
                type="pil", 
                format="png", 
                interactive=False,
                visible=True
            )

    process_btn.click(
        fn=process_image,
        inputs=[input_image, scale_slider],
        outputs=[output_slider, download_output]
    )

demo.launch(debug=True)