download option

New UI:
- download element
- output folder
- max image element height (no more scrolling needed)

app.py

@@ -1,3 +1,5 @@
+import os
+import time
 import spaces
 import gradio as gr
 from gradio_imageslider import ImageSlider
@@ -21,6 +23,14 @@ aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2", device=DEVICE)
 # Restore original torch.load
 torch.load = original_load
 
+# Create output folder if not exists
+output_folder = '../outputs'
+os.makedirs(output_folder, exist_ok=True)
+
+def generate_output_filename():
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    return f"upscaled_{timestamp}.png"
+
 def process_image(input_image):
     if input_image is None:
         raise gr.Error("Please provide an image to upscale.")
@@ -34,7 +44,12 @@ def process_image(input_image):
     # Convert result to numpy array if it's not already
     result_array = np.array(upscaled_image)
 
-    return [input_image, result_array]
+    # Save result as PNG
+    output_filename = generate_output_filename()
+    output_path = os.path.join(output_folder, output_filename)
+    upscaled_image.save(output_path, format="PNG")
+
+    return [input_image, result_array], output_path
 
 @spaces.GPU
 def process_image_on_gpu(pil_image):
@@ -48,21 +63,28 @@ title = """<h1 align="center">AuraSR-v2 - An open reproduction of the GigaGAN Up
 </center></p>
 """
 
-with gr.Blocks() as demo:
+css = """
+    .img { max-height: 80vh !Important }
+    #slider-container { overflow: hidden; display: flex; justify-content: center }
+    #row-height { height: 65px !important }
+"""
+
+with gr.Blocks(css=css) as demo:
     
     gr.HTML(title)
     
     with gr.Row():
         with gr.Column(scale=1):
-            input_image = gr.Image(label="Input Image", type="numpy")
+            input_image = gr.Image(label="Input Image", type="numpy", elem_classes="img")
             process_btn = gr.Button("Upscale Image")
         with gr.Column(scale=1):
-            output_slider = ImageSlider(label="Before / After", type="numpy")
+            output_slider = ImageSlider(label="Before / After", type="numpy", elem_classes="img", elem_id="slider-container")
+            download_btn = gr.File(label="Download Image", elem_id="row-height")
 
     process_btn.click(
         fn=process_image,
         inputs=[input_image],
-        outputs=output_slider
+        outputs=[output_slider, download_btn]
     )
 
     # Add examples