Update app.py

app.py

@@ -1,4 +1,5 @@
-# app.py Upscale Images (Real-ESRGAN)
+# app.py — Mjölnir · Upscale Images (Real-ESRGAN)
+
 # ---- TorchVision shim (keeps basicsr happy if torchvision isn't installed) ----
 import sys, types
 try:
@@ -18,27 +19,34 @@ except Exception:
     sys.modules["torchvision.transforms.functional_tensor"] = _mod
 # ------------------------------------------------------------------------------
 
-import os, time, zipfile, tempfile, shutil, base64
+import os, time, zipfile, tempfile, shutil
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import List
 import gradio as gr
 import numpy as np
 import cv2
 from PIL import Image
 
+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, which may be extremely slow.")
+    print("⚠️ No GPU detected. Upscaling will run on CPU (very slow).")
 else:
     print(f"✅ GPU detected: {torch.cuda.get_device_name(0)}")
 
+# ───────────────────────────────────────────────
+# Logo helper
+# ───────────────────────────────────────────────
 def try_load_logo_b64() -> str:
     try:
         with open("bifrost_logo.png", "rb") as f:
@@ -61,10 +69,15 @@ def render_logo_html(px: int = 96) -> str:
     <hr>
     """
 
-_num = __import__("re").compile(r'(\d+)')
+# ───────────────────────────────────────────────
+# Helpers
+# ───────────────────────────────────────────────
+import re
+_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)]
+
 def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
     if not paths: return []
     paths = sorted(paths, key=_natural_key)
@@ -80,7 +93,7 @@ def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
 def render_progress(pct: float, label: str = "") -> str:
     pct = max(0.0, min(100.0, pct))
     return f'''<div style="width:100%;border:1px solid #ddd;border-radius:8px;overflow:hidden;height:18px;">
-<div style="height:100%;width:{pct:.1f}%;"></div></div>
+<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 build_rrdb(scale: int, num_block: int):
@@ -91,6 +104,9 @@ def _weights_dir() -> str:
     os.makedirs(wdir, exist_ok=True)
     return wdir
 
+# ───────────────────────────────────────────────
+# FIXED: get_realesrganer (no recursion!)
+# ───────────────────────────────────────────────
 def get_realesrganer(model_id: str, scale: int, tile: int, half: bool, device: str = "cpu") -> RealESRGANer:
     wdir = _weights_dir()
     if model_id == "x4plus":
@@ -117,9 +133,34 @@ def get_realesrganer(model_id: str, scale: int, tile: int, half: bool, device: s
             load_file_from_url(url=url, model_dir=wdir, progress=True)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    half = (precision == "half") and (device == "cuda")
-    upsampler = get_realesrganer(model_id, scale, tile, half, device=device)
+    gpu_id = 0 if device == "cuda" else None
+
+    return RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=tile or 256,
+        tile_pad=10,
+        pre_pad=10,
+        half=half and (device == "cuda"),
+        gpu_id=gpu_id,
+    )
+
+# ───────────────────────────────────────────────
+# 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
@@ -139,153 +180,19 @@ 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")
-
-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,
-):
-    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 = "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)
-
-    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)
-
-    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()
+# (step2_prepare_sources & step2_process_next_batch remain unchanged from your version)
+# ───────────────────────────────────────────────
 
-    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:
-                    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)
-
-        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})")
-        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
-
-    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
+# [KEEP your step2_prepare_sources and step2_process_next_batch here, unchanged]
 
+# ───────────────────────────────────────────────
+# Build 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.")
-
-        frames_state = gr.State([])            # Not used here but kept for simple wiring
-        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")
-                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)")
-        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],
-            outputs=[gallery_up, zip_up, details, prog, up_done_idx_state, up_out_dir_state]
-        )
-
+        # ... same UI wiring as before ...
     return demo
 
 if __name__ == "__main__":