""" Upscale service — auto-detects best available method and runs it. Auto-installs Real-ESRGAN NCNN Vulkan binary on first use if no AI upscaler found. Priority (auto mode): 1. Real-ESRGAN PyTorch + CUDA GPU — fastest, best quality 2. Real-ESRGAN PyTorch + Apple MPS — fast on Apple Silicon 3. Real-ESRGAN NCNN Vulkan binary — fast on any GPU (Intel/AMD/integrated) 4. Real-ESRGAN PyTorch CPU — works, slow (warn user) 5. Lanczos — always available, instant Capability probe is run once at first call and cached. NCNN binary is auto-downloaded if no AI upscaler is found. """ import asyncio import os import platform import shutil import stat import subprocess import sys import tempfile import urllib.request import zipfile from dataclasses import dataclass, field from enum import Enum from io import BytesIO from pathlib import Path from typing import Optional from PIL import Image # ── NCNN auto-install ───────────────────────────────────────────────────────── NCNN_DEST_DIR = Path("/app/data/models/realesrgan") NCNN_VERSION = "v0.2.5.0" NCNN_BASE_URL = f"https://github.com/xinntao/Real-ESRGAN/releases/download/{NCNN_VERSION}" _PLATFORM_ZIP = { "linux": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-ubuntu.zip", "darwin": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-macos.zip", "win32": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-windows.zip", "windows": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-windows.zip", } class InstallState(str, Enum): idle = "idle" downloading = "downloading" extracting = "extracting" done = "done" failed = "failed" @dataclass class InstallStatus: state: InstallState = InstallState.idle progress: int = 0 # 0-100 message: str = "" error: str = "" _install_status = InstallStatus() _install_lock = asyncio.Lock() def get_install_status() -> dict: s = _install_status return { "state": s.state.value, "progress": s.progress, "message": s.message, "error": s.error, } def _ncnn_binary_name() -> str: plat = sys.platform.lower() return "realesrgan-ncnn-vulkan.exe" if "win" in plat else "realesrgan-ncnn-vulkan" async def ensure_ncnn_installed() -> Optional[Path]: """ Check if NCNN binary is present; if not, download and install it. Returns the binary Path on success, None on failure. Serialised via _install_lock so concurrent callers wait for a single install. """ global _install_status binary_path = NCNN_DEST_DIR / _ncnn_binary_name() if binary_path.exists() and os.access(binary_path, os.X_OK): _install_status = InstallStatus(state=InstallState.done, progress=100, message="Already installed.") return binary_path async with _install_lock: # Re-check after acquiring lock (another coroutine may have just finished) if binary_path.exists() and os.access(binary_path, os.X_OK): _install_status = InstallStatus(state=InstallState.done, progress=100, message="Already installed.") return binary_path if _install_status.state == InstallState.downloading: return None # install already in progress plat = sys.platform.lower() zip_name = _PLATFORM_ZIP.get(plat) if not zip_name: _install_status = InstallStatus( state=InstallState.failed, error=f"Unsupported platform: {plat}", ) return None url = f"{NCNN_BASE_URL}/{zip_name}" try: NCNN_DEST_DIR.mkdir(parents=True, exist_ok=True) zip_path = NCNN_DEST_DIR / zip_name # Download _install_status = InstallStatus( state=InstallState.downloading, progress=0, message=f"Downloading Real-ESRGAN NCNN {NCNN_VERSION}…", ) def _do_download(): def _progress(count, block, total): if total > 0: pct = min(90, int(count * block * 90 / total)) _install_status.progress = pct urllib.request.urlretrieve(url, zip_path, _progress) loop = asyncio.get_event_loop() await loop.run_in_executor(None, _do_download) # Extract _install_status.state = InstallState.extracting _install_status.progress = 92 _install_status.message = "Extracting…" def _do_extract(): with zipfile.ZipFile(zip_path, "r") as zf: zf.extractall(NCNN_DEST_DIR) # Find binary (may be in a subdirectory) found = list(NCNN_DEST_DIR.rglob(_ncnn_binary_name())) if not found: raise FileNotFoundError(f"Binary not found after extract: {_ncnn_binary_name()}") extracted = found[0] if extracted != binary_path: extracted.rename(binary_path) # Make executable if "win" not in sys.platform.lower(): binary_path.chmod( binary_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH ) zip_path.unlink(missing_ok=True) await loop.run_in_executor(None, _do_extract) _install_status = InstallStatus( state=InstallState.done, progress=100, message=f"Installed: {binary_path}", ) # Bust caps cache so probe picks up new binary invalidate_caps_cache() return binary_path except Exception as exc: _install_status = InstallStatus( state=InstallState.failed, error=str(exc), message="Installation failed.", ) print(f"[upscale] NCNN auto-install failed: {exc}") return None # ── Capability detection ────────────────────────────────────────────────────── _caps: Optional[dict] = None def probe_upscale_capabilities() -> dict: """ Detect what upscaling hardware and software is available. Result is cached after first call. """ global _caps if _caps is not None: return _caps caps = { "lanczos": True, "realesrgan_pytorch": False, "realesrgan_pytorch_device": None, "realesrgan_ncnn": False, "realesrgan_ncnn_path": None, "recommended": "lanczos", "recommended_label": "Lanczos (no AI upscaler found)", "methods": ["lanczos"], "ncnn_install_status": get_install_status(), } # ── PyTorch path ────────────────────────────────────────────────────────── pytorch_device = None try: import torch if torch.cuda.is_available(): pytorch_device = "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): pytorch_device = "mps" else: pytorch_device = "cpu" except ImportError: pass if pytorch_device: try: import realesrgan # noqa: F401 from basicsr.archs.rrdbnet_arch import RRDBNet # noqa: F401 caps["realesrgan_pytorch"] = True caps["realesrgan_pytorch_device"] = pytorch_device caps["methods"].append("realesrgan_pytorch") except ImportError: pass # ── NCNN Vulkan binary ──────────────────────────────────────────────────── ncnn_path = _find_ncnn_binary() if ncnn_path: caps["realesrgan_ncnn"] = True caps["realesrgan_ncnn_path"] = str(ncnn_path) caps["methods"].append("realesrgan_ncnn") # ── Pick recommended ────────────────────────────────────────────────────── if caps["realesrgan_pytorch"] and pytorch_device in ("cuda", "mps"): device_label = "CUDA GPU" if pytorch_device == "cuda" else "Apple Silicon" caps["recommended"] = "realesrgan_pytorch" caps["recommended_label"] = f"Real-ESRGAN ({device_label})" elif caps["realesrgan_ncnn"]: caps["recommended"] = "realesrgan_ncnn" caps["recommended_label"] = "Real-ESRGAN NCNN (Vulkan)" elif caps["realesrgan_pytorch"] and pytorch_device == "cpu": caps["recommended"] = "realesrgan_pytorch" caps["recommended_label"] = "Real-ESRGAN (CPU — may be slow)" else: caps["recommended"] = "lanczos" caps["recommended_label"] = "Lanczos (installing Real-ESRGAN…)" _caps = caps return caps def _find_ncnn_binary() -> Optional[Path]: """Find realesrgan-ncnn-vulkan binary on the system.""" found = shutil.which("realesrgan-ncnn-vulkan") if found: return Path(found) candidates = [ NCNN_DEST_DIR / _ncnn_binary_name(), Path("/usr/local/bin/realesrgan-ncnn-vulkan"), Path.home() / ".local/bin/realesrgan-ncnn-vulkan", Path(r"C:/realesrgan-ncnn-vulkan/realesrgan-ncnn-vulkan.exe"), Path("/opt/homebrew/bin/realesrgan-ncnn-vulkan"), ] for p in candidates: if p.exists() and os.access(p, os.X_OK): return p return None def invalidate_caps_cache(): """Call after installing new software so next probe picks it up.""" global _caps _caps = None # ── Upscale implementations ─────────────────────────────────────────────────── def _to_png_bytes(img: Image.Image) -> bytes: buf = BytesIO() img.save(buf, format="PNG") return buf.getvalue() def upscale_lanczos(image: Image.Image, scale: float) -> tuple[bytes, str]: """Pure Pillow Lanczos — instant, always available.""" new_w = round(image.width * scale) new_h = round(image.height * scale) result = image.resize((new_w, new_h), Image.Resampling.LANCZOS) return _to_png_bytes(result), "lanczos" def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, str]: """ Real-ESRGAN via PyTorch. Uses CUDA > MPS > CPU automatically based on what's available. """ import torch from basicsr.archs.rrdbnet_arch import RRDBNet from realesrgan import RealESRGANer caps = probe_upscale_capabilities() device = caps.get("realesrgan_pytorch_device", "cpu") model_scale = 2 if scale <= 2.5 else 4 model = RRDBNet( num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=model_scale ) model_dir = Path("/app/data/models/realesrgan") model_dir.mkdir(parents=True, exist_ok=True) model_name = f"RealESRGAN_x{model_scale}plus.pth" model_path = model_dir / model_name if not model_path.exists(): model_path = None upsampler = RealESRGANer( scale=model_scale, model_path=str(model_path) if model_path else None, model=model, tile=512, tile_pad=10, pre_pad=0, half=(device == "cuda"), device=torch.device(device), ) import numpy as np img_bgr = np.array(image)[:, :, ::-1].copy() enhanced, _ = upsampler.enhance(img_bgr, outscale=scale) result = Image.fromarray(enhanced[:, :, ::-1]) label = f"realesrgan_pytorch_{device}" return _to_png_bytes(result), label def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, str]: """ Real-ESRGAN via NCNN Vulkan binary — works on any GPU. Runs as subprocess with temp file I/O. """ caps = probe_upscale_capabilities() binary = caps.get("realesrgan_ncnn_path") if not binary: raise RuntimeError("realesrgan-ncnn-vulkan binary not found") model_scale = 4 if scale > 2.5 else 2 target_w = round(image.width * scale) target_h = round(image.height * scale) with tempfile.TemporaryDirectory() as tmpdir: in_path = Path(tmpdir) / "input.png" out_path = Path(tmpdir) / "output.png" image.save(in_path, format="PNG") model_name = f"realesrgan-x{model_scale}plus" cmd = [ binary, "-i", str(in_path), "-o", str(out_path), "-s", str(model_scale), "-n", model_name, "-f", "png", ] result_proc = subprocess.run(cmd, capture_output=True, timeout=300) if result_proc.returncode != 0: raise RuntimeError( f"realesrgan-ncnn-vulkan failed: {result_proc.stderr.decode()}" ) result = Image.open(out_path).convert("RGB") if result.width != target_w or result.height != target_h: result = result.resize((target_w, target_h), Image.Resampling.LANCZOS) return _to_png_bytes(result), "realesrgan_ncnn" # ── Public entry point ──────────────────────────────────────────────────────── def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]: """Upscale image synchronously. Returns (png_bytes, method_used_label).""" caps = probe_upscale_capabilities() if method == "auto": method = caps["recommended"] if method == "realesrgan_pytorch": if caps["realesrgan_pytorch"]: try: return upscale_realesrgan_pytorch(image, scale) except Exception as e: print(f"Real-ESRGAN PyTorch failed, falling back: {e}") if caps["realesrgan_ncnn"]: try: return upscale_realesrgan_ncnn(image, scale) except Exception as e: print(f"Real-ESRGAN NCNN fallback failed: {e}") return upscale_lanczos(image, scale) if method == "realesrgan_ncnn": if caps["realesrgan_ncnn"]: try: return upscale_realesrgan_ncnn(image, scale) except Exception as e: print(f"Real-ESRGAN NCNN failed, falling back: {e}") if caps["realesrgan_pytorch"]: try: return upscale_realesrgan_pytorch(image, scale) except Exception as e: print(f"Real-ESRGAN PyTorch fallback failed: {e}") return upscale_lanczos(image, scale) return upscale_lanczos(image, scale) async def upscale_image(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]: """Async wrapper — runs upscale in thread pool to avoid blocking the event loop.""" loop = asyncio.get_event_loop() return await loop.run_in_executor(None, upscale_sync, image, scale, method)