Update app.py

app.py

@@ -6,14 +6,17 @@ import numpy as np
 from PIL import Image
 from image_gen_aux import UpscaleWithModel
 
-# -----------------------------
+# ---------------------------------
 # Model configuration
-# -----------------------------
+# ---------------------------------
 
-MODEL_ID = "2xEvangelion_dat2"
-UPSCALE_MODEL = None
+MODEL_MAP = {
+    "Anime Enhancer": "OzzyGT/4xRemacri",
+    "Photo Enhancer": "OzzyGT/4xNomosWebPhoto_RealPLKSR",
+}
+
+MODEL_CACHE = {}
 
-# Aspect ratios
 RATIO_MAP = {
     "16:9": (16, 9),
     "9:16": (9, 16),
@@ -24,91 +27,83 @@ RATIO_MAP = {
     "3:2": (3, 2),
 }
 
-# -----------------------------
-# Load model once (lazy load)
-# -----------------------------
+# ---------------------------------
+# Helpers
+# ---------------------------------
 
-def get_upscaler():
-    global UPSCALE_MODEL
-    if UPSCALE_MODEL is None:
-        UPSCALE_MODEL = UpscaleWithModel.from_pretrained(MODEL_ID).to("cuda")
-    return UPSCALE_MODEL
+def get_upscaler(enhancer_type: str):
+    global MODEL_CACHE
 
-# -----------------------------
-# Resize image to selected ratio
-# -----------------------------
+    if enhancer_type not in MODEL_CACHE:
+        model_id = MODEL_MAP[enhancer_type]
+        MODEL_CACHE[enhancer_type] = UpscaleWithModel.from_pretrained(model_id).to("cuda")
 
-def fit_to_ratio(image: Image.Image, ratio_name: str, long_side: int):
+    return MODEL_CACHE[enhancer_type]
 
+
+def fit_to_ratio(image: Image.Image, ratio_name: str, long_side: int) -> Image.Image:
     rw, rh = RATIO_MAP[ratio_name]
 
     if rw >= rh:
-        width = long_side
-        height = int(long_side * rh / rw)
+        width = int(long_side)
+        height = round(long_side * rh / rw)
     else:
-        height = long_side
-        width = int(long_side * rw / rh)
+        height = int(long_side)
+        width = round(long_side * rw / rh)
 
-    width -= width % 2
-    height -= height % 2
+    width = max(2, width - (width % 2))
+    height = max(2, height - (height % 2))
 
     return image.resize((width, height), Image.LANCZOS)
 
-# -----------------------------
-# GPU upscale function
-# -----------------------------
+
+# ---------------------------------
+# GPU function
+# ---------------------------------
 
 @spaces.GPU
 def upscale(
     image,
+    enhancer_type,
     tile_width,
     tile_height,
     use_ratio,
     ratio_name,
-    target_long_side
+    target_long_side,
 ):
-
     if image is 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()
+    upscaler = get_upscaler(enhancer_type)
 
     out = upscaler(
         img,
         tiling=True,
         tile_width=int(tile_width),
-        tile_height=int(tile_height)
+        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
-    )
+    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")
+    gr.Markdown("# Image Enhancer")
 
     input_image = gr.Image(
         type="numpy",
@@ -116,8 +111,13 @@ with gr.Blocks() as demo:
         height=400
     )
 
-    with gr.Row():
+    enhancer_type = gr.Radio(
+        choices=["Anime Enhancer", "Photo Enhancer"],
+        value="Anime Enhancer",
+        label="Enhancer Type"
+    )
 
+    with gr.Row():
         tile_width = gr.Number(
             value=1024,
             label="Tile Width",
@@ -131,7 +131,6 @@ with gr.Blocks() as demo:
         )
 
     with gr.Group():
-
         use_ratio = gr.Checkbox(
             label="Use Aspect Ratio Preset",
             value=False
@@ -149,11 +148,11 @@ with gr.Blocks() as demo:
             precision=0
         )
 
-    run_button = gr.Button("Upscale")
+    run_button = gr.Button("Enhance 4x")
 
     output_image = gr.Image(
         type="numpy",
-        label="Upscaled Preview",
+        label="Enhanced Preview",
         height=400
     )
 
@@ -165,15 +164,16 @@ with gr.Blocks() as demo:
         fn=upscale,
         inputs=[
             input_image,
+            enhancer_type,
             tile_width,
             tile_height,
             use_ratio,
             ratio_name,
-            target_long_side
+            target_long_side,
         ],
         outputs=[
             output_image,
-            download_file
+            download_file,
         ],
         show_progress=True
     )