diff --git a/backend/app/main.py b/backend/app/main.py index 0d02fb8..70d9b5a 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -4,6 +4,7 @@ from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, HTMLResponse from contextlib import asynccontextmanager from pathlib import Path +import asyncio import os from app.config import settings @@ -13,8 +14,13 @@ from app.routers import projects, edits, images, patches, generate, tools, ai_to @asynccontextmanager async def lifespan(app: FastAPI): - """Initialize database on startup""" + """Initialize database on startup; auto-install Real-ESRGAN NCNN in background.""" init_db() + # Kick off NCNN install in background if no AI upscaler detected + from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed + caps = probe_upscale_capabilities() + if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]: + asyncio.create_task(ensure_ncnn_installed()) yield diff --git a/backend/app/routers/print_tools.py b/backend/app/routers/print_tools.py index 2d5b1ec..52fe476 100644 --- a/backend/app/routers/print_tools.py +++ b/backend/app/routers/print_tools.py @@ -303,17 +303,38 @@ def upscale_refresh_caps(): @router.get("/upscale/available") -def upscale_available(): +async def upscale_available(): """ Return capability probe: which upscale methods are available, which device will be used, and which method is recommended. + If no AI upscaler is found, triggers background NCNN auto-install. Frontend uses this to populate the method selector. """ - from app.services.upscale import probe_upscale_capabilities + from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed, get_install_status caps = probe_upscale_capabilities() + # Auto-install NCNN if no AI upscaler is available yet + if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]: + asyncio.create_task(ensure_ncnn_installed()) + caps["ncnn_install_status"] = get_install_status() return caps +@router.get("/upscale/install-status") +def upscale_install_status(): + """Poll for Real-ESRGAN NCNN auto-install progress.""" + from app.services.upscale import get_install_status, probe_upscale_capabilities, _find_ncnn_binary + status = get_install_status() + # If install just finished, refresh caps + if status["state"] == "done": + from app.services.upscale import invalidate_caps_cache + invalidate_caps_cache() + caps = probe_upscale_capabilities() + status["ncnn_available"] = caps["realesrgan_ncnn"] + else: + status["ncnn_available"] = False + return status + + @router.post("/upscale") async def upscale(req: UpscaleRequest): """ diff --git a/backend/app/services/upscale.py b/backend/app/services/upscale.py index 9e8bbba..92d43b4 100644 --- a/backend/app/services/upscale.py +++ b/backend/app/services/upscale.py @@ -1,5 +1,6 @@ """ 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 @@ -9,20 +10,173 @@ Priority (auto mode): 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 @@ -40,12 +194,13 @@ def probe_upscale_capabilities() -> dict: caps = { "lanczos": True, "realesrgan_pytorch": False, - "realesrgan_pytorch_device": None, # "cuda" | "mps" | "cpu" + "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 ────────────────────────────────────────────────────────── @@ -91,7 +246,7 @@ def probe_upscale_capabilities() -> dict: caps["recommended_label"] = "Real-ESRGAN (CPU — may be slow)" else: caps["recommended"] = "lanczos" - caps["recommended_label"] = "Lanczos (install Real-ESRGAN for AI quality)" + caps["recommended_label"] = "Lanczos (installing Real-ESRGAN…)" _caps = caps return caps @@ -99,26 +254,20 @@ def probe_upscale_capabilities() -> dict: def _find_ncnn_binary() -> Optional[Path]: """Find realesrgan-ncnn-vulkan binary on the system.""" - # Check PATH first found = shutil.which("realesrgan-ncnn-vulkan") if found: return Path(found) - # Check known install locations candidates = [ - Path("/app/data/models/realesrgan/realesrgan-ncnn-vulkan"), + NCNN_DEST_DIR / _ncnn_binary_name(), Path("/usr/local/bin/realesrgan-ncnn-vulkan"), Path.home() / ".local/bin/realesrgan-ncnn-vulkan", - # Windows Path(r"C:/realesrgan-ncnn-vulkan/realesrgan-ncnn-vulkan.exe"), - # macOS Homebrew Path("/opt/homebrew/bin/realesrgan-ncnn-vulkan"), - Path("/usr/local/bin/realesrgan-ncnn-vulkan"), ] for p in candidates: if p.exists() and os.access(p, os.X_OK): return p - return None @@ -148,7 +297,6 @@ def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, """ Real-ESRGAN via PyTorch. Uses CUDA > MPS > CPU automatically based on what's available. - Scale factors: any float — upscales to nearest 2x or 4x model, then resizes to exact target. """ import torch from basicsr.archs.rrdbnet_arch import RRDBNet @@ -157,20 +305,18 @@ def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, caps = probe_upscale_capabilities() device = caps.get("realesrgan_pytorch_device", "cpu") - # Choose model: x2 for scale <= 2.5, x4 otherwise 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 path: check local cache first, then let RealESRGANer auto-download 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 # RealESRGANer will download to its default cache + model_path = None upsampler = RealESRGANer( scale=model_scale, @@ -179,14 +325,14 @@ def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, tile=512, tile_pad=10, pre_pad=0, - half=(device == "cuda"), # fp16 only on CUDA + half=(device == "cuda"), device=torch.device(device), ) import numpy as np - img_bgr = np.array(image)[:, :, ::-1].copy() # RGB→BGR + img_bgr = np.array(image)[:, :, ::-1].copy() enhanced, _ = upsampler.enhance(img_bgr, outscale=scale) - result = Image.fromarray(enhanced[:, :, ::-1]) # BGR→RGB + result = Image.fromarray(enhanced[:, :, ::-1]) label = f"realesrgan_pytorch_{device}" return _to_png_bytes(result), label @@ -194,7 +340,7 @@ def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, str]: """ - Real-ESRGAN via NCNN Vulkan binary — works on any GPU (Intel/AMD/integrated/Apple). + Real-ESRGAN via NCNN Vulkan binary — works on any GPU. Runs as subprocess with temp file I/O. """ caps = probe_upscale_capabilities() @@ -202,8 +348,6 @@ def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, st if not binary: raise RuntimeError("realesrgan-ncnn-vulkan binary not found") - # NCNN only supports integer scales (2, 3, 4) natively - # For non-integer scales: upscale to nearest integer, then resize to exact target model_scale = 4 if scale > 2.5 else 2 target_w = round(image.width * scale) target_h = round(image.height * scale) @@ -214,29 +358,19 @@ def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, st image.save(in_path, format="PNG") - # Model name for NCNN (bundled with binary) 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", + 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 - ) + 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") - - # Resize to exact target if scale was non-integer if result.width != target_w or result.height != target_h: result = result.resize((target_w, target_h), Image.Resampling.LANCZOS) @@ -246,17 +380,7 @@ def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, st # ── Public entry point ──────────────────────────────────────────────────────── def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]: - """ - Upscale image synchronously. Call via run_in_executor from async context. - - method values: - "auto" — pick best available automatically - "realesrgan_pytorch" — force PyTorch path - "realesrgan_ncnn" — force NCNN binary path - "lanczos" — force Lanczos - - Returns (png_bytes, method_used_label). - """ + """Upscale image synchronously. Returns (png_bytes, method_used_label).""" caps = probe_upscale_capabilities() if method == "auto": @@ -268,7 +392,6 @@ def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tupl return upscale_realesrgan_pytorch(image, scale) except Exception as e: print(f"Real-ESRGAN PyTorch failed, falling back: {e}") - # Fall through to next best if caps["realesrgan_ncnn"]: try: return upscale_realesrgan_ncnn(image, scale) @@ -282,7 +405,6 @@ def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tupl return upscale_realesrgan_ncnn(image, scale) except Exception as e: print(f"Real-ESRGAN NCNN failed, falling back: {e}") - # Fall through if caps["realesrgan_pytorch"]: try: return upscale_realesrgan_pytorch(image, scale) @@ -290,7 +412,6 @@ def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tupl print(f"Real-ESRGAN PyTorch fallback failed: {e}") return upscale_lanczos(image, scale) - # Default / lanczos return upscale_lanczos(image, scale) diff --git a/frontend/src/js/modules/image/upscale.js b/frontend/src/js/modules/image/upscale.js index 759b652..6ee805a 100644 --- a/frontend/src/js/modules/image/upscale.js +++ b/frontend/src/js/modules/image/upscale.js @@ -2,12 +2,8 @@ * Upscale — increase image resolution. * Fetches available methods from /api/print/upscale/available on first open. * Auto-selects the recommended method; user can override. - * - * Methods (in priority order, server picks best): - * auto — server picks best available - * realesrgan_pytorch — Real-ESRGAN via PyTorch (CUDA > MPS > CPU) - * realesrgan_ncnn — Real-ESRGAN NCNN Vulkan binary (any GPU) - * lanczos — always available, instant + * If no AI upscaler is found, polls /api/print/upscale/install-status while + * the backend auto-installs Real-ESRGAN NCNN Vulkan, then refreshes and continues. * * Menu target: image/upscale.upscale */ @@ -20,7 +16,6 @@ import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.j var instance = null; -// Method display labels const METHOD_LABELS = { auto: 'Auto (best available)', realesrgan_pytorch: 'Real-ESRGAN — PyTorch', @@ -45,15 +40,25 @@ class Image_upscale_class { return; } + // If a previous caps fetch showed no AI upscaler, check install progress var caps = await this._fetchCaps(); + if (!caps.realesrgan_pytorch && !caps.realesrgan_ncnn) { + await this._waitForInstall(caps); + // Re-fetch caps after install + this._caps = null; + caps = await this._fetchCaps(); + } + + this._showDialog(caps); + } + + _showDialog(caps) { var W = config.layer.width_original; var H = config.layer.height_original; - // Build method selector — only show what's available + auto var available = ['auto', ...caps.methods]; - var methodValues = [...new Set(available)]; // dedupe + var methodValues = [...new Set(available)]; - // Label each option, mark recommended var methodLabels = methodValues.map(m => { var label = METHOD_LABELS[m] || m; if (m === 'auto') { @@ -64,7 +69,6 @@ class Image_upscale_class { return label; }); - // Annotate with device info var deviceNote = ''; if (caps.realesrgan_pytorch) { var dev = caps.realesrgan_pytorch_device; @@ -77,8 +81,7 @@ class Image_upscale_class { deviceNote += 'NCNN Vulkan binary found. '; } if (!caps.realesrgan_pytorch && !caps.realesrgan_ncnn) { - deviceNote = 'No AI upscaler detected — Lanczos only. ' + - 'Install Real-ESRGAN for AI quality (see docs).'; + deviceNote = 'No AI upscaler available — Lanczos only.'; } var _this = this; @@ -103,7 +106,7 @@ class Image_upscale_class { { name: 'method', title: 'Method:', - value: methodLabels[0], // auto + value: methodLabels[0], values: methodLabels, type: 'select', }, @@ -114,7 +117,6 @@ class Image_upscale_class { }, ], on_finish: async function (params) { - // Map label back to method key var labelIdx = methodLabels.indexOf(params.method); var methodKey = labelIdx >= 0 ? methodValues[labelIdx] : 'auto'; var scale = parseFloat(params.scale); @@ -123,6 +125,52 @@ class Image_upscale_class { }); } + /** + * Poll install-status until done/failed, showing a progress bar notification. + */ + async _waitForInstall(caps) { + var installStatus = caps.ncnn_install_status || {}; + if (installStatus.state === 'done' || installStatus.state === 'failed') { + return; + } + + return new Promise((resolve) => { + var msg = alertify.message( + `
Installing Real-ESRGAN AI upscaler…
+ + 0%
`, + 0 + ); + + var poll = setInterval(async () => { + try { + var base = window.API_BASE_URL || ''; + var r = await fetch(`${base}/api/print/upscale/install-status`); + if (!r.ok) return; + var s = await r.json(); + + var bar = document.getElementById('esrgan-install-progress'); + var pct = document.getElementById('esrgan-install-pct'); + if (bar) bar.value = s.progress || 0; + if (pct) pct.textContent = `${s.progress || 0}%`; + + if (s.state === 'done') { + clearInterval(poll); + alertify.dismissAll(); + alertify.success('Real-ESRGAN NCNN installed ✓'); + resolve(); + } else if (s.state === 'failed') { + clearInterval(poll); + alertify.dismissAll(); + alertify.warning('AI upscaler install failed — using Lanczos.'); + resolve(); + } + } catch { /* network hiccup, keep polling */ } + }, 1500); + }); + } + async _fetchCaps() { if (this._caps) return this._caps; try { @@ -133,7 +181,6 @@ class Image_upscale_class { } } catch { /* ignore */ } - // Safe default if fetch failed if (!this._caps) { this._caps = { lanczos: true, @@ -142,6 +189,7 @@ class Image_upscale_class { recommended: 'lanczos', recommended_label: 'Lanczos', methods: ['lanczos'], + ncnn_install_status: { state: 'idle', progress: 0 }, }; } return this._caps; @@ -156,7 +204,7 @@ class Image_upscale_class { ? `Auto (${caps.recommended_label || 'best available'})` : (METHOD_LABELS[method] || method); - alertify.message(`Upscaling ${scale}× · ${methodLabel}...`, 0); + alertify.message(`Upscaling ${scale}× · ${methodLabel}…`, 0); try { var layerCanvas = document.createElement('canvas'); @@ -185,7 +233,6 @@ class Image_upscale_class { resultCanvas.height = img.naturalHeight; resultCanvas.getContext('2d').drawImage(img, 0, 0); - // Human-readable method label for undo history var usedLabel = result.method.replace('realesrgan_pytorch_', 'ESRGAN/') .replace('realesrgan_ncnn', 'ESRGAN/NCNN');