Update app.py

app.py

@@ -183,7 +183,7 @@ else:
 # Helpers
 # ─────────────────────────────────────────────────────────────
 
-    # Map UI model names (demo) to our internal model IDs
+# Map UI model names (demo) to our internal model IDs
 def map_ui_model_to_internal(ui_name: str) -> str:
     mapping = {
         "RealESRGAN_x4plus": "x4plus",
@@ -202,26 +202,6 @@ def clamp_scale_for_model(outscale: int, model_id: str) -> int:
     # For x4plus / x4plus-anime, force 4 (ignore 5–6)
     return 4
 
-
-def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
-    """Evenly sample up to n items across the entire list (deterministic), in natural numeric order."""
-    if not paths:
-        return []
-    paths = sorted(paths, key=_natural_key)  # ensure numeric order first
-    total = len(paths)
-    n = max(1, min(n, total))
-    if n == total:
-        return [str(p) for p in paths]
-    step = (total - 1) / (n - 1)  # cover both ends
-    idxs = [round(i * step) for i in range(n)]
-    out, seen = [], set()
-    for i in idxs:
-        if i not in seen:
-            out.append(str(paths[int(i)]))
-            seen.add(int(i))
-    return out
-
-
 def sanitize_prefix(txt: str) -> str:
     txt = (txt or "").strip()
     if not txt:
@@ -594,14 +574,33 @@ def step2_process_next_batch(
 ):
     if not up_src_paths or not up_out_dir:
         return None, None, "Load sources first.", render_progress(0.0, "Idle"), up_done_idx, up_out_dir
+
     model_id = map_ui_model_to_internal(ui_model_name)
     scale = clamp_scale_for_model(int(outscale or 4), model_id)
     device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
     half = (precision == "half") and (device == "cuda")
     tile = int(tile or 256)
     batch_size = max(1, int(batch_size or 8))
+
+    # Build upsampler
     upsampler = get_realesrganer(model_id, scale, tile, half, device=device)
 
+    # Optional: GFPGAN face enhancer
+    face_enhancer = None
+    if face_enhance:
+        try:
+            from gfpgan import GFPGANer
+            face_enhancer = GFPGANer(
+                model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
+                upscale=scale,
+                arch="clean",
+                channel_multiplier=2,
+                bg_upsampler=upsampler
+            )
+        except Exception as e:
+            print("GFPGAN load failed:", e)
+            face_enhancer = None
+
     start = int(up_done_idx or 0)
     end = min(start + batch_size, int(up_total or 0))
     out_dir = Path(up_out_dir)
@@ -619,16 +618,26 @@ def step2_process_next_batch(
         try:
             with Image.open(fp) as im:
                 img = im.convert("RGB")
-                output, _ = upsampler.enhance(np.array(img), outscale=scale)
-            Image.fromarray(output).save(out_dir / (Path(fp).stem + ".jpg"), quality=95)
-        except Exception:
-            pass
+                cv_img = np.array(img)
+
+                if face_enhancer:
+                    _, _, output = face_enhancer.enhance(
+                        cv_img, has_aligned=False, only_center_face=False, paste_back=True
+                    )
+                else:
+                    output, _ = upsampler.enhance(cv_img, outscale=scale, denoise_strength=float(denoise_strength or 0.5))
+
+                Image.fromarray(output).save(out_dir / (Path(fp).stem + ".jpg"), quality=95)
+
+        except Exception as e:
+            print("Upscale error:", e)
         processed_now += 1
 
     next_idx = end
     pct = int(round((next_idx / up_total) * 100)) if up_total else 0
     label = (f"Processed {processed_now} image(s) this batch. "
-             f"{next_idx}/{up_total} done (x{scale}, model={ui_model_name}).")
+             f"{next_idx}/{up_total} done (x{scale}, model={ui_model_name}, "
+             f"denoise={denoise_strength}, face={face_enhance}).")
     prog = render_progress(pct, f"Upscaling… {pct}%")
 
     gallery = _build_gallery_from_dir(out_dir, 30)
@@ -638,6 +647,7 @@ def step2_process_next_batch(
 
 
 
+
 def save_uploaded_images(files: List[gr.File] | None, prefix: str = "upload") -> Tuple[List[Path], Path]:
     tmp = Path(tempfile.mkdtemp(prefix="imgup_"))
     in_dir = tmp / "input"; in_dir.mkdir(parents=True, exist_ok=True)
@@ -654,19 +664,8 @@ def save_uploaded_images(files: List[gr.File] | None, prefix: str = "upload") ->
 
 
 # Map UI model names to internal IDs (make sure MODEL_MAP exists above)
-MODEL_MAP = {
-    "RealESRGAN_x4plus": "x4plus",
-    "RealESRGAN_x4plus_anime_6B": "x4plus-anime",
-    "RealESRGAN_x2plus": "x2plus",
-    # fallbacks for entries we don't wire separately here:
-    "RealESRNet_x4plus": "x4plus",
-    "realesr-general-x4v3": "x4plus",
-}
-
-def _clamp_scale_for_model(outscale: int, model_id: str) -> int:
-    if model_id == "x2plus":
-        return 2
-    return 4  # x4 models
+model_id = map_ui_model_to_internal(ui_model_name)
+scale = clamp_scale_for_model(int(outscale or 4), model_id)
 
 def step2_upscale(
     frames_list: List[str] | None,
@@ -717,17 +716,9 @@ def step2_upscale(
         batch_size = 16
 
     # Map UI model -> internal id; clamp scale to model
-    MODEL_MAP = {
-        "RealESRGAN_x4plus": "x4plus",
-        "RealESRNet_x4plus": "x4plus",           # fallback
-        "RealESRGAN_x4plus_anime_6B": "x4plus-anime",
-        "RealESRGAN_x2plus": "x2plus",
-        "realesr-general-x4v3": "x4plus",        # fallback
-    }
-    model_id = MODEL_MAP.get(ui_model_name, "x4plus")
+    model_id = map_ui_model_to_internal(ui_model_name)
+    scale = clamp_scale_for_model(int(outscale or 4), model_id)
 
-    def _clamp_scale_for_model(s: int, mid: str) -> int:
-        return 2 if mid == "x2plus" else 4
 
     scale = _clamp_scale_for_model(int(outscale or 4), model_id)
     device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"