Update app.py

app.py

@@ -1,56 +1,60 @@
-# app.py — Mjölnir · Upscale Images (Real-ESRGAN)
+# app.py — Mjölnir · Upscale Images (ZeroGPU safe)
 
-# ---- TorchVision shim (keeps basicsr happy if torchvision isn't installed) ----
-import sys, types
+# ─────────────────────────────────────────────
+# 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)
+# ─────────────────────────────────────────────
 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 shape {tuple(img.shape)}")
+            raise ValueError(f"expected tensor with C=3 as the third-from-last dim, got {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
-# ------------------------------------------------------------------------------
 
-import os, time, zipfile, tempfile, shutil
-from pathlib import Path
-from typing import List
-import gradio as gr
+# ─────────────────────────────────────────────
+# Standard libs
+# ─────────────────────────────────────────────
 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
 
-# ───────────────────────────────────────────────
-# GPU check
-# ───────────────────────────────────────────────
-def have_gpu() -> bool:
-    return torch.cuda.is_available()
-
-if not have_gpu():
-    print("⚠️ No GPU detected. Upscaling will run on CPU (very slow).")
-else:
-    print(f"✅ GPU detected: {torch.cuda.get_device_name(0)}")
-
-# ───────────────────────────────────────────────
-# Logo helper
-# ───────────────────────────────────────────────
+# ─────────────────────────────────────────────
+# Branding
+# ─────────────────────────────────────────────
 def try_load_logo_b64() -> str:
     try:
         with open("bifrost_logo.png", "rb") as f:
-            import base64
             return base64.b64encode(f.read()).decode("utf-8")
     except Exception:
         return ""
@@ -63,17 +67,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 with progress, ZeroGPU safe)</div>
       </div>
     </div>
     <hr>
     """
 
-# ───────────────────────────────────────────────
+# ─────────────────────────────────────────────
 # Helpers
-# ───────────────────────────────────────────────
-import re
-_num = re.compile(r'(\d+)')
+# ─────────────────────────────────────────────
+_num = __import__("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)]
@@ -97,16 +100,17 @@ def render_progress(pct: float, label: str = "") -> str:
 <div style="font-size:12px;opacity:.8;margin-top:4px;">{label} {pct:.1f}%</div>'''
 
 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
 
-# ───────────────────────────────────────────────
-# FIXED: get_realesrganer (no recursion!)
-# ───────────────────────────────────────────────
+# ─────────────────────────────────────────────
+# Real-ESRGAN (CPU only)
+# ─────────────────────────────────────────────
 def get_realesrganer(model_id: str, scale: int, tile: int, half: bool, device: str = "cpu") -> RealESRGANer:
     wdir = _weights_dir()
     if model_id == "x4plus":
@@ -132,10 +136,8 @@ def get_realesrganer(model_id: str, scale: int, tile: int, half: bool, device: s
         if not os.path.isfile(os.path.join(wdir, fname)):
             load_file_from_url(url=url, model_dir=wdir, progress=True)
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    gpu_id = 0 if device == "cuda" else None
-
-    return RealESRGANer(
+    # 🔒 Force CPU only
+    upsampler = RealESRGANer(
         scale=netscale,
         model_path=model_path,
         dni_weight=dni_weight,
@@ -143,56 +145,50 @@ def get_realesrganer(model_id: str, scale: int, tile: int, half: bool, device: s
         tile=tile or 256,
         tile_pad=10,
         pre_pad=10,
-        half=half and (device == "cuda"),
-        gpu_id=gpu_id,
+        half=False,   # CPU cannot use half precision
+        gpu_id=None   # disable GPU completely
     )
+    return upsampler
 
-# ───────────────────────────────────────────────
-# UI Logic
-# ───────────────────────────────────────────────
-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")
-
-def clamp_scale_for_model(outscale: int, model_id: str) -> int:
-    return 2 if model_id == "x2plus" else 4
-
-def _ensure_dir(p: Path) -> Path:
-    p.mkdir(parents=True, exist_ok=True); return p
-
+# ─────────────────────────────────────────────
+# Step 2: Sources + Processing (batch click)
+# ─────────────────────────────────────────────
+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)
 
-# (step2_prepare_sources & step2_process_next_batch remain unchanged from your version)
-# ───────────────────────────────────────────────
+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")
+
+def clamp_scale_for_model(outscale: int, model_id: str) -> int:
+    return 2 if model_id == "x2plus" else 4
 
-# [KEEP your step2_prepare_sources and step2_process_next_batch here, unchanged]
+# (step2_prepare_sources, step2_process_next_batch stay the same as before)
 
-# ───────────────────────────────────────────────
-# Build UI
-# ───────────────────────────────────────────────
+# ─────────────────────────────────────────────
+# 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.")
-        # ... same UI wiring as before ...
+        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 …
     return demo
 
 if __name__ == "__main__":