Update app.py

app.py

@@ -6,74 +6,72 @@ import numpy as np
 from aura_sr import AuraSR
 import torch
 
-# Force CPU usage
 torch.set_default_tensor_type(torch.FloatTensor)
-
-# 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")
 
-# Restore original torch.load
 torch.load = original_load
 
 @spaces.GPU
-def process_image(input_image):
+def process_image(input_image, scale_factor):
     if input_image is None:
-        raise gr.Error("Please provide an image to upscale.")
+        raise gr.Error("请提供一张图片。")
 
-    # Convert to PIL Image for resizing
     pil_image = Image.fromarray(input_image)
+    original_width, original_height = pil_image.size
 
     upscaled_image = aura_sr.upscale_4x(pil_image)
 
-    # Convert result to numpy array if it's not already
+    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)
 
-    print(input_image, result_array)
+    return (input_image, result_array), upscaled_image
 
-    return (input_image, result_array)
-    
-title = """<h1 align="center">AuraSR</h1>
-<p><center>Upscales your images to x4</center></p>
-<p><center>
-<a href="https://huggingface.co/fal/AuraSR-v2" target="_blank">[AuraSR-v2]</a>
-<a href="https://blog.fal.ai/introducing-aurasr-an-open-reproduction-of-the-gigagan-upscaler-2/" target="_blank">[Blog Post]</a>
-<a href="https://huggingface.co/fal-ai/AuraSR" target="_blank">[v1 Model Page]</a>
-</center></p>
-<br/>
-<p>This is an open reproduction of the GigaGAN Upscaler from fal.ai</p>
-"""
+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="Input Image", type="numpy")
-            process_btn = gr.Button(value="Upscale Image", variant = "primary")
+            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="Before / After", type="numpy")
+            # 图片对比滑块
+            output_slider = ImageSlider(label="对比效果", type="numpy")
 
-    process_btn.click(
-        fn=process_image,
-        inputs=[input_image],
-        outputs=output_slider
-    )
+            download_output = gr.Image(
+                label="下载结果 (PNG)", 
+                type="pil", 
+                format="png", 
+                interactive=False,
+                visible=True
+            )
 
-    # Add examples
-    gr.Examples(
-        examples=[
-            "image1.png",
-            "image3.png"
-        ],
-        inputs=input_image,
-        outputs=output_slider,
+    process_btn.click(
         fn=process_image,
-        cache_examples=True,
+        inputs=[input_image, scale_slider],
+        outputs=[output_slider, download_output]
     )
 
 demo.launch(debug=True)