Update app.py

app.py

@@ -1,6 +1,6 @@
-# app.py — Mjölnir · Upscale Images (Real-ESRGAN) — GPU/CPU safe
+# app.py — Mjölnir · Upscale Images (ZeroGPU-ready, batch click, with logo)
 
-# ---- TorchVision shim (keeps basicsr happy if torchvision isn't installed) ----
+# ---- TorchVision shim (so basicsr can import even without torchvision) ----
 import sys, types
 try:
     import torchvision.transforms.functional_tensor as _ft  # noqa: F401
@@ -17,32 +17,29 @@ except Exception:
         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
-# ------------------------------------------------------------------------------
+# --------------------------------------------------------------------------
 
 import os, time, zipfile, tempfile, shutil, base64
 from pathlib import Path
-from typing import List, Optional
-import re
+from typing import List, Optional, Tuple
+
 import gradio as gr
 import numpy as np
 import cv2
 from PIL import Image
 
-import torch
+# ZeroGPU hook
+import spaces
+
+# Real-ESRGAN / basicsr (importing these at top is OK; just don't touch CUDA here)
 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
 
 # ────────────────────────────────────────────────────────
-# Small utils
+# Branding (logo)
 # ────────────────────────────────────────────────────────
-
-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:
@@ -58,12 +55,16 @@ 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)</div>
+        <div style="opacity:0.8;">Real-ESRGAN · batch click · live progress · ZeroGPU</div>
       </div>
     </div>
     <hr>
     """
 
+# ────────────────────────────────────────────────────────
+# Small helpers
+# ────────────────────────────────────────────────────────
+import re
 _num = re.compile(r'(\d+)')
 def _natural_key(p: Path | str):
     s = str(p)
@@ -77,9 +78,8 @@ 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[i])); seen.add(i)
+        if int(i) not in seen:
+            out.append(str(paths[int(i)])); seen.add(int(i))
     return out
 
 def render_progress(pct: float, label: str = "") -> str:
@@ -88,10 +88,27 @@ 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>'''
 
+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)
+
 # ────────────────────────────────────────────────────────
-# Real-ESRGAN wiring (GPU/CPU safe)
+# Models & weight management (CPU-safe; no CUDA used here)
 # ────────────────────────────────────────────────────────
-
 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)
 
@@ -100,128 +117,63 @@ def _weights_dir() -> str:
     os.makedirs(wdir, exist_ok=True)
     return wdir
 
-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:
+def ensure_weights(model_id: str) -> Tuple[object, int, str, Optional[list]]:
     """
-    Returns a RealESRGANer for:
-      - x4plus (RRDB)
-      - x4plus-anime (RRDB)
-      - x2plus (RRDB)
-      - general-x4v3 (SRVGG, supports denoise_strength via DNI weights)
+    Returns: (model, netscale, model_path_or_list, dni_weight_placeholder)
     """
     wdir = _weights_dir()
-
     if model_id == "x4plus":
-        model = build_rrdb(scale=4, num_block=23)
-        netscale = 4
+        model = build_rrdb(scale=4, num_block=23); 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
+        return model, netscale, model_path, None
 
-    elif model_id == "x4plus-anime":
-        model = build_rrdb(scale=4, num_block=6)
-        netscale = 4
+    if 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
+        return model, netscale, model_path, None
 
-    elif model_id == "x2plus":
-        model = build_rrdb(scale=2, num_block=23)
-        netscale = 2
+    if 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
+        return model, netscale, model_path, None
 
-    elif model_id == "general-x4v3":
-        # SRVGG + two weights for DNI blend (denoise control)
+    # For UI compatibility only: map general-x4v3 to x4plus backend (or implement SRVGG if you prefer)
+    if model_id == "general-x4v3":
+        # Proper SRVGG model:
         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}")
-
-    # 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_arg,
-        dni_weight=dni_weight,
-        model=model,
-        tile=int(tile or 256),
-        tile_pad=10,
-        pre_pad=0,
-        half=half,
-        device=device,
-    )
+        base_pth = os.path.join(wdir, "realesr-general-x4v3.pth")
+        if not os.path.isfile(base_pth):
+            load_file_from_url(url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
+                               model_dir=wdir, progress=True)
+        return model, netscale, base_pth, None
 
-# ────────────────────────────────────────────────────────
-# Batch upscaling helpers
-# ────────────────────────────────────────────────────────
+    raise ValueError(f"Unknown model_id: {model_id}")
 
-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 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
+        "realesr-general-x4v3": "general-x4v3", # SRVGG
+    }.get(ui_name, "x4plus")
 
-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 clamp_scale_for_model(outscale: int, model_id: str) -> int:
+    return 2 if model_id == "x2plus" else 4
 
 # ────────────────────────────────────────────────────────
-# Step 2 · Prepare + Process (generator for streaming)
+# Step 2 · Prepare sources (CPU)
 # ────────────────────────────────────────────────────────
-
 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:
@@ -237,56 +189,58 @@ def step2_prepare_sources(frames_list, uploaded_imgs, max_images):
     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")
 
+# ────────────────────────────────────────────────────────
+# Step 2 · Process next batch (GPU) — ZeroGPU entry point
+# ────────────────────────────────────────────────────────
+@spaces.GPU
 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
 ):
+    """
+    Runs on ZeroGPU. Heavy parts (model load + enhance) are done inside this function.
+    Yields progress after each image in the current batch.
+    """
+    # Validate inputs
     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
 
+    # Resolve model & scale
     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))
+    use_half = (precision == "half")  # we'll honor this on CUDA only
 
-    # Build upsampler (handles general-x4v3 as well)
-    upsampler = get_realesrganer(
-        model_id=model_id,
+    # Ensure weights & build model (still CPU-safe) then instantiate ESRGANer on GPU
+    model, netscale, model_path, dni_weight = ensure_weights(model_id)
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
         tile=tile,
-        half=use_half,
-        device=device,
-        denoise_strength=float(denoise_strength or 0.5)
+        tile_pad=10,
+        pre_pad=10,
+        half=use_half,     # when ZeroGPU gives CUDA, this enables fp16
+        gpu_id=0           # request the single available GPU
     )
 
-    # Optional: GFPGAN face enhancer
+    # Optional face enhancer (kept off by default as it adds weight download)
     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
+                upscale=scale, arch="clean", channel_multiplier=2, bg_upsampler=upsampler
             )
         except Exception as e:
             print("GFPGAN load failed:", e)
             face_enhancer = None
 
+    # Batch window
     start = int(up_done_idx or 0)
     end = min(start + batch_size, int(up_total or 0))
     out_dir = Path(up_out_dir)
@@ -307,23 +261,27 @@ def step2_process_next_batch(
                 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)
+                    _, _, 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)
+                    # denoise_strength has effect with general-x4v3; harmless otherwise
+                    output, _ = upsampler.enhance(cv_img, outscale=scale, denoise_strength=float(denoise_strength or 0.5))
                 Image.fromarray(output).save(out_dir / (Path(fp).stem + ".jpg"), quality=95)
         except Exception as e:
             print("Upscale error:", e)
 
+        # Progress for THIS batch
         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})")
+                 f"global {start+idx}/{up_total} (x{scale}, model={ui_model_name}, tile={tile}, 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
 
+    # Batch complete
     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)
@@ -333,14 +291,13 @@ def step2_process_next_batch(
 # ────────────────────────────────────────────────────────
 # 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. Process in batches with live progress. GPU if available, CPU fallback.")
+        gr.Markdown("Upload images and upscale with Real-ESRGAN. **ZeroGPU** spins up a GPU **only** during batch processing.")
 
         # States
-        frames_state = gr.State([])            # kept for parity
+        frames_state = gr.State([])            # present for parity; not used here
         up_src_paths_state = gr.State([])
         up_out_dir_state   = gr.State("")
         up_done_idx_state  = gr.State(0)
@@ -356,38 +313,44 @@ def build_ui():
                     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)")
+                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=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)")
+            precision = gr.Dropdown(["auto", "half", "full"], value="half", label="Precision (GPU: half, CPU ignored)")
         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")
+            btn_next    = gr.Button("Process Next Batch (uses GPU)", variant="primary")
 
         prog = gr.HTML(render_progress(0.0, "Idle"))
-        gallery_up = gr.Gallery(label="Upscaled preview (30 sampled)", columns=6, height=480)
+        gallery_up = gr.Gallery(label="Upscaled preview (sampled 30)", columns=6, height=480)
         zip_up = gr.File(label="Download upscaled ZIP")
         details = gr.Markdown("")
 
+        # 1) load/reset sources (CPU)
         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]
         )
 
+        # 2) process one batch per click (GPU)
         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],
+                    ui_model_name, outscale, tile, precision, denoise_strength, face_enhance, batch_size],
             outputs=[gallery_up, zip_up, details, prog, up_done_idx_state, up_out_dir_state]
         )
 
+        gr.Markdown(
+            "> ℹ️ **ZeroGPU tips**: Larger tiles are faster but use more VRAM. If you hit OOM, try `tile=128`, "
+            "`batch size=4–8`, and keep `Precision=half`."
+        )
+
     return demo
 
 if __name__ == "__main__":