Claude
b2d064573f
Matches the existing vendor/easy-asterisk convention (used by services/asterisk.sh) instead of two one-off top-level directories that cluttered the repo root and didn't look like anything else next to setup.sh, lib/, services/, extras/. Only the two services' own SRC_DIR path resolution and header comments needed updating — nothing else in the repo referenced the old ./ai-stack / ./paintplus paths. Also documents vendor/ in README.md's Layout section.
495 lines
18 KiB
Python
495 lines
18 KiB
Python
"""
|
|
Upscale service — auto-detects best available method and runs it.
|
|
Auto-installs Real-ESRGAN NCNN Vulkan binary when Vulkan GPU is available.
|
|
Skips NCNN on headless/CPU-only machines and uses PyTorch CPU or Lanczos instead.
|
|
|
|
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 Vulkan GPU
|
|
4. Real-ESRGAN PyTorch CPU — AI quality, slow (~1-3 min)
|
|
5. Lanczos — always available, instant
|
|
|
|
Capability probe is run once at first call and cached.
|
|
NCNN binary is auto-downloaded only when Vulkan is detected.
|
|
Set REALESRGAN_NCNN=force env var to override the Vulkan check.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import shutil
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import urllib.request
|
|
import zipfile
|
|
from dataclasses import dataclass
|
|
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"
|
|
skipped = "skipped" # headless / no Vulkan
|
|
downloading = "downloading"
|
|
extracting = "extracting"
|
|
verifying = "verifying"
|
|
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:
|
|
return "realesrgan-ncnn-vulkan.exe" if "win" in sys.platform.lower() else "realesrgan-ncnn-vulkan"
|
|
|
|
|
|
def _vulkan_available() -> bool:
|
|
"""
|
|
Check whether a Vulkan-capable GPU is accessible.
|
|
Returns True if confident a GPU with Vulkan exists; False on headless/CPU-only.
|
|
Set REALESRGAN_NCNN=force to bypass this check.
|
|
"""
|
|
if os.environ.get("REALESRGAN_NCNN", "").lower() == "force":
|
|
return True
|
|
|
|
plat = sys.platform.lower()
|
|
|
|
if plat == "linux":
|
|
# DRI render nodes exist when a GPU is present and drivers loaded
|
|
dri = Path("/dev/dri")
|
|
if dri.exists() and list(dri.glob("renderD*")):
|
|
return True
|
|
# Fallback: vulkaninfo (not always installed)
|
|
if shutil.which("vulkaninfo"):
|
|
r = subprocess.run(["vulkaninfo", "--summary"],
|
|
capture_output=True, timeout=5)
|
|
if r.returncode == 0 and b"GPU" in r.stdout:
|
|
return True
|
|
return False
|
|
|
|
if plat == "darwin":
|
|
# macOS with Metal/MPS — Vulkan via MoltenVK always present on Apple Silicon/modern Intel
|
|
return True
|
|
|
|
if "win" in plat:
|
|
# Windows always has a display adapter; assume Vulkan available
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def _test_ncnn_binary(binary_path: Path) -> bool:
|
|
"""Run binary with --help to confirm it actually works (Vulkan loads ok)."""
|
|
try:
|
|
r = subprocess.run(
|
|
[str(binary_path), "--help"],
|
|
capture_output=True, timeout=15,
|
|
)
|
|
# NCNN binary exits 255 for --help but prints usage; that's fine.
|
|
# A Vulkan init failure produces "no vulkan device" on stderr.
|
|
stderr = r.stderr.decode(errors="replace").lower()
|
|
if "no vulkan" in stderr or "failed to create" in stderr:
|
|
return False
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
async def ensure_ncnn_installed() -> Optional[Path]:
|
|
"""
|
|
Check for Vulkan, then download+install the NCNN binary if needed.
|
|
Skips silently on headless/CPU-only machines.
|
|
Returns binary Path on success, None otherwise.
|
|
"""
|
|
global _install_status
|
|
|
|
binary_path = NCNN_DEST_DIR / _ncnn_binary_name()
|
|
|
|
# Already installed — quick verify it still works
|
|
if binary_path.exists() and os.access(binary_path, os.X_OK):
|
|
loop = asyncio.get_event_loop()
|
|
ok = await loop.run_in_executor(None, _test_ncnn_binary, binary_path)
|
|
if ok:
|
|
_install_status = InstallStatus(state=InstallState.done, progress=100,
|
|
message="Already installed.")
|
|
return binary_path
|
|
else:
|
|
# Binary exists but Vulkan broken — treat as headless
|
|
_install_status = InstallStatus(
|
|
state=InstallState.skipped,
|
|
message="Vulkan unavailable — skipping NCNN (using PyTorch CPU or Lanczos).",
|
|
)
|
|
return None
|
|
|
|
async with _install_lock:
|
|
# Re-check after lock
|
|
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 in (InstallState.downloading, InstallState.extracting,
|
|
InstallState.verifying):
|
|
return None # already running
|
|
|
|
# Check Vulkan before downloading anything
|
|
loop = asyncio.get_event_loop()
|
|
has_vulkan = await loop.run_in_executor(None, _vulkan_available)
|
|
if not has_vulkan:
|
|
_install_status = InstallStatus(
|
|
state=InstallState.skipped,
|
|
message="No Vulkan GPU detected — skipping NCNN install. "
|
|
"AI upscaling via PyTorch CPU or set REALESRGAN_NCNN=force to override.",
|
|
)
|
|
print("[upscale] Headless/no-Vulkan detected — skipping NCNN download.")
|
|
return None
|
|
|
|
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:
|
|
_install_status.progress = min(85, int(count * block * 85 / total))
|
|
urllib.request.urlretrieve(url, zip_path, _progress)
|
|
|
|
await loop.run_in_executor(None, _do_download)
|
|
|
|
# Extract
|
|
_install_status.state = InstallState.extracting
|
|
_install_status.progress = 88
|
|
_install_status.message = "Extracting…"
|
|
|
|
def _do_extract():
|
|
with zipfile.ZipFile(zip_path, "r") as zf:
|
|
zf.extractall(NCNN_DEST_DIR)
|
|
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)
|
|
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)
|
|
|
|
# Verify binary actually works
|
|
_install_status.state = InstallState.verifying
|
|
_install_status.progress = 95
|
|
_install_status.message = "Verifying Vulkan…"
|
|
|
|
ok = await loop.run_in_executor(None, _test_ncnn_binary, binary_path)
|
|
if not ok:
|
|
binary_path.unlink(missing_ok=True)
|
|
_install_status = InstallStatus(
|
|
state=InstallState.skipped,
|
|
message="Binary installed but Vulkan unavailable at runtime — "
|
|
"falling back to PyTorch CPU / Lanczos.",
|
|
)
|
|
print("[upscale] NCNN binary installed but Vulkan check failed — skipping.")
|
|
return None
|
|
|
|
_install_status = InstallStatus(
|
|
state=InstallState.done, progress=100,
|
|
message=f"Real-ESRGAN NCNN installed: {binary_path}",
|
|
)
|
|
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 available upscaling methods. 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:
|
|
install_state = _install_status.state
|
|
if install_state in (InstallState.downloading, InstallState.extracting, InstallState.verifying):
|
|
caps["recommended_label"] = "Lanczos (AI upscaler installing…)"
|
|
elif install_state == InstallState.skipped:
|
|
caps["recommended_label"] = "Lanczos (headless — no Vulkan GPU)"
|
|
else:
|
|
caps["recommended_label"] = "Lanczos (no AI upscaler found)"
|
|
|
|
_caps = caps
|
|
return caps
|
|
|
|
|
|
def _find_ncnn_binary() -> Optional[Path]:
|
|
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():
|
|
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]:
|
|
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]:
|
|
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_path = model_dir / f"RealESRGAN_x{model_scale}plus.pth"
|
|
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])
|
|
return _to_png_bytes(result), f"realesrgan_pytorch_{device}"
|
|
|
|
|
|
def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
|
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")
|
|
|
|
cmd = [
|
|
binary,
|
|
"-i", str(in_path), "-o", str(out_path),
|
|
"-s", str(model_scale), "-n", f"realesrgan-x{model_scale}plus", "-f", "png",
|
|
]
|
|
r = subprocess.run(cmd, capture_output=True, timeout=300)
|
|
if r.returncode != 0:
|
|
raise RuntimeError(f"realesrgan-ncnn-vulkan failed: {r.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 synchronously. Returns (png_bytes, method_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."""
|
|
loop = asyncio.get_event_loop()
|
|
return await loop.run_in_executor(None, upscale_sync, image, scale, method)
|