Create app.py

app.py

@@ -0,0 +1,285 @@
+# app.py Upscale Images (Real-ESRGAN)
+# ---- 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 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
+# ------------------------------------------------------------------------------
+
+import os, time, zipfile, tempfile, shutil, base64
+from pathlib import Path
+from typing import List, Optional, Tuple
+import gradio as gr
+import numpy as np
+import cv2
+from PIL import Image
+
+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
+
+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 ""
+LOGO_B64 = try_load_logo_b64()
+
+def render_logo_html(px: int = 96) -> str:
+    img = f'<img src="data:image/png;base64,{LOGO_B64}" style="height:{px}px;width:auto;" />' if LOGO_B64 else ""
+    return f"""
+    <div style="display:flex;align-items:center;gap:16px;">
+      {img}
+      <div>
+        <div style="font-size:1.6rem;font-weight:800;">Bifröst · Upscale Images</div>
+        <div style="opacity:0.8;">Real-ESRGAN (batch click with progress)</div>
+      </div>
+    </div>
+    <hr>
+    """
+
+_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)]
+def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
+    if not paths: return []
+    paths = sorted(paths, key=_natural_key)
+    total = len(paths); n = max(1, min(n, total))
+    if n == total: return [str(p) for p in paths]
+    step = (total - 1) / (n - 1); idxs = [round(i * step) for i in range(n)]
+    out, seen = [], set()
+    for i in idxs:
+        if i not in seen:
+            out.append(str(paths[int(i)])); seen.add(int(i))
+    return out
+
+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="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)
+
+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:
+    wdir = _weights_dir()
+    if model_id == "x4plus":
+        model = build_rrdb(scale=4, num_block=23); netscale = 4
+        urls = ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"]
+        model_path = os.path.join(wdir, "RealESRGAN_x4plus.pth")
+        dni_weight = None
+    elif model_id == "x4plus-anime":
+        model = build_rrdb(scale=4, num_block=6); netscale = 4
+        urls = ["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")
+        dni_weight = None
+    elif model_id == "x2plus":
+        model = build_rrdb(scale=2, num_block=23); netscale = 2
+        urls = ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"]
+        model_path = os.path.join(wdir, "RealESRGAN_x2plus.pth")
+        dni_weight = None
+    else:
+        raise ValueError(f"Unknown model_id: {model_id}")
+
+    for url in urls:
+        fname = os.path.basename(url)
+        if not os.path.isfile(os.path.join(wdir, fname)):
+            load_file_from_url(url=url, model_dir=wdir, progress=True)
+
+    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=bool(half and device == "cuda"), gpu_id=gpu_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 _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()
+
+    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
+
+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]
+        )
+
+    return demo
+
+if __name__ == "__main__":
+    build_ui().queue().launch()