Update app.py

app.py

@@ -532,6 +532,94 @@ def step1_extract(
 
 # ───────────────── Upscale (Step 2) — supports uploaded images OR frames from Step 1
 
+# Manual-batch Step 2 helpers (resumable, click-to-advance)
+def _ensure_dir(p: Path) -> Path:
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+def _save_zip_of_dir(dir_path: Path, zip_path: Path) -> str:
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in sorted(dir_path.glob("*.*"), key=_natural_key):
+            if p.suffix.lower() in [".jpg", ".jpeg", ".png"]:
+                zf.write(p, p.name)
+    return str(zip_path)
+
+def _list_image_paths_from_upload(files: List[gr.File] | None) -> List[str]:
+    if not files: return []
+    return [str(Path(f.name)) for f in files if Path(f.name).suffix.lower() in [".jpg",".jpeg",".png"]]
+
+def _build_gallery_from_dir(dir_path: Path, n: int = 30) -> List[str]:
+    paths = sorted(list(dir_path.glob("*.jpg")) + list(dir_path.glob("*.png")), key=_natural_key)
+    return sample_paths(paths, n)
+
+def step2_prepare_sources(frames_list, uploaded_imgs, max_images):
+    src = _list_image_paths_from_upload(uploaded_imgs) or (frames_list or [])
+    if not src:
+        return [], "", 0, 0, "No images found. Upload files or run Step 1 first.", render_progress(0.0, "Idle")
+    try:
+        max_images = int(max_images or 0)
+    except Exception:
+        max_images = 0
+    if max_images > 0:
+        src = src[:max_images]
+    work = Path(tempfile.mkdtemp(prefix="up_manual_"))
+    out_dir = _ensure_dir(work / "upscaled")
+    total = len(src)
+    done_idx = 0
+    msg = f"Sources loaded: {total} image(s). Click 'Process Next Batch' to start."
+    prog = render_progress(0.0, "Ready")
+    return src, str(out_dir), done_idx, total, msg, prog
+
+def step2_process_next_batch(
+    up_src_paths, up_out_dir, up_done_idx, up_total,
+    ui_model_name, outscale, tile, precision, denoise_strength, face_enhance, batch_size,
+):
+    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))
+    upsampler = get_realesrganer(model_id, scale, tile, half, device=device)
+
+    start = int(up_done_idx or 0)
+    end = min(start + batch_size, int(up_total or 0))
+    out_dir = Path(up_out_dir)
+
+    if start >= up_total:
+        gallery = _build_gallery_from_dir(out_dir, 30)
+        zip_path = Path(out_dir.parent) / "upscaled.zip"
+        zip_file = _save_zip_of_dir(out_dir, zip_path)
+        prog = render_progress(100.0, "All images processed")
+        details = f"Done. Total upscaled: {len(list(out_dir.glob('*.jpg')))+len(list(out_dir.glob('*.png')))}"
+        return gallery, zip_file, details, prog, start, up_out_dir
+
+    processed_now = 0
+    for fp in up_src_paths[start:end]:
+        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
+        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}).")
+    prog = render_progress(pct, f"Upscaling… {pct}%")
+
+    gallery = _build_gallery_from_dir(out_dir, 30)
+    zip_path = Path(out_dir.parent) / "upscaled.zip"
+    zip_file = _save_zip_of_dir(out_dir, zip_path)
+    return gallery, zip_file, label, prog, next_idx, up_out_dir
+
+
+
 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)
@@ -874,7 +962,12 @@ def build_ui():
         frames_dir_state = gr.State("")   # str
         prefix_state = gr.State("")       # str
         fps_state = gr.State(30.0)         # float
-
+        # Shared Step 2 states (manual batching)
+        up_src_paths_state = gr.State([])     # list[str] absolute paths to process
+        up_out_dir_state   = gr.State("")     # str: output dir path
+        up_done_idx_state  = gr.State(0)      # int: next index to start from
+        up_total_state     = gr.State(0)      # int: total images
+        
         with gr.Tabs():
             # STEP 1
             with gr.Tab("Step 1 · Extract Frames"):
@@ -970,37 +1063,44 @@ def build_ui():
                             value="RealESRGAN_x4plus",
                             show_label=True
                         )
-                        denoise_strength = gr.Slider(
-                            label="Denoise Strength (only for realesr-general-x4v3)",
-                            minimum=0, maximum=1, step=0.1, value=0.5
-                        )
-                        outscale = gr.Slider(
-                            label="Resolution upscale",
-                            minimum=1, maximum=6, step=1, value=4, show_label=True
-                        )
-                        face_enhance = gr.Checkbox(label="Face Enhancement (GFPGAN)", value=False)
+                        denoise_strength = gr.Slider(0, 1, value=0.5, step=0.1, label="Denoise (only general-x4v3)")
+                        outscale = gr.Slider(1, 6, value=4, step=1, label="Resolution upscale")
+                        face_enhance = gr.Checkbox(value=False, label="Face Enhancement (GFPGAN)")
                 
                 with gr.Row():
-                    tile = gr.Number(value=0, label="Tile size (0 = auto)")
+                    tile = gr.Number(value=256, label="Tile size (try 128 if OOM; 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)")
+                    batch_size = gr.Number(value=12, precision=0, label="Batch size per click")
                     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")
+                    btn_prepare = gr.Button("Step 2: Load / Reset Sources", variant="secondary")
+                    btn_next    = gr.Button("Process Next Batch", variant="primary")
                 
                 prog2 = gr.HTML(render_progress(0.0, "Idle"))
                 gallery_up = gr.Gallery(label="Upscaled preview (30 sampled)", columns=6, height=480)
                 zip_up = gr.File(label="Download upscaled ZIP")
                 details2 = gr.Markdown("")
-
-                btn_upscale.click(
-                    step2_upscale,
-                    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],
+                
+                # 1) load/reset sources
+                btn_prepare.click(
+                    step2_prepare_sources,
+                    inputs=[frames_state, imgs_override, max_images],
+                    outputs=[up_src_paths_state, up_out_dir_state, up_done_idx_state, up_total_state, details2, prog2]
+                )
+                
+                # 2) process one batch per click
+                btn_next.click(
+                    step2_process_next_batch,
+                    inputs=[
+                        up_src_paths_state, up_out_dir_state, up_done_idx_state, up_total_state,
+                        ui_model_name, outscale, tile, precision, denoise_strength, face_enhance, batch_size
+                    ],
+                    outputs=[gallery_up, zip_up, details2, prog2, up_done_idx_state, up_out_dir_state]
                 )
 
+
             # STEP 3 — Re-encode
             with gr.Tab("Step 3 · Re-encode Video"):
                 gr.Markdown("Use frames from Step 1 **or** upload a frames ZIP / images. Optionally provide a video for audio track.")