PaintPlus/backend/app/services/upscale.py

"""
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)