Update app.py

app.py

@@ -1,9 +1,7 @@
 import os
 import uuid
-import math
 import gradio as gr
 import spaces
-import numpy as np
 from PIL import Image
 from image_gen_aux import UpscaleWithModel
 
@@ -13,17 +11,16 @@ from image_gen_aux import UpscaleWithModel
 
 BASE_TMP_DIR = "/tmp/image_enhancer"
 ENHANCED_DIR = os.path.join(BASE_TMP_DIR, "enhanced")
-
 os.makedirs(ENHANCED_DIR, exist_ok=True)
 
 # ---------------------------------
 # Model configuration
 # ---------------------------------
+# Swap only these repo IDs if you want to test different models later.
 
-MODEL_MAP = {
-    "Anime Enhancer": "OzzyGT/4xRemacri",
-    "Photo Enhancer": "OzzyGT/4xNomosWebPhoto_RealPLKSR",
-}
+ANIME_HIFI_2X_MODEL = "Phips/2xHFA2kCompact"
+FAST_ANIME_4X_MODEL = "OzzyGT/4xRemacri"
+FAST_PHOTO_4X_MODEL = "OzzyGT/4xNomosWebPhoto_RealPLKSR"
 
 MODEL_CACHE = {}
 
@@ -37,18 +34,30 @@ RATIO_MAP = {
     "3:2": (3, 2),
 }
 
+MODE_CHOICES = [
+    "2x High Fidelity",
+    "4x High Fidelity",
+    "4x Fast / Medium Fidelity",
+    "8x Multi-Pass (Drift Likely)",
+]
+
+REDUCTION_CHOICES = ["Off", "80%", "85%", "90%"]
+
+REDUCTION_DISCLAIMER = (
+    "Hi-Fi Output offers Post-Processing Size Reduction to further improve results. "
+    "Mode is entirely Optional and is defaulted at Off. Please note, this does not work "
+    "with Fast modes. If toggled, it will not be applied."
+)
+
 # ---------------------------------
 # Helpers
 # ---------------------------------
 
-def get_upscaler(enhancer_type: str):
+def get_model(model_id: str):
     global MODEL_CACHE
-
-    if enhancer_type not in MODEL_CACHE:
-        model_id = MODEL_MAP[enhancer_type]
-        MODEL_CACHE[enhancer_type] = UpscaleWithModel.from_pretrained(model_id).to("cuda")
-
-    return MODEL_CACHE[enhancer_type]
+    if model_id not in MODEL_CACHE:
+        MODEL_CACHE[model_id] = UpscaleWithModel.from_pretrained(model_id).to("cuda")
+    return MODEL_CACHE[model_id]
 
 
 def get_tile_dimensions(ratio_name: str, tile_preset: str):
@@ -77,8 +86,7 @@ def update_tile_display(ratio_name: str, tile_preset: str):
 
 
 def format_megapixels(width: int, height: int) -> str:
-    mp = (width * height) / 1_000_000
-    return f"{mp:.2f} MP"
+    return f"{(width * height) / 1_000_000:.2f} MP"
 
 
 def format_file_size(num_bytes: int) -> str:
@@ -98,25 +106,52 @@ def build_stats_markdown(
     enhanced_height: int,
     file_size_bytes: int,
     export_format: str,
-    scale_mode: str,
+    mode_name: str,
+    reduction_choice: str,
+    reduction_applied: bool,
 ):
+    reduction_status = reduction_choice if reduction_applied else "Ignored / Not Applied"
+
     return (
         f"**Original Dimensions:** {original_width} × {original_height}px  \n"
         f"**Original Megapixels:** {format_megapixels(original_width, original_height)}  \n\n"
         f"**Enhanced Dimensions:** {enhanced_width} × {enhanced_height}px  \n"
         f"**Enhanced Megapixels:** {format_megapixels(enhanced_width, enhanced_height)}  \n\n"
-        f"**Scale Mode:** {scale_mode}  \n"
+        f"**Mode:** {mode_name}  \n"
         f"**Export Format:** {export_format}  \n"
+        f"**Hi-Fi Output Reduction:** {reduction_status}  \n"
         f"**Saved File Size:** {format_file_size(file_size_bytes)}"
     )
 
 
-def upscale_once_with_model(
-    img: Image.Image,
-    upscaler,
-    tile_width: int,
-    tile_height: int,
-) -> Image.Image:
+def reduction_factor_from_choice(choice: str):
+    mapping = {
+        "80%": 0.80,
+        "85%": 0.85,
+        "90%": 0.90,
+    }
+    return mapping.get(choice, 1.0)
+
+
+def apply_output_reduction(img: Image.Image, reduction_choice: str):
+    factor = reduction_factor_from_choice(reduction_choice)
+    if factor >= 1.0:
+        return img
+
+    new_width = max(2, int(round(img.width * factor)))
+    new_height = max(2, int(round(img.height * factor)))
+
+    new_width -= new_width % 2
+    new_height -= new_height % 2
+
+    new_width = max(2, new_width)
+    new_height = max(2, new_height)
+
+    return img.resize((new_width, new_height), Image.LANCZOS)
+
+
+def upscale_once(img: Image.Image, model_id: str, tile_width: int, tile_height: int):
+    upscaler = get_model(model_id)
     out = upscaler(
         img,
         tiling=True,
@@ -130,31 +165,35 @@ def upscale_once_with_model(
     return out.convert("RGB")
 
 
-def run_scale_pipeline(
+def run_mode_pipeline(
     img: Image.Image,
-    upscaler,
+    enhancer_type: str,
+    mode_name: str,
     tile_width: int,
     tile_height: int,
-    scale_mode: str,
-) -> Image.Image:
-    # 4x = one model pass
-    if scale_mode == "4x":
-        return upscale_once_with_model(img, upscaler, tile_width, tile_height)
+):
+    # Conservative Hi-Fi branch uses true 2x model.
+    if mode_name == "2x High Fidelity":
+        return upscale_once(img, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+
+    if mode_name == "4x High Fidelity":
+        first = upscale_once(img, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+        second = upscale_once(first, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+        return second
 
-    # 8x Experimental = one real 4x pass, then high-quality resize to 8x total
-    if scale_mode == "8x Experimental":
-        first = upscale_once_with_model(img, upscaler, tile_width, tile_height)
-        target_w = img.width * 8
-        target_h = img.height * 8
-        return first.resize((target_w, target_h), Image.LANCZOS).convert("RGB")
+    # Rough capability / direct 4x branch
+    if mode_name == "4x Fast / Medium Fidelity":
+        if enhancer_type == "Photo Enhancer":
+            return upscale_once(img, FAST_PHOTO_4X_MODEL, tile_width, tile_height)
+        return upscale_once(img, FAST_ANIME_4X_MODEL, tile_width, tile_height)
 
-    # 16x Experimental = two real 4x passes = 16x total
-    if scale_mode == "16x Experimental":
-        first = upscale_once_with_model(img, upscaler, tile_width, tile_height)
-        second = upscale_once_with_model(first, upscaler, tile_width, tile_height)
-        return second.convert("RGB")
+    if mode_name == "8x Multi-Pass (Drift Likely)":
+        first = upscale_once(img, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+        second = upscale_once(first, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+        third = upscale_once(second, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+        return third
 
-    return upscale_once_with_model(img, upscaler, tile_width, tile_height)
+    return upscale_once(img, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
 
 
 def save_output_image(output_img: Image.Image, export_format: str):
@@ -171,40 +210,45 @@ def save_output_image(output_img: Image.Image, export_format: str):
     return path
 
 
+def is_hifi_mode(mode_name: str):
+    return mode_name in {"2x High Fidelity", "4x High Fidelity"}
+
+
 # ---------------------------------
 # GPU function
 # ---------------------------------
 
 @spaces.GPU
 def enhance_image(
+    reduction_choice,
     enhancer_type,
+    mode_name,
     ratio_name,
     tile_preset,
-    scale_mode,
     export_format,
     input_image,
 ):
     if input_image is None:
-        return (
-            None,
-            None,
-            "No stats available yet."
-        )
+        return None, None, "No stats available yet."
 
     original_img = Image.fromarray(input_image).convert("RGB")
     original_width, original_height = original_img.size
 
     tile_width, tile_height = get_tile_dimensions(ratio_name, tile_preset)
-    upscaler = get_upscaler(enhancer_type)
 
-    enhanced_img = run_scale_pipeline(
+    enhanced_img = run_mode_pipeline(
         img=original_img,
-        upscaler=upscaler,
+        enhancer_type=enhancer_type,
+        mode_name=mode_name,
         tile_width=tile_width,
         tile_height=tile_height,
-        scale_mode=scale_mode,
     )
 
+    reduction_applied = False
+    if is_hifi_mode(mode_name) and reduction_choice != "Off":
+        enhanced_img = apply_output_reduction(enhanced_img, reduction_choice)
+        reduction_applied = True
+
     enhanced_width, enhanced_height = enhanced_img.size
 
     output_path = save_output_image(enhanced_img, export_format)
@@ -217,7 +261,9 @@ def enhance_image(
         enhanced_height=enhanced_height,
         file_size_bytes=file_size_bytes,
         export_format=export_format,
-        scale_mode=scale_mode,
+        mode_name=mode_name,
+        reduction_choice=reduction_choice,
+        reduction_applied=reduction_applied,
     )
 
     return enhanced_img, output_path, stats_markdown
@@ -230,7 +276,18 @@ def enhance_image(
 with gr.Blocks() as demo:
     gr.Markdown("# Image Enhancer")
 
-    # 1. Enhancer Type
+    # 0. Hi-Fi Output Reduction
+    with gr.Group():
+        gr.Markdown("### Hi-Fi Output Reduction")
+        gr.Markdown(REDUCTION_DISCLAIMER)
+
+        reduction_choice = gr.Radio(
+            choices=REDUCTION_CHOICES,
+            value="Off",
+            label="Reduction Amount"
+        )
+
+    # 1. Enhancer / Mode box
     with gr.Group():
         enhancer_type = gr.Radio(
             choices=["Anime Enhancer", "Photo Enhancer"],
@@ -238,7 +295,13 @@ with gr.Blocks() as demo:
             label="Enhancer Type"
         )
 
-    # 2. Combined Tile Settings box
+        mode_name = gr.Radio(
+            choices=MODE_CHOICES,
+            value="2x High Fidelity",
+            label="Processing Mode"
+        )
+
+    # 2. Combined Tile Settings
     with gr.Group():
         gr.Markdown("### Tile Settings")
 
@@ -258,16 +321,10 @@ with gr.Blocks() as demo:
             value=update_tile_display("1:1", "768")
         )
 
-    # Extra processing settings
+    # 2.5 Output Settings
     with gr.Group():
         gr.Markdown("### Output Settings")
 
-        scale_mode = gr.Radio(
-            choices=["4x", "8x Experimental", "16x Experimental"],
-            value="4x",
-            label="Scale Mode"
-        )
-
         export_format = gr.Radio(
             choices=["PNG", "TIFF"],
             value="PNG",
@@ -300,7 +357,7 @@ with gr.Blocks() as demo:
             label="Download new enhanced image file"
         )
 
-    # Stats box
+    # Stats
     with gr.Group():
         gr.Markdown("### Image Stats")
         stats_box = gr.Markdown(
@@ -322,10 +379,11 @@ with gr.Blocks() as demo:
     run_button.click(
         fn=enhance_image,
         inputs=[
+            reduction_choice,
             enhancer_type,
+            mode_name,
             ratio_name,
             tile_preset,
-            scale_mode,
             export_format,
             input_image,
         ],