Update app.py

app.py

@@ -31,6 +31,10 @@ import numpy
 import os
 import random
 import inspect
+from pathlib import Path
+import zipfile
+import tempfile
+
 
 from basicsr.archs.rrdbnet_arch import RRDBNet as _RRDBNet
 from basicsr.utils.download_util import load_file_from_url
@@ -160,6 +164,69 @@ def build_rrdb(scale: int, num_block: int):
     except TypeError as e:
         raise TypeError(f"RRDBNet signature not recognized: {e}")
 
+#Factor an upsampler builder
+def get_upsampler(model_name: str, outscale: int, tile: int = 256):
+    # Build the same backbone/weights as in realesrgan(), but return a ready RealESRGANer
+    if model_name == 'RealESRGAN_x4plus':
+        model = build_rrdb(scale=4, num_block=23); netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    elif model_name == 'RealESRNet_x4plus':
+        model = build_rrdb(scale=4, num_block=23); netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':
+        model = build_rrdb(scale=4, num_block=6); netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'RealESRGAN_x2plus':
+        model = build_rrdb(scale=2, num_block=23); netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif model_name == 'realesr-general-x4v3':
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu'); netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth'
+        ]
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+
+    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+    weights_dir = os.path.join(ROOT_DIR, 'weights')
+    os.makedirs(weights_dir, exist_ok=True)
+    for url in file_url:
+        fname = os.path.basename(url)
+        local_path = os.path.join(weights_dir, fname)
+        if not os.path.isfile(local_path):
+            load_file_from_url(url=url, model_dir=weights_dir, progress=True)
+
+    if model_name == 'realesr-general-x4v3':
+        model_path = [
+            os.path.join(weights_dir, 'realesr-general-x4v3.pth'),
+            os.path.join(weights_dir, 'realesr-general-wdn-x4v3.pth'),
+        ]
+        dni_weight = None  # supplied at call site if using denoise blend
+    else:
+        model_path = os.path.join(weights_dir, f"{model_name}.pth")
+        dni_weight = None
+
+    use_cuda = False
+    try:
+        use_cuda = hasattr(cv2, "cuda") and cv2.cuda.getCudaEnabledDeviceCount() > 0
+    except Exception:
+        use_cuda = False
+    gpu_id = 0 if use_cuda else None
+
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=tile or 256,
+        tile_pad=10,
+        pre_pad=10,
+        half=bool(use_cuda),
+        gpu_id=gpu_id
+    )
+    return upsampler, netscale, use_cuda, model_path
+
 
 # ────────────────────────────────────────────────────────
 # Core upscaling
@@ -285,6 +352,141 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
 
     return display_img
 
+#Add a batch upscaler that preserves filenames
+def render_progress(pct: float, text: str = "") -> str:
+    pct = max(0.0, min(100.0, float(pct)))
+    bar = f"<div style='width:100%;border:1px solid #ddd;border-radius:6px;overflow:hidden;height:12px;'><div style='height:100%;width:{pct:.1f}%;background:#3b82f6;'></div></div>"
+    label = f"<div style='font-size:12px;opacity:.8;margin-top:4px;'>{text} {pct:.1f}%</div>"
+    return bar + label
+
+def batch_realesrgan(
+    files: list,             # from gr.Files (type='filepath')
+    model_name: str,
+    denoise_strength: float,
+    face_enhance: bool,
+    outscale: int,
+    tile: int,
+    batch_size: int = 16,
+):
+    """
+    Processes multiple images in batches, preserves original file names for outputs,
+    and returns (gallery, zip_file, details, progress_html) with streamed progress.
+    """
+    # Validate
+    if not files or len(files) == 0:
+        yield None, None, "No files uploaded.", render_progress(0, "Idle")
+        return
+
+    # Build upsampler once (much faster than per-image)
+    upsampler, netscale, use_cuda, model_path = get_upsampler(model_name, outscale, tile=tile)
+
+    # Optional: face enhancer (same as your single-image path)
+    face_enhancer = None
+    if face_enhance:
+        from gfpgan import GFPGANer
+        face_enhancer = GFPGANer(
+            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler
+        )
+
+    # Prepare work/output dirs
+    work = Path(tempfile.mkdtemp(prefix="batch_up_"))
+    out_dir = work / "upscaled"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Normalize list of input paths
+    src_paths = [Path(f.name if hasattr(f, "name") else f) for f in files]
+
+    total = len(src_paths)
+    done = 0
+    out_paths = []
+
+    # If realesr-general-x4v3: support blending base + WDN via dni (optional)
+    dni_weight = None
+    if model_name == "realesr-general-x4v3":
+        # Blend [base, WDN] with user's slider
+        denoise_strength = float(denoise_strength)
+        dni_weight = [1.0 - denoise_strength, denoise_strength]
+        # RealESRGANer.enhance accepts dni_weight override via attribute on the instance
+        try:
+            upsampler.dni_weight = dni_weight
+        except Exception:
+            pass
+
+    # Process in batches (I/O and PIL open are still per-file)
+    for i in range(0, total, int(max(1, batch_size))):
+        batch = src_paths[i:i + int(max(1, batch_size))]
+        for src in batch:
+            try:
+                # Load as RGB consistently
+                from PIL import Image
+                with Image.open(src) as im:
+                    img = im.convert("RGB")
+                arr = numpy.array(img)
+                arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
+
+                if face_enhancer:
+                    _, _, output = face_enhancer.enhance(arr, has_aligned=False, only_center_face=False, paste_back=True)
+                else:
+                    output, _ = upsampler.enhance(arr, outscale=int(outscale))
+
+                # Preserve original file name & (reasonable) extension
+                orig_ext = src.suffix.lower()
+                ext = orig_ext if orig_ext in (".png", ".jpg", ".jpeg") else ".png"
+                out_path = out_dir / (src.stem + ext)
+
+                # Save (keep alpha if produced, else RGB)
+                if output.ndim == 3 and output.shape[2] == 4:
+                    cv2.imwrite(str(out_path.with_suffix(".png")), output)  # 4ch → PNG
+                    out_path = out_path.with_suffix(".png")
+                else:
+                    if ext in (".jpg", ".jpeg"):
+                        cv2.imwrite(str(out_path), output, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
+                    else:
+                        cv2.imwrite(str(out_path), output)  # PNG default
+                out_paths.append(out_path)
+            except Exception as e:
+                # Continue on errors
+                print(f"[batch] Error on {src}: {e}")
+            finally:
+                done += 1
+
+        pct = (done / total) * 100.0 if total else 0.0
+        remaining = max(0, total - done)
+        msg = f"Upscaling… {done}/{total} done · {remaining} remaining (batch {(i//batch_size)+1}/{(total+batch_size-1)//batch_size})"
+        yield None, None, msg, render_progress(pct, msg)
+
+    if not out_paths:
+        yield None, None, "No outputs produced.", render_progress(100, "Finished")
+        return
+
+    # Small even-sampled gallery for preview
+    def _sample_even(seq, n=30):
+        if not seq: return []
+        if len(seq) <= n: return [str(p) for p in seq]
+        step = (len(seq)-1) / (n-1)
+        idxs = [round(i*step) for i in range(n)]
+        seen, out = set(), []
+        for i in idxs:
+            if i not in seen:
+                out.append(str(seq[int(i)])); seen.add(int(i))
+        return out
+
+    out_paths = sorted(out_paths)  # stable
+    gallery = _sample_even(out_paths, 30)
+
+    # Zip with same file names
+    zip_path = work / "upscaled.zip"
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in out_paths:
+            zf.write(p, arcname=p.name)
+
+    details = f"Upscaled {len(out_paths)} images → {out_dir}"
+    yield gallery, str(zip_path), details, render_progress(100, "Complete")
+
 
 # ────────────────────────────────────────────────────────
 # UI
@@ -339,11 +541,44 @@ def main():
         )
         reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])
 
+                # --- Batch Upscale (multi-image) ---
+        gr.Markdown("### Batch Upscale")
+        with gr.Accordion("Batch options", open=True):
+            with gr.Row():
+                batch_files = gr.Files(
+                    label="Upload multiple images (PNG/JPG/JPEG)",
+                    type="filepath",
+                    file_types=[".png", ".jpg", ".jpeg"],
+                )
+            with gr.Row():
+                batch_tile = gr.Number(label="Tile size (0/auto → 256)", value=256, precision=0)
+                batch_size = gr.Number(label="Batch size (images per batch)", value=16, precision=0)
+
+        with gr.Row():
+            batch_btn = gr.Button("Upscale Batch", variant="primary")
+
+        batch_prog = gr.HTML(render_progress(0.0, "Idle"))
+        batch_gallery = gr.Gallery(label="Preview (sampled 30)", columns=6, height=420)
+        batch_zip = gr.File(label="Download upscaled.zip")
+        batch_details = gr.Markdown("")
+
+        # Wire it up (generator → streaming)
+        batch_btn.click(
+            fn=batch_realesrgan,
+            inputs=[batch_files, model_name, denoise_strength, face_enhance, outscale, batch_tile, batch_size],
+            outputs=[batch_gallery, batch_zip, batch_details, batch_prog],
+        )
         gr.Markdown("")  # spacer
 
     # Disable SSR (ZeroGPU + Gradio logs suggested turning this off)
     demo.launch(ssr_mode=False)  # set share=True for a public link
 
 
+def main():
+    with gr.Blocks(title="Real-ESRGAN Gradio Demo", theme="ParityError/Interstellar") as demo:
+        # ... your current UI (plus batch section) ...
+        return demo
+
 if __name__ == "__main__":
-    main()
+    demo = main()
+    demo.queue().launch(ssr_mode=False)