Claude
cea9ee9d6c
Detection priority (probed once, cached): 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 via Vulkan (no CUDA needed) 4. Real-ESRGAN PyTorch CPU → works, slow (warned in UI) 5. Lanczos → always available, instant fallback Backend: - services/upscale.py: full capability probe (probe_upscale_capabilities), implementations for PyTorch (CUDA/MPS/CPU auto-device) and NCNN binary, upscale_sync() resolves method with fallback chain, async upscale_image() runs in thread pool - print_tools.py: /api/print/upscale uses new service; method="auto" by default; GET /api/print/upscale/available returns full capability map with device info and recommended_label; POST /api/print/upscale/refresh-caps busts cache without restart (useful after installing NCNN binary into container) Frontend: - upscale.js: fetches capability map on first open; builds method selector showing only available options; labels recommended method with ★; shows device info (CUDA/MPS/CPU/NCNN) in dialog; maps display label back to method key on submit; shows actual method used in success toast and undo history entry Scripts: - scripts/download_realesrgan.py: downloads NCNN Vulkan binary for current platform (Linux/macOS/Windows) to /app/data/models/realesrgan/; makes executable; run inside container or locally https://claude.ai/code/session_01B58MaJCU1R6KwBDJCp8AfN
301 lines
11 KiB
Python
301 lines
11 KiB
Python
"""
|
|
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)
|