Update app.py

app.py

@@ -3,6 +3,7 @@ import uuid
 import gradio as gr
 import spaces
 from PIL import Image
+from huggingface_hub import hf_hub_download
 from image_gen_aux import UpscaleWithModel
 
 # ---------------------------------
@@ -11,16 +12,29 @@ from image_gen_aux import UpscaleWithModel
 
 BASE_TMP_DIR = "/tmp/image_enhancer"
 ENHANCED_DIR = os.path.join(BASE_TMP_DIR, "enhanced")
+MODEL_DIR = os.path.join(BASE_TMP_DIR, "models")
+
 os.makedirs(ENHANCED_DIR, exist_ok=True)
+os.makedirs(MODEL_DIR, exist_ok=True)
 
 # ---------------------------------
 # Model configuration
 # ---------------------------------
-# Swap only these repo IDs if you want to test different models later.
 
-ANIME_HIFI_2X_MODEL = "Phips/2xHFA2kCompact"
-FAST_ANIME_4X_MODEL = "OzzyGT/4xRemacri"
-FAST_PHOTO_4X_MODEL = "OzzyGT/4xNomosWebPhoto_RealPLKSR"
+MODEL_SPECS = {
+    "AnimeSharp": {
+        "repo_id": "Kim2091/AnimeSharp",
+        "filename": "4x-AnimeSharp.pth",
+    },
+    "UltraSharp": {
+        "repo_id": "Kim2091/UltraSharp",
+        "filename": "4x-UltraSharp.pth",
+    },
+    "UltraMix Balanced": {
+        "repo_id": "LykosAI/Upscalers",
+        "filename": "UltraMix/4x-UltraMix_Balanced.pth",
+    },
+}
 
 MODEL_CACHE = {}
 
@@ -35,9 +49,7 @@ RATIO_MAP = {
 }
 
 MODE_CHOICES = [
-    "2x High Fidelity",
     "4x High Fidelity",
-    "4x Fast / Medium Fidelity",
     "8x Multi-Pass (Drift Likely)",
 ]
 
@@ -53,11 +65,23 @@ REDUCTION_DISCLAIMER = (
 # Helpers
 # ---------------------------------
 
-def get_model(model_id: str):
+def get_model(model_name: str):
     global MODEL_CACHE
-    if model_id not in MODEL_CACHE:
-        MODEL_CACHE[model_id] = UpscaleWithModel.from_pretrained(model_id).to("cuda")
-    return MODEL_CACHE[model_id]
+
+    if model_name in MODEL_CACHE:
+        return MODEL_CACHE[model_name]
+
+    spec = MODEL_SPECS[model_name]
+
+    local_path = hf_hub_download(
+        repo_id=spec["repo_id"],
+        filename=spec["filename"],
+        local_dir=MODEL_DIR,
+        local_dir_use_symlinks=False,
+    )
+
+    MODEL_CACHE[model_name] = UpscaleWithModel.from_pretrained(local_path).to("cuda")
+    return MODEL_CACHE[model_name]
 
 
 def get_tile_dimensions(ratio_name: str, tile_preset: str):
@@ -109,17 +133,19 @@ def build_stats_markdown(
     mode_name: str,
     reduction_choice: str,
     reduction_applied: bool,
+    model_name: str,
 ):
     reduction_status = reduction_choice if reduction_applied else "Ignored / Not Applied"
 
     return (
+        f"**Model:** {model_name}  \n"
+        f"**Mode:** {mode_name}  \n"
+        f"**Export Format:** {export_format}  \n"
+        f"**Hi-Fi Output Reduction:** {reduction_status}  \n\n"
         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"**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)}"
     )
 
@@ -150,8 +176,8 @@ def apply_output_reduction(img: Image.Image, reduction_choice: str):
     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)
+def upscale_once(img: Image.Image, model_name: str, tile_width: int, tile_height: int):
+    upscaler = get_model(model_name)
     out = upscaler(
         img,
         tiling=True,
@@ -167,33 +193,20 @@ def upscale_once(img: Image.Image, model_id: str, tile_width: int, tile_height:
 
 def run_mode_pipeline(
     img: Image.Image,
-    enhancer_type: str,
+    model_name: str,
     mode_name: str,
     tile_width: int,
     tile_height: int,
 ):
-    # 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
-
-    # 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)
+        return upscale_once(img, model_name, tile_width, tile_height)
 
     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
+        first = upscale_once(img, model_name, tile_width, tile_height)
+        second = first.resize((img.width * 8, img.height * 8), Image.LANCZOS)
+        return second.convert("RGB")
 
-    return upscale_once(img, ANIME_HIFI_2X_MODEL, tile_width, tile_height)
+    return upscale_once(img, model_name, tile_width, tile_height)
 
 
 def save_output_image(output_img: Image.Image, export_format: str):
@@ -210,10 +223,6 @@ 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
 # ---------------------------------
@@ -221,7 +230,7 @@ def is_hifi_mode(mode_name: str):
 @spaces.GPU
 def enhance_image(
     reduction_choice,
-    enhancer_type,
+    model_name,
     mode_name,
     ratio_name,
     tile_preset,
@@ -238,14 +247,14 @@ def enhance_image(
 
     enhanced_img = run_mode_pipeline(
         img=original_img,
-        enhancer_type=enhancer_type,
+        model_name=model_name,
         mode_name=mode_name,
         tile_width=tile_width,
         tile_height=tile_height,
     )
 
     reduction_applied = False
-    if is_hifi_mode(mode_name) and reduction_choice != "Off":
+    if mode_name == "4x High Fidelity" and reduction_choice != "Off":
         enhanced_img = apply_output_reduction(enhanced_img, reduction_choice)
         reduction_applied = True
 
@@ -264,6 +273,7 @@ def enhance_image(
         mode_name=mode_name,
         reduction_choice=reduction_choice,
         reduction_applied=reduction_applied,
+        model_name=model_name,
     )
 
     return enhanced_img, output_path, stats_markdown
@@ -287,17 +297,17 @@ with gr.Blocks() as demo:
             label="Reduction Amount"
         )
 
-    # 1. Enhancer / Mode box
+    # 1. Model / Mode box
     with gr.Group():
-        enhancer_type = gr.Radio(
-            choices=["Anime Enhancer", "Photo Enhancer"],
-            value="Anime Enhancer",
-            label="Enhancer Type"
+        model_name = gr.Radio(
+            choices=["AnimeSharp", "UltraSharp", "UltraMix Balanced"],
+            value="AnimeSharp",
+            label="Reconstruction Model"
         )
 
         mode_name = gr.Radio(
             choices=MODE_CHOICES,
-            value="2x High Fidelity",
+            value="4x High Fidelity",
             label="Processing Mode"
         )
 
@@ -307,7 +317,7 @@ with gr.Blocks() as demo:
 
         ratio_name = gr.Radio(
             choices=["16:9", "9:16", "4:5", "1:1", "5:4", "2:3", "3:2"],
-            value="1:1",
+            value="2:3",
             label="Aspect Ratio"
         )
 
@@ -318,7 +328,7 @@ with gr.Blocks() as demo:
         )
 
         tile_display = gr.Markdown(
-            value=update_tile_display("1:1", "768")
+            value=update_tile_display("2:3", "768")
         )
 
     # 2.5 Output Settings
@@ -380,7 +390,7 @@ with gr.Blocks() as demo:
         fn=enhance_image,
         inputs=[
             reduction_choice,
-            enhancer_type,
+            model_name,
             mode_name,
             ratio_name,
             tile_preset,