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