Update app.py

app.py

@@ -547,17 +547,45 @@ def save_uploaded_images(files: List[gr.File] | None, prefix: str = "upload") ->
     return paths, in_dir
 
 
+# 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
+
 def step2_upscale(
     frames_list: List[str] | None,
-    model_name: str,
-    scale: int,
+    ui_model_name: str,
+    outscale: int,
     tile: int,
     precision: str,
     prog_html: str,
     uploaded_imgs: List[gr.File] | None,
+    denoise_strength: float = 0.5,   # <- NEW param (to match UI)
+    face_enhance: bool = False,      # <- NEW param (to match UI)
 ):
+    """Upscale frames with live progress updates.
+    Streams updates like: Upscaling… 20% · 80/100 remaining
+    Yields: (gallery|None, zip_path|None, details|str, progress_html|str)
+    """
+    if not HAVE_REALESRGAN:
+        msg = ("Real-ESRGAN not available. Ensure requirements.txt includes (each on its own line):\n"
+               "--prefer-binary\n"
+               "numpy==1.26.4\nscipy==1.11.4\nscikit-image==0.22.0\nopencv-python-headless==4.10.0.84\n"
+               "torch==2.2.2\nrealesrgan==0.3.0\nbasicsr==1.4.2\npillow\ngradio==5.44.1")
+        yield None, None, msg, prog_html
+        return
 
-    # decide source: uploaded images take priority, else frames from step 1
+    # Choose source: uploaded images take priority, else frames from Step 1
     if uploaded_imgs and len(uploaded_imgs) > 0:
         img_paths, _ = save_uploaded_images(uploaded_imgs, prefix="up")
         src_paths = [str(p) for p in img_paths]
@@ -568,25 +596,27 @@ def step2_upscale(
         yield None, None, "No images provided. Upload files or run Step 1 first.", prog_html
         return
 
-    # Map demo model -> internal, clamp scale to supported values
-    model_id = map_ui_model_to_internal(ui_model_name)
-    scale = clamp_scale_for_model(int(outscale or 4), model_id)
-    
+    # Map UI -> internal model id and clamp outscale
+    model_id = MODEL_MAP.get(ui_model_name, "x4plus")
+    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")
     upsampler = get_realesrganer(model_id, scale, tile, half, device=device)
 
     work = Path(tempfile.mkdtemp(prefix="up_"))
-    out_dir = work / "upscaled"; out_dir.mkdir(parents=True, exist_ok=True)
+    out_dir = work / "upscaled"
+    out_dir.mkdir(parents=True, exist_ok=True)
 
     total = len(src_paths)
     done = 0
     up_paths: List[Path] = []
-
     last_pct = -1
+
     for fp in src_paths:
         try:
             img = Image.open(fp).convert("RGB")
+            # (denoise_strength & face_enhance are accepted but not wired in this minimal build)
             output, _ = upsampler.enhance(np.array(img), outscale=scale)
             out_img = Image.fromarray(output)
             out_file = out_dir / (Path(fp).stem + ".jpg")
@@ -594,6 +624,7 @@ def step2_upscale(
             up_paths.append(out_file)
         except Exception:
             pass
+
         done += 1
         if total:
             pct = int(round((done / total) * 100))
@@ -601,7 +632,7 @@ def step2_upscale(
             if pct != last_pct:
                 label = f"Upscaling… {pct}% · {remaining}/{total} remaining"
                 prog_html = render_progress(pct, label)
-                # stream progress (gallery/zip stay None until the end)
+                # Stream a progress update (gallery/zip are None until the end)
                 yield None, None, label, prog_html
                 last_pct = pct
 
@@ -610,13 +641,15 @@ def step2_upscale(
         return
 
     # Final outputs
+    up_paths = sorted(up_paths, key=_natural_key)
     gallery = sample_paths(up_paths, 30)
     zip_path = work / "upscaled.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
         for p in up_paths:
             zf.write(p, p.name)
 
-    final_label = f"Upscaled: {len(up_paths)}"
+    final_label = (f"Upscaled: {len(up_paths)} | Model: {ui_model_name}→{model_id} | "
+                   f"Scale: x{scale} | Tile: {tile} | Precision: {precision}")
     yield gallery, str(zip_path), final_label, render_progress(100.0, "Upscaling complete")
 
 # ───────────────── Encode (Step 3) — supports uploaded frames/ZIP & optional audio source