backend/app/services/upscale_engine.py

"""
upscale_engine.py — Real-ESRGAN x4plus inference engine.

This module is INDEPENDENT of the LaMa/watermark-removal pipeline.
Do not import or reference any LaMa-related modules from here.

Architecture
------------
BaseUpscaleEngine(ABC)
    └── RealESRGANEngine
            ├── __init__: loads RRDBNet + params_ema weights once, sets eval mode
            ├── upscale:  dispatches to whole-image or tiled path based on image size
            ├── _infer_tensor: PIL → tensor → model → PIL (shared by both paths)
            └── _upscale_tiled: splits large images into overlapping tiles, runs
                                _infer_tensor per tile, stitches results back

Tiling algorithm
----------------
For each tile at grid position (row, col):

  Input-space:
    x0, y0        — start of logical (non-overlapping) tile
    x1, y1        — end of logical tile (clamped to image size)
    rx0..ry1      — read region = logical tile expanded by `overlap` on each
                    side, clamped to image bounds so we never read outside

  Per-side overlap (varies at image boundary — 0 at edges, `overlap` elsewhere):
    left_ovlp     = x0 - rx0
    top_ovlp      = y0 - ry0
    right_ovlp    = rx1 - x1
    bottom_ovlp   = ry1 - y1

  SR-space (after model, each dimension × SCALE=4):
    Full SR patch is (rx1-rx0)*4 × (ry1-ry0)*4.
    Crop off left_ovlp*4 .. right_ovlp*4 pixels on each axis to discard the
    "low-context" border produced by the model's limited receptive field near
    patch edges.  The remaining region is exactly (x1-x0)*4 × (y1-y0)*4 and
    is written to output[y0*4:y1*4, x0*4:x1*4].

Math check for a middle tile (tile_size=512, overlap=32, scale=4):
  read = 576×576 → SR = 2304×2304 → crop 128px each side → 2048×2048 = 512×4 ✓

Edge tile (left boundary, left_ovlp=0):
  read = 544×... → SR = 2176×... → crop 0 left, 128 right → 2048 wide = 512×4 ✓

RGB note
--------
RealESRGAN_x4plus.pth was trained on RGB images.  The reference implementation
reads with cv2 (BGR) and swaps to RGB before the model call.  Since PIL is
natively RGB, no channel-swap is needed here — we feed the model directly from
PIL's array representation.
"""

import logging
import math
import time
from abc import ABC, abstractmethod
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from app.services.rrdbnet_arch import RRDBNet
from app.config import (
    UPSCALE_TILE_THRESHOLD,
    UPSCALE_TILE_SIZE,
    UPSCALE_TILE_OVERLAP,
)

logger = logging.getLogger(__name__)

# RealESRGAN x4plus is always a 4× upscaler — this is not a runtime setting.
_SCALE = 4


# ---------------------------------------------------------------------------
# Abstract base
# ---------------------------------------------------------------------------

class BaseUpscaleEngine(ABC):
    """Interface for all upscale engines."""

    @abstractmethod
    def upscale(self, image: Image.Image) -> Image.Image:
        """Upscale *image* and return the result as a new PIL Image (RGB).

        The caller is responsible for ensuring *image* is a valid RGB image.
        The returned image will have dimensions (width * scale, height * scale).
        """


# ---------------------------------------------------------------------------
# Real-ESRGAN implementation
# ---------------------------------------------------------------------------

class RealESRGANEngine(BaseUpscaleEngine):
    """4× upscaler backed by RealESRGAN_x4plus.pth (RRDBNet, 23 blocks)."""

    def __init__(self, model_path: str, device: str = "cpu") -> None:
        """Load the RRDBNet model from *model_path*.

        Args:
            model_path: Absolute or project-relative path to
                        RealESRGAN_x4plus.pth.
            device:     PyTorch device string ("cpu", "cuda", etc.).

        Raises:
            FileNotFoundError: If *model_path* does not exist on disk.
        """
        path = Path(model_path)
        if not path.exists():
            raise FileNotFoundError(
                f"RealESRGAN model checkpoint not found at: {path}\n"
                "Download RealESRGAN_x4plus.pth and place it at that path."
            )

        self.device = torch.device(device)

        # Build the network (x4plus config — must exactly match the checkpoint).
        self._model = RRDBNet(
            num_in_ch   = 3,
            num_out_ch  = 3,
            scale       = 4,
            num_feat    = 64,
            num_block   = 23,
            num_grow_ch = 32,
        )

        # Load checkpoint — the params_ema key holds the EMA weights, which
        # are the higher-quality copy saved during training.  Fall back to
        # "params" if "params_ema" is absent (non-EMA checkpoint variant).
        ckpt = torch.load(str(path), map_location="cpu", weights_only=False)
        if "params_ema" in ckpt:
            state_dict = ckpt["params_ema"]
        elif "params" in ckpt:
            state_dict = ckpt["params"]
        else:
            # Bare state dict (no wrapper key — unusual but handle gracefully).
            state_dict = ckpt

        self._model.load_state_dict(state_dict, strict=True)
        self._model.eval()
        self._model.to(self.device)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def upscale(self, image: Image.Image) -> Image.Image:
        """Return a 4× upscaled version of *image*.

        Chooses between whole-image inference (fast, suitable for small inputs)
        and tiled inference (memory-safe for large inputs) based on
        UPSCALE_TILE_THRESHOLD from config.

        Args:
            image: Source PIL Image (any mode — converted to RGB internally).

        Returns:
            PIL Image (RGB) with dimensions (image.width*4, image.height*4).
        """
        img = image.convert("RGB")
        w, h = img.size

        if max(w, h) < UPSCALE_TILE_THRESHOLD:
            return self._infer_tensor(img)
        return self._upscale_tiled(img)

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _infer_tensor(self, pil_image: Image.Image) -> Image.Image:
        """Run one PIL Image (RGB) through the model and return the SR result.

        This is called both for whole-image inference and for individual tiles
        during tiled inference.  *pil_image* must already be in RGB mode.

        Conversion pipeline:
            PIL (RGB, uint8) → numpy float32 [0,1] → torch (1,3,H,W)
                → model → torch (1,3,H*4,W*4) → numpy uint8 → PIL (RGB)
        """
        # HWC uint8 → HWC float32 [0,1] → CHW → (1,C,H,W)
        img_np = np.asarray(pil_image, dtype=np.float32) / 255.0
        img_t  = torch.from_numpy(img_np.transpose(2, 0, 1)).unsqueeze(0).to(self.device)

        with torch.no_grad():
            out_t = self._model(img_t)

        # (1,C,H,W) → HWC, clamp, scale to uint8
        out_np = out_t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy()
        out_u8 = (out_np * 255.0).round().astype(np.uint8)
        return Image.fromarray(out_u8, mode="RGB")

    def _upscale_tiled(self, image: Image.Image) -> Image.Image:
        """Tiled 4× upscale for images whose max dimension exceeds the threshold.

        Splits *image* (already RGB) into overlapping tiles, runs each through
        _infer_tensor, discards the scaled overlap borders, and stitches the
        cropped tiles into a full-size output canvas.

        See module docstring for the tiling math.
        """
        w, h = image.size
        tile  = UPSCALE_TILE_SIZE
        ovlp  = UPSCALE_TILE_OVERLAP
        sc    = _SCALE

        out_w, out_h = w * sc, h * sc
        # Pre-allocate the output canvas in numpy (fill with zeros — every pixel
        # will be written exactly once by a tile; no pixel is ever left at 0
        # unless the input itself was smaller than one tile, which can't happen
        # here because we only tile when max(w,h) > UPSCALE_TILE_THRESHOLD ≥ tile).
        canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8)

        n_cols   = math.ceil(w / tile)
        n_rows   = math.ceil(h / tile)
        n_tiles  = n_cols * n_rows
        t_tiled_start = time.monotonic()

        logger.info(
            "Tiled upscale: image %dx%d → %d cols × %d rows = %d tiles "
            "(tile=%d overlap=%d scale=%d → output %dx%d)",
            w, h, n_cols, n_rows, n_tiles, tile, ovlp, sc, out_w, out_h,
        )

        for row in range(n_rows):
            for col in range(n_cols):
                tile_num = row * n_cols + col
                t_tile = time.monotonic()

                # ── Logical tile bounds (input-image coords) ──────────────
                x0 = col * tile
                y0 = row * tile
                x1 = min(x0 + tile, w)
                y1 = min(y0 + tile, h)

                # ── Padded read region (clamped to image bounds) ──────────
                rx0 = max(0, x0 - ovlp)
                ry0 = max(0, y0 - ovlp)
                rx1 = min(w, x1 + ovlp)
                ry1 = min(h, y1 + ovlp)

                # ── Actual overlap applied on each side (0 at image edge) ─
                left_ovlp   = x0 - rx0
                top_ovlp    = y0 - ry0
                right_ovlp  = rx1 - x1
                bottom_ovlp = ry1 - y1

                # ── Extract patch, upscale ─────────────────────────────────
                patch    = image.crop((rx0, ry0, rx1, ry1))
                sr_patch = self._infer_tensor(patch)   # PIL, size = (rx1-rx0)*4 × (ry1-ry0)*4

                t_tile_elapsed = time.monotonic() - t_tile
                logger.info(
                    "  tile %2d/%d (row=%d col=%d) read=%dx%d: %.1fs",
                    tile_num + 1, n_tiles, row, col,
                    rx1 - rx0, ry1 - ry0,
                    t_tile_elapsed,
                )

                # ── Crop scaled overlap borders from SR output ─────────────
                # Discard the model's less-accurate near-edge output;
                # keep only the center region corresponding to the logical tile.
                sr_x0 = left_ovlp   * sc
                sr_y0 = top_ovlp    * sc
                sr_x1 = sr_patch.width  - right_ovlp  * sc
                sr_y1 = sr_patch.height - bottom_ovlp * sc
                tile_sr = sr_patch.crop((sr_x0, sr_y0, sr_x1, sr_y1))

                # ── Write into output canvas ───────────────────────────────
                out_x0 = x0 * sc
                out_y0 = y0 * sc
                tile_arr = np.asarray(tile_sr)
                canvas[out_y0 : out_y0 + tile_arr.shape[0],
                       out_x0 : out_x0 + tile_arr.shape[1]] = tile_arr

        logger.info(
            "Tiled upscale complete: %d tiles in %.1fs (avg %.1fs/tile)",
            n_tiles,
            time.monotonic() - t_tiled_start,
            (time.monotonic() - t_tiled_start) / max(n_tiles, 1),
        )
        return Image.fromarray(canvas, mode="RGB")