Update app.py

app.py

@@ -570,25 +570,25 @@ def step2_upscale(
     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)
+    denoise_strength: float = 0.5,
+    face_enhance: bool = False,
+    batch_size: int = 16,          # << NEW
+    max_images: int = 0,           # << NEW (0 = all)
 ):
-    """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)
+    """Upscale frames **in batches** with live progress updates.
+    Streams: "Upscaling… 20% · 80/100 remaining (batch 2/10)"
     """
     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")
+        msg = ("Real-ESRGAN not available. Check requirements.txt includes: --prefer-binary, "
+               "numpy==1.26.4, scipy==1.11.4, scikit-image==0.22.0, opencv-python-headless, "
+               "torch==2.2.2, realesrgan==0.3.0, basicsr==1.4.2, pillow, gradio.")
         yield None, None, msg, prog_html
         return
 
-    # Choose source: uploaded images take priority, else frames from Step 1
+    # Source: uploaded > 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]
+        # Use direct file paths; no extra staging copy
+        src_paths = [str(Path(f.name)) for f in uploaded_imgs]
     else:
         src_paths = frames_list or []
 
@@ -596,10 +596,34 @@ def step2_upscale(
         yield None, None, "No images provided. Upload files or run Step 1 first.", prog_html
         return
 
-    # Map UI -> internal model id and clamp outscale
+    # Optional cap
+    try:
+        max_images = int(max_images or 0)
+    except Exception:
+        max_images = 0
+    if max_images > 0:
+        src_paths = src_paths[:max_images]
+
+    # Batch size
+    try:
+        batch_size = max(1, int(batch_size or 1))
+    except Exception:
+        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")
-    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"
     half = (precision == "half") and (device == "cuda")
     upsampler = get_realesrganer(model_id, scale, tile, half, device=device)
@@ -611,36 +635,37 @@ def step2_upscale(
     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")
-            out_img.save(out_file, quality=95)
-            up_paths.append(out_file)
-        except Exception:
-            pass
+    # Process in batches
+    for i in range(0, total, batch_size):
+        batch = src_paths[i:i+batch_size]
 
-        done += 1
-        if total:
-            pct = int(round((done / total) * 100))
-            remaining = max(0, total - done)
-            if pct != last_pct:
-                label = f"Upscaling… {pct}% · {remaining}/{total} remaining"
-                prog_html = render_progress(pct, label)
-                # Stream a progress update (gallery/zip are None until the end)
-                yield None, None, label, prog_html
-                last_pct = pct
+        for fp in batch:
+            try:
+                with Image.open(fp) as im:
+                    img = im.convert("RGB")
+                    output, _ = upsampler.enhance(np.array(img), outscale=scale)
+                out_img = Image.fromarray(output)
+                out_file = out_dir / (Path(fp).stem + ".jpg")
+                out_img.save(out_file, quality=95)
+                up_paths.append(out_file)
+            except Exception:
+                # continue on errors
+                pass
+            finally:
+                done += 1
+
+        # Emit progress after each batch
+        pct = int(round((done / total) * 100)) if total else 0
+        remaining = max(0, total - done)
+        label = f"Upscaling… {pct}% · {remaining}/{total} remaining (batch {i//batch_size+1}/{(total+batch_size-1)//batch_size})"
+        prog_html = render_progress(pct, label)
+        yield None, None, label, prog_html
 
     if not up_paths:
         yield None, None, "Upscaling produced no outputs.", prog_html
         return
 
-    # Final outputs
     up_paths = sorted(up_paths, key=_natural_key)
     gallery = sample_paths(up_paths, 30)
     zip_path = work / "upscaled.zip"
@@ -649,9 +674,10 @@ def step2_upscale(
             zf.write(p, p.name)
 
     final_label = (f"Upscaled: {len(up_paths)} | Model: {ui_model_name}→{model_id} | "
-                   f"Scale: x{scale} | Tile: {tile} | Precision: {precision}")
+                   f"Scale: x{scale} | Tile: {tile} | Precision: {precision} | Batch: {batch_size}")
     yield gallery, str(zip_path), final_label, render_progress(100.0, "Upscaling complete")
 
+
 # ───────────────── Encode (Step 3) — supports uploaded frames/ZIP & optional audio source
 
 def prepare_frames_from_upload(files: List[gr.File] | None, prefix: str = "enc") -> Tuple[Optional[str], Optional[str]]:
@@ -957,7 +983,10 @@ def build_ui():
                 with gr.Row():
                     tile = gr.Number(value=0, label="Tile size (0 = auto)")
                     precision = gr.Dropdown(["auto", "half", "full"], value="auto", label="Precision (GPU=half, CPU=full)")
-                
+                with gr.Row():
+                    batch_size = gr.Number(value=16, precision=0, label="Batch size (images per batch)")
+                    max_images = gr.Number(value=0, precision=0, label="Max images to process (0 = all)")
+
                 with gr.Row():
                     btn_upscale = gr.Button("Step 2: Upscale", variant="primary")
                 
@@ -968,7 +997,7 @@ def build_ui():
 
                 btn_upscale.click(
                     step2_upscale,
-                    inputs=[frames_state, ui_model_name, outscale, tile, precision, prog2, imgs_override, denoise_strength, face_enhance],
+                    inputs=[frames_state, ui_model_name, outscale, tile, precision, prog2, imgs_override, denoise_strength, face_enhance, batch_size, max_images],
                     outputs=[gallery_up, zip_up, details2, prog2],
                 )