Update app.py

app.py

@@ -1,11 +1,32 @@
+import os
+import uuid
 import gradio as gr
 import spaces
 import numpy as np
 from PIL import Image
 from image_gen_aux import UpscaleWithModel
 
+# -----------------------------
+# Model configuration
+# -----------------------------
+
+MODEL_ID = "2xEvangelion_dat2"
 UPSCALE_MODEL = None
-MODEL_ID = "Phips/2xEvangelion_dat2"  # anime-oriented example from the working Space list
+
+# Aspect ratios
+RATIO_MAP = {
+    "16:9": (16, 9),
+    "9:16": (9, 16),
+    "4:5": (4, 5),
+    "1:1": (1, 1),
+    "5:4": (5, 4),
+    "2:3": (2, 3),
+    "3:2": (3, 2),
+}
+
+# -----------------------------
+# Load model once (lazy load)
+# -----------------------------
 
 def get_upscaler():
     global UPSCALE_MODEL
@@ -13,29 +34,147 @@ def get_upscaler():
         UPSCALE_MODEL = UpscaleWithModel.from_pretrained(MODEL_ID).to("cuda")
     return UPSCALE_MODEL
 
+# -----------------------------
+# Resize image to selected ratio
+# -----------------------------
+
+def fit_to_ratio(image: Image.Image, ratio_name: str, long_side: int):
+
+    rw, rh = RATIO_MAP[ratio_name]
+
+    if rw >= rh:
+        width = long_side
+        height = int(long_side * rh / rw)
+    else:
+        height = long_side
+        width = int(long_side * rw / rh)
+
+    width -= width % 2
+    height -= height % 2
+
+    return image.resize((width, height), Image.LANCZOS)
+
+# -----------------------------
+# GPU upscale function
+# -----------------------------
+
 @spaces.GPU
-def upscale(image):
+def upscale(
+    image,
+    tile_width,
+    tile_height,
+    use_ratio,
+    ratio_name,
+    target_long_side
+):
+
     if image is None:
-        return None
+        return None, None
 
     img = Image.fromarray(image).convert("RGB")
+
+    # apply ratio preset if enabled
+    if use_ratio:
+        img = fit_to_ratio(img, ratio_name, int(target_long_side))
+
     upscaler = get_upscaler()
-    out = upscaler(img, tiling=True, tile_width=1024, tile_height=1024)
 
-    if isinstance(out, Image.Image):
-        return np.array(out)
-    return out
+    out = upscaler(
+        img,
+        tiling=True,
+        tile_width=int(tile_width),
+        tile_height=int(tile_height)
+    )
+
+    if not isinstance(out, Image.Image):
+        out = Image.fromarray(out)
+
+    os.makedirs("/tmp/upscaled", exist_ok=True)
+
+    file_path = f"/tmp/upscaled/{uuid.uuid4().hex}.png"
+
+    out.save(
+        file_path,
+        format="PNG",
+        compress_level=0
+    )
+
+    return np.array(out), file_path
+
+
+# -----------------------------
+# UI
+# -----------------------------
 
 with gr.Blocks() as demo:
+
     gr.Markdown("# Anime Upscale Alpha")
-    input_image = gr.Image(type="numpy", label="Input Image", height=400)
+
+    input_image = gr.Image(
+        type="numpy",
+        label="Input Image",
+        height=400
+    )
+
+    with gr.Row():
+
+        tile_width = gr.Number(
+            value=1024,
+            label="Tile Width",
+            precision=0
+        )
+
+        tile_height = gr.Number(
+            value=1024,
+            label="Tile Height",
+            precision=0
+        )
+
+    with gr.Group():
+
+        use_ratio = gr.Checkbox(
+            label="Use Aspect Ratio Preset",
+            value=False
+        )
+
+        ratio_name = gr.Radio(
+            choices=["16:9", "9:16", "4:5", "1:1", "5:4", "2:3", "3:2"],
+            value="1:1",
+            label="Aspect Ratio"
+        )
+
+        target_long_side = gr.Number(
+            value=2048,
+            label="Target Long Side (px)",
+            precision=0
+        )
+
     run_button = gr.Button("Upscale")
-    output_image = gr.Image(type="numpy", label="Upscaled Image", height=400)
+
+    output_image = gr.Image(
+        type="numpy",
+        label="Upscaled Preview",
+        height=400
+    )
+
+    download_file = gr.File(
+        label="Download Full PNG"
+    )
 
     run_button.click(
         fn=upscale,
-        inputs=input_image,
-        outputs=output_image,
+        inputs=[
+            input_image,
+            tile_width,
+            tile_height,
+            use_ratio,
+            ratio_name,
+            target_long_side
+        ],
+        outputs=[
+            output_image,
+            download_file
+        ],
         show_progress=True
     )