root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Auto-install Real-ESRGAN NCNN Vulkan binary on first use

Browse Source 
- upscale.py: add InstallStatus dataclass + ensure_ncnn_installed() async
  function that downloads and extracts the NCNN binary for the current
  platform (Linux/macOS/Windows), tracks progress (0-100%), and busts the
  caps cache when done
- main.py: trigger ensure_ncnn_installed() as a background task on app
  startup when no AI upscaler is detected
- print_tools.py: /upscale/available triggers install task when no AI
  upscaler found; new GET /upscale/install-status endpoint for polling
- upscale.js: if no AI upscaler on open, poll install-status showing a
  progress bar notification, then refresh caps and proceed when done

https://claude.ai/code/session_01B58MaJCU1R6KwBDJCp8AfN
This commit is contained in:
Claude
2026-06-09 18:38:45 +00:00
parent cea9ee9d6c
commit ed2a0d7f0c
4 changed files with 262 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/main.py
+7 -1
View File
@@ -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
backend/app/routers/print_tools.py
+23 -2
View File
@@ -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):
"""
backend/app/services/upscale.py
+167 -46
View File
@@ -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)