diff --git a/backend/app/routers/print_tools.py b/backend/app/routers/print_tools.py index 4486d3d..2d5b1ec 100644 --- a/backend/app/routers/print_tools.py +++ b/backend/app/routers/print_tools.py @@ -61,7 +61,8 @@ class FrameFitRequest(BaseModel): class UpscaleRequest(BaseModel): image: str # base64 scale: float = 2.0 # 1.5, 2, 3, 4 - method: Literal["lanczos", "ai"] = "lanczos" + # auto = pick best available; lanczos = always works; realesrgan_pytorch / realesrgan_ncnn = explicit + method: str = "auto" # ── Frame sizes endpoint ─────────────────────────────────────────────────── @@ -293,83 +294,63 @@ def _mirror_fill(canvas, mask, scaled, gap_dir, gap_a, gap_b, target_w, target_h # ── Upscale ──────────────────────────────────────────────────────────────── +@router.post("/upscale/refresh-caps") +def upscale_refresh_caps(): + """Bust the capability cache (call after installing Real-ESRGAN without restarting).""" + from app.services.upscale import invalidate_caps_cache, probe_upscale_capabilities + invalidate_caps_cache() + return probe_upscale_capabilities() + + +@router.get("/upscale/available") +def upscale_available(): + """ + Return capability probe: which upscale methods are available, + which device will be used, and which method is recommended. + Frontend uses this to populate the method selector. + """ + from app.services.upscale import probe_upscale_capabilities + caps = probe_upscale_capabilities() + return caps + + @router.post("/upscale") async def upscale(req: UpscaleRequest): """ - Upscale image. - method=lanczos — always available, fast, good for clean images - method=ai — Real-ESRGAN if installed, else falls back to lanczos + Upscale image. method values: + auto — pick best available (recommended) + realesrgan_pytorch — Real-ESRGAN via PyTorch (CUDA/MPS/CPU) + realesrgan_ncnn — Real-ESRGAN NCNN Vulkan binary + lanczos — always available, instant + Any AI method falls back to the next best if unavailable. """ if not (1.1 <= req.scale <= 8.0): raise HTTPException(status_code=400, detail="scale must be 1.1–8.0") + valid_methods = {"auto", "realesrgan_pytorch", "realesrgan_ncnn", "lanczos"} + if req.method not in valid_methods: + raise HTTPException(status_code=400, + detail=f"method must be one of {sorted(valid_methods)}") + try: image = Image.open(BytesIO(_decode(req.image))).convert("RGB") except Exception as e: raise HTTPException(status_code=400, detail=f"Could not decode image: {e}") orig_w, orig_h = image.size - new_w = round(orig_w * req.scale) - new_h = round(orig_h * req.scale) - method_used = req.method - - if req.method == "ai": - try: - result_bytes = await asyncio.get_event_loop().run_in_executor( - None, _realesrgan_upscale, image, req.scale - ) - result = Image.open(BytesIO(result_bytes)).convert("RGB") - method_used = "realesrgan" - except Exception as e: - print(f"Real-ESRGAN failed, using Lanczos: {e}") - result = image.resize((new_w, new_h), Image.Resampling.LANCZOS) - method_used = "lanczos_fallback" - else: - result = image.resize((new_w, new_h), Image.Resampling.LANCZOS) + try: + from app.services.upscale import upscale_image + result_bytes, method_used = await upscale_image(image, req.scale, req.method) + result = Image.open(BytesIO(result_bytes)) + except Exception as e: + import traceback; traceback.print_exc() + raise HTTPException(status_code=500, detail=str(e)) return { - "result": _encode(_to_png(result)), + "result": _encode(result_bytes), "method": method_used, "original": {"width": orig_w, "height": orig_h}, - "output": {"width": result.width, "height": result.height}, - "scale": req.scale, + "output": {"width": result.width, "height": result.height}, + "scale": req.scale, } - - -def _realesrgan_upscale(image: Image.Image, scale: float) -> bytes: - """Run Real-ESRGAN upscaling. Raises if not installed.""" - from basicsr.archs.rrdbnet_arch import RRDBNet - from realesrgan import RealESRGANer - import torch - import numpy as np - - model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, - num_block=23, num_grow_ch=32, scale=4) - upsampler = RealESRGANer( - scale=4, - model_path=None, # auto-download - model=model, - tile=400, - tile_pad=10, - pre_pad=0, - half=torch.cuda.is_available(), - ) - img_np = np.array(image)[:, :, ::-1] # RGB→BGR for cv2 - output, _ = upsampler.enhance(img_np, outscale=scale) - result = Image.fromarray(output[:, :, ::-1]) # BGR→RGB - buf = BytesIO() - result.save(buf, format="PNG") - return buf.getvalue() - - -@router.get("/upscale/available") -def upscale_available(): - """Check which upscale methods are available.""" - ai_available = False - try: - import realesrgan # noqa: F401 - ai_available = True - except ImportError: - pass - return {"lanczos": True, "realesrgan": ai_available} diff --git a/backend/app/services/upscale.py b/backend/app/services/upscale.py new file mode 100644 index 0000000..9e8bbba --- /dev/null +++ b/backend/app/services/upscale.py @@ -0,0 +1,300 @@ +""" +Upscale service — auto-detects best available method and runs it. + +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. +""" + +import asyncio +import os +import shutil +import subprocess +import sys +import tempfile +from io import BytesIO +from pathlib import Path +from typing import Optional + +from PIL import Image + +# ── 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, # "cuda" | "mps" | "cpu" + "realesrgan_ncnn": False, + "realesrgan_ncnn_path": None, + "recommended": "lanczos", + "recommended_label": "Lanczos (no AI upscaler found)", + "methods": ["lanczos"], + } + + # ── 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 (install Real-ESRGAN for AI quality)" + + _caps = caps + return caps + + +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"), + 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 + + +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. + 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 + from realesrgan import RealESRGANer + + 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 + + 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"), # fp16 only on CUDA + device=torch.device(device), + ) + + import numpy as np + img_bgr = np.array(image)[:, :, ::-1].copy() # RGB→BGR + enhanced, _ = upsampler.enhance(img_bgr, outscale=scale) + result = Image.fromarray(enhanced[:, :, ::-1]) # BGR→RGB + + 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 (Intel/AMD/integrated/Apple). + 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") + + # 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) + + 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 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", + ] + + 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) + + 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. 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). + """ + 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}") + # Fall through to next best + 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}") + # Fall through + 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) + + # Default / lanczos + 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) diff --git a/frontend/src/js/modules/image/upscale.js b/frontend/src/js/modules/image/upscale.js index fd172cf..759b652 100644 --- a/frontend/src/js/modules/image/upscale.js +++ b/frontend/src/js/modules/image/upscale.js @@ -1,9 +1,13 @@ /** * Upscale — increase image resolution. + * Fetches available methods from /api/print/upscale/available on first open. + * Auto-selects the recommended method; user can override. * - * Lanczos: always available, fast, good for clean/sharp images. - * AI (Real-ESRGAN): much better for photos — restores texture, sharpness. - * Requires `realesrgan-ncnn-vulkan` or `basicsr` + `realesrgan` Python packages. + * 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 * * Menu target: image/upscale.upscale */ @@ -16,6 +20,14 @@ 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', + realesrgan_ncnn: 'Real-ESRGAN — NCNN Vulkan', + lanczos: 'Lanczos (fast, no AI)', +}; + class Image_upscale_class { constructor() { @@ -24,7 +36,7 @@ class Image_upscale_class { this.Base_layers = new Base_layers_class(); this.Dialog = new Dialog_class(); this.isProcessing = false; - this._aiAvailable = null; + this._caps = null; } async upscale() { @@ -33,24 +45,41 @@ class Image_upscale_class { return; } + var caps = await this._fetchCaps(); var W = config.layer.width_original; var H = config.layer.height_original; - // Check AI availability once, cache it - if (this._aiAvailable === null) { - try { - var base = window.API_BASE_URL || ''; - var r = await fetch(`${base}/api/print/upscale/available`); - var data = r.ok ? await r.json() : {}; - this._aiAvailable = data.realesrgan || false; - } catch { - this._aiAvailable = false; - } - } + // Build method selector — only show what's available + auto + var available = ['auto', ...caps.methods]; + var methodValues = [...new Set(available)]; // dedupe - var aiNote = this._aiAvailable - ? 'Real-ESRGAN AI upscaling available.' - : 'AI upscaling not installed (Real-ESRGAN). Using Lanczos only.'; + // Label each option, mark recommended + var methodLabels = methodValues.map(m => { + var label = METHOD_LABELS[m] || m; + if (m === 'auto') { + label = `Auto → ${caps.recommended_label}`; + } else if (m === caps.recommended && m !== 'auto') { + label += ' ★'; + } + return label; + }); + + // Annotate with device info + var deviceNote = ''; + if (caps.realesrgan_pytorch) { + var dev = caps.realesrgan_pytorch_device; + var devLabel = dev === 'cuda' ? 'CUDA GPU' + : dev === 'mps' ? 'Apple Silicon' + : 'CPU (slow — ~1–3 min for large images)'; + deviceNote += `PyTorch: ${devLabel}. `; + } + if (caps.realesrgan_ncnn) { + 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).'; + } var _this = this; @@ -60,7 +89,8 @@ class Image_upscale_class { { title: '', html: `