Update app.py

app.py

@@ -1,57 +1,48 @@
-# app.py — Mjölnir · Upscale Images (ZeroGPU safe)
+# app.py — Mjölnir · Upscale Images (Real-ESRGAN) — GPU/CPU safe
 
-# ─────────────────────────────────────────────
-# Force CPU only (ZeroGPU mode, no CUDA allowed)
-# ─────────────────────────────────────────────
-import os, sys, types, time, zipfile, tempfile, shutil, base64
-from pathlib import Path
-from typing import List
-
-os.environ["CUDA_VISIBLE_DEVICES"] = ""  # hide GPUs completely
-
-import torch
-def have_gpu() -> bool:
-    return torch.cuda.is_available()
-
-if not have_gpu():
-    print("⚠️ ZeroGPU mode: Running on CPU only (slow, but stable).")
-else:
-    print(f"✅ GPU detected: {torch.cuda.get_device_name(0)}")
-
-# ─────────────────────────────────────────────
-# TorchVision shim (keeps basicsr happy)
-# ─────────────────────────────────────────────
+# ---- TorchVision shim (keeps basicsr happy if torchvision isn't installed) ----
+import sys, types
 try:
     import torchvision.transforms.functional_tensor as _ft  # noqa: F401
 except Exception:
+    import torch
     _mod = types.ModuleType("torchvision.transforms.functional_tensor")
     def rgb_to_grayscale(img: "torch.Tensor", num_output_channels: int = 1) -> "torch.Tensor":
         if not torch.is_tensor(img):
             raise TypeError("rgb_to_grayscale expects a torch.Tensor")
         if img.ndim < 3 or img.shape[-3] != 3:
-            raise ValueError(f"expected tensor with C=3 as the third-from-last dim, got {tuple(img.shape)}")
+            raise ValueError(f"expected tensor with C=3 as the third-from-last dim, got shape {tuple(img.shape)}")
         r, g, b = img[..., -3, :, :], img[..., -2, :, :], img[..., -1, :, :]
         gray = 0.2989*r + 0.5870*g + 0.1140*b
         return torch.stack([gray, gray, gray], dim=-3) if num_output_channels == 3 else gray.unsqueeze(-3)
     _mod.rgb_to_grayscale = rgb_to_grayscale
     sys.modules["torchvision.transforms.functional_tensor"] = _mod
+# ------------------------------------------------------------------------------
 
-# ─────────────────────────────────────────────
-# Standard libs
-# ─────────────────────────────────────────────
+import os, time, zipfile, tempfile, shutil, base64
+from pathlib import Path
+from typing import List, Optional
+import re
+import gradio as gr
 import numpy as np
 import cv2
 from PIL import Image
-import gradio as gr
 
+import torch
 from basicsr.archs.rrdbnet_arch import RRDBNet as _RRDBNet
 from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
-# ─────────────────────────────────────────────
-# Branding
-# ─────────────────────────────────────────────
+# ────────────────────────────────────────────────────────
+# Small utils
+# ────────────────────────────────────────────────────────
+
+def have_gpu() -> bool:
+    return torch.cuda.is_available()
+
+print("✅ GPU available" if have_gpu() else "⚠️ No GPU detected. Will use CPU (slow).")
+
 def try_load_logo_b64() -> str:
     try:
         with open("bifrost_logo.png", "rb") as f:
@@ -67,16 +58,13 @@ def render_logo_html(px: int = 96) -> str:
       {img}
       <div>
         <div style="font-size:1.6rem;font-weight:800;">Mjölnir · Upscale Images</div>
-        <div style="opacity:0.8;">Real-ESRGAN (batch click with progress, ZeroGPU safe)</div>
+        <div style="opacity:0.8;">Real-ESRGAN (batch click with progress)</div>
       </div>
     </div>
     <hr>
     """
 
-# ─────────────────────────────────────────────
-# Helpers
-# ─────────────────────────────────────────────
-_num = __import__("re").compile(r'(\d+)')
+_num = re.compile(r'(\d+)')
 def _natural_key(p: Path | str):
     s = str(p)
     return [int(t) if t.isdigit() else t.lower() for t in _num.split(s)]
@@ -89,8 +77,9 @@ def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
     step = (total - 1) / (n - 1); idxs = [round(i * step) for i in range(n)]
     out, seen = [], set()
     for i in idxs:
+        i = int(i)
         if i not in seen:
-            out.append(str(paths[int(i)])); seen.add(int(i))
+            out.append(str(paths[i])); seen.add(i)
     return out
 
 def render_progress(pct: float, label: str = "") -> str:
@@ -99,12 +88,46 @@ def render_progress(pct: float, label: str = "") -> str:
 <div style="height:100%;width:{pct:.1f}%;background:#3b82f6;"></div></div>
 <div style="font-size:12px;opacity:.8;margin-top:4px;">{label} {pct:.1f}%</div>'''
 
+# ────────────────────────────────────────────────────────
+# Real-ESRGAN wiring (GPU/CPU safe)
+# ────────────────────────────────────────────────────────
+
 def build_rrdb(scale: int, num_block: int):
-    return RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
-                   num_block=num_block, num_grow_ch=32, scale=scale)
+    return _RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=scale)
+
+def _weights_dir() -> str:
+    wdir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights")
+    os.makedirs(wdir, exist_ok=True)
+    return wdir
 
-def get_realesrganer(model_id: str, scale: int, tile: int,
-                     half: bool, device: str = "cpu") -> RealESRGANer:
+def map_ui_model_to_internal(ui_name: str) -> str:
+    return {
+        "RealESRGAN_x4plus": "x4plus",
+        "RealESRGAN_x4plus_anime_6B": "x4plus-anime",
+        "RealESRGAN_x2plus": "x2plus",
+        "RealESRNet_x4plus": "x4plus",            # fallback to RRDB x4
+        "realesr-general-x4v3": "general-x4v3",   # SRVGG
+    }.get(ui_name, "x4plus")
+
+def clamp_scale_for_model(outscale: int, model_id: str) -> int:
+    if model_id == "x2plus":  # true 2x model
+        return 2
+    return 4  # rest are 4x
+
+def get_realesrganer(
+    model_id: str,
+    tile: int,
+    half: bool,
+    device: str,
+    denoise_strength: float = 0.5
+) -> RealESRGANer:
+    """
+    Returns a RealESRGANer for:
+      - x4plus (RRDB)
+      - x4plus-anime (RRDB)
+      - x2plus (RRDB)
+      - general-x4v3 (SRVGG, supports denoise_strength via DNI weights)
+    """
     wdir = _weights_dir()
 
     if model_id == "x4plus":
@@ -112,76 +135,259 @@ def get_realesrganer(model_id: str, scale: int, tile: int,
         netscale = 4
         url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
         model_path = os.path.join(wdir, "RealESRGAN_x4plus.pth")
+        if not os.path.isfile(model_path):
+            load_file_from_url(url=url, model_dir=wdir, progress=True)
+        model_path_arg = model_path
+        dni_weight = None
+
     elif model_id == "x4plus-anime":
         model = build_rrdb(scale=4, num_block=6)
         netscale = 4
         url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
         model_path = os.path.join(wdir, "RealESRGAN_x4plus_anime_6B.pth")
+        if not os.path.isfile(model_path):
+            load_file_from_url(url=url, model_dir=wdir, progress=True)
+        model_path_arg = model_path
+        dni_weight = None
+
     elif model_id == "x2plus":
         model = build_rrdb(scale=2, num_block=23)
         netscale = 2
         url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
         model_path = os.path.join(wdir, "RealESRGAN_x2plus.pth")
+        if not os.path.isfile(model_path):
+            load_file_from_url(url=url, model_dir=wdir, progress=True)
+        model_path_arg = model_path
+        dni_weight = None
+
+    elif model_id == "general-x4v3":
+        # SRVGG + two weights for DNI blend (denoise control)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        base_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/"
+        base_path = os.path.join(wdir, "realesr-general-x4v3.pth")
+        wdn_path  = os.path.join(wdir, "realesr-general-wdn-x4v3.pth")
+        if not os.path.isfile(base_path):
+            load_file_from_url(url=base_url + "realesr-general-x4v3.pth", model_dir=wdir, progress=True)
+        if not os.path.isfile(wdn_path):
+            load_file_from_url(url=base_url + "realesr-general-wdn-x4v3.pth", model_dir=wdir, progress=True)
+        model_path_arg = [base_path, wdn_path]
+        # blend base vs denoised
+        d = float(denoise_strength)
+        d = max(0.0, min(1.0, d))
+        dni_weight = [1.0 - d, d]
+
     else:
         raise ValueError(f"Unknown model_id: {model_id}")
 
-    if not os.path.isfile(model_path):
-        load_file_from_url(url=url, model_dir=wdir, progress=True)
-
-    # auto device
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    half = half and device == "cuda"
+    # Final device policy
+    device = device if device in ("cuda", "cpu") else ("cuda" if torch.cuda.is_available() else "cpu")
+    half = bool(half and device == "cuda")
 
     return RealESRGANer(
         scale=netscale,
-        model_path=model_path,
+        model_path=model_path_arg,
+        dni_weight=dni_weight,
         model=model,
-        tile=tile,
+        tile=int(tile or 256),
         tile_pad=10,
         pre_pad=0,
         half=half,
         device=device,
     )
 
-# ─────────────────────────────────────────────
-# Step 2: Sources + Processing (batch click)
-# ─────────────────────────────────────────────
-def _ensure_dir(p: Path) -> Path: p.mkdir(parents=True, exist_ok=True); return p
+# ────────────────────────────────────────────────────────
+# Batch upscaling helpers
+# ────────────────────────────────────────────────────────
+
+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 map_ui_model_to_internal(ui_name: str) -> str:
-    return {
-        "RealESRGAN_x4plus": "x4plus",
-        "RealESRGAN_x4plus_anime_6B": "x4plus-anime",
-        "RealESRGAN_x2plus": "x2plus",
-        "RealESRNet_x4plus": "x4plus",
-        "realesr-general-x4v3": "x4plus",
-    }.get(ui_name, "x4plus")
+# ────────────────────────────────────────────────────────
+# Step 2 · Prepare + Process (generator for streaming)
+# ────────────────────────────────────────────────────────
 
-def clamp_scale_for_model(outscale: int, model_id: str) -> int:
-    return 2 if model_id == "x2plus" else 4
+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 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
+    return src, str(out_dir), done_idx, total, f"Sources loaded: {total} image(s). Click 'Process Next Batch'.", render_progress(0.0, "Ready")
+
+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,
+    force_cpu
+):
+    if not up_src_paths or not up_out_dir:
+        yield None, None, "Load sources first.", render_progress(0.0, "Idle"), up_done_idx, up_out_dir
+        return
+
+    model_id = map_ui_model_to_internal(ui_model_name)
+    scale = clamp_scale_for_model(int(outscale or 4), model_id)
+
+    # Device policy
+    device = "cpu" if force_cpu else ("cuda" if torch.cuda.is_available() else "cpu")
+    # Precision policy
+    if precision == "half":
+        use_half = (device == "cuda")
+    elif precision == "full":
+        use_half = False
+    else:  # auto
+        use_half = (device == "cuda")
+
+    tile = int(tile or 256)
+    batch_size = max(1, int(batch_size or 8))
+
+    # Build upsampler (handles general-x4v3 as well)
+    upsampler = get_realesrganer(
+        model_id=model_id,
+        tile=tile,
+        half=use_half,
+        device=device,
+        denoise_strength=float(denoise_strength or 0.5)
+    )
+
+    # 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)
+
+    if start >= up_total:
+        gallery = _build_gallery_from_dir(out_dir, 30)
+        zip_file = _save_zip_of_dir(out_dir, Path(out_dir.parent) / "upscaled.zip")
+        yield gallery, zip_file, "All images processed.", render_progress(100.0, "Done"), start, up_out_dir
+        return
+
+    batch_paths = up_src_paths[start:end]
+    total_in_batch = len(batch_paths)
+    t0 = time.time()
+
+    for idx, fp in enumerate(batch_paths, start=1):
+        try:
+            with Image.open(fp) as im:
+                img = im.convert("RGB")
+                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:
+                    # For general-x4v3, denoise_strength is already in DNI weights
+                    output, _ = upsampler.enhance(cv_img, outscale=scale)
+                Image.fromarray(output).save(out_dir / (Path(fp).stem + ".jpg"), quality=95)
+        except Exception as e:
+            print("Upscale error:", e)
+
+        elapsed = time.time() - t0
+        pct_batch = (idx / total_in_batch) * 100.0
+        eta = (total_in_batch - idx) * (elapsed / max(1, idx))
+        label = (f"Batch: {idx}/{total_in_batch} · ~{eta:.1f}s ETA · "
+                 f"global {start+idx}/{up_total} (x{scale}, model={ui_model_name}, device={device}, half={use_half})")
+        gallery = _build_gallery_from_dir(out_dir, 30)
+        zip_file = _save_zip_of_dir(out_dir, Path(out_dir.parent) / "upscaled.zip")
+        yield gallery, zip_file, label, render_progress(pct_batch, f"Upscaling {pct_batch:.0f}% (batch)"), start+idx, up_out_dir
 
-# (step2_prepare_sources, step2_process_next_batch stay the same as before)
+    next_idx = end
+    pct_global = (next_idx / up_total) * 100.0 if up_total else 100.0
+    gallery = _build_gallery_from_dir(out_dir, 30)
+    zip_file = _save_zip_of_dir(out_dir, Path(out_dir.parent) / "upscaled.zip")
+    yield gallery, zip_file, f"Processed batch of {total_in_batch}. {next_idx}/{up_total} done.", render_progress(pct_global, "Upscaling… (global)"), next_idx, up_out_dir
 
-# ─────────────────────────────────────────────
+# ────────────────────────────────────────────────────────
 # UI
-# ─────────────────────────────────────────────
+# ────────────────────────────────────────────────────────
+
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.HTML(render_logo_html(88))
-        gr.Markdown("Upload images and upscale with Real-ESRGAN. Runs in CPU-only ZeroGPU mode (slow).")
-        # … keep the rest of your batch-click UI wiring unchanged …
+        gr.Markdown("Upload images and upscale with Real-ESRGAN. Process in batches with live progress. GPU if available, CPU fallback.")
+
+        # States
+        frames_state = gr.State([])            # kept for parity
+        up_src_paths_state = gr.State([])
+        up_out_dir_state   = gr.State("")
+        up_done_idx_state  = gr.State(0)
+        up_total_state     = gr.State(0)
+
+        imgs_override = gr.Files(label="Upload images (JPG/PNG)", file_types=[".jpg",".jpeg",".png"], type="filepath")
+
+        with gr.Accordion("Upscaling options", open=True):
+            with gr.Row():
+                ui_model_name = gr.Dropdown(
+                    label="Upscaler model",
+                    choices=["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B", "RealESRGAN_x2plus", "realesr-general-x4v3"],
+                    value="RealESRGAN_x4plus"
+                )
+                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 (model-limited)")
+                face_enhance = gr.Checkbox(value=False, label="Face Enhancement (GFPGAN)")
+        with gr.Row():
+            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)")
+            force_cpu = gr.Checkbox(value=False, label="Zero-GPU Mode (force CPU)")
+        with gr.Row():
+            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_prepare = gr.Button("Load / Reset Sources", variant="secondary")
+            btn_next    = gr.Button("Process Next Batch", variant="primary")
+
+        prog = 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")
+        details = gr.Markdown("")
+
+        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, details, prog]
+        )
+
+        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, force_cpu],
+            outputs=[gallery_up, zip_up, details, prog, up_done_idx_state, up_out_dir_state]
+        )
+
     return demo
 
 if __name__ == "__main__":