upscale.py

#!/usr/bin/env python3
"""Universal AI image upscaler (RGBA-aware) — a small, reusable CLI.

  python upscale.py <input> <output> [--scale N] [--weights path] [--ensemble]

Loads ANY super-resolution model via spandrel (ESRGAN/Real-ESRGAN, HAT, DAT, SwinIR, SRFormer,
OmniSR, …); falls back to a built-in RRDBNet for ESRGAN-family weights if spandrel is absent.
RGB is upscaled by the model at its native scale; any alpha channel is resized with Lanczos so
transparency survives. --ensemble runs the geometric self-ensemble (8 orientations, averaged) for
a cleaner/sharper result. Runs on CUDA when available, else CPU.
"""
import argparse
import re

import numpy as np
import torch
from PIL import Image


def _convert_old_esrgan(sd: dict) -> dict:
    """Remap an OLD-ESRGAN state dict (`model.1.sub.N…`) to the RealESRGAN RRDBNet layout."""
    out = {}
    for k, v in sd.items():
        m = re.match(r"model\.1\.sub\.(\d+)\.RDB(\d)\.conv(\d)\.0\.(weight|bias)$", k)
        if m:
            i, j, c, wb = m.groups()
            out[f"body.{i}.rdb{j}.conv{c}.{wb}"] = v
            continue
        out[
            k.replace("model.1.sub.23.", "conv_body.")
            .replace("model.0.", "conv_first.")
            .replace("model.3.", "conv_up1.")
            .replace("model.6.", "conv_up2.")
            .replace("model.8.", "conv_hr.")
            .replace("model.10.", "conv_last.")
        ] = v
    return out


def load_model(weights: str, device: str):
    """Return (runner, scale). `runner(tensor NCHW 0-1)` -> upscaled tensor NCHW."""
    try:
        import spandrel

        try:
            import spandrel_extra_arches

            spandrel_extra_arches.install()
        except Exception:
            pass
        desc = spandrel.ModelLoader().load_from_file(weights)
        desc.to(device).eval()
        print(f"model: {desc.architecture.name} x{desc.scale} ({weights.split('/')[-1]})")
        return desc, int(desc.scale)
    except ImportError:
        from rrdbnet import RRDBNet

        model = RRDBNet(3, 3, 64, 23, 32)
        sd = torch.load(weights, map_location="cpu")
        if isinstance(sd, dict):
            sd = sd.get("params_ema", sd.get("params", sd))
        if any(k.startswith("model.") for k in sd):
            sd = _convert_old_esrgan(sd)
        model.load_state_dict(sd, strict=True)
        model.eval().to(device)
        print(f"model: RRDBNet x4 ({weights.split('/')[-1]}) [manual]")
        return model, 4


# Dihedral (flip/transpose) augmentations for geometric self-ensemble.
def _aug(x: torch.Tensor, i: int) -> torch.Tensor:
    if i & 1:
        x = x.flip(-1)
    if i & 2:
        x = x.flip(-2)
    if i & 4:
        x = x.transpose(-2, -1)
    return x.contiguous()


def _unaug(x: torch.Tensor, i: int) -> torch.Tensor:
    if i & 4:
        x = x.transpose(-2, -1)
    if i & 2:
        x = x.flip(-2)
    if i & 1:
        x = x.flip(-1)
    return x


@torch.no_grad()
def run_model(runner, rgb: np.ndarray, device: str, ensemble: bool = False) -> np.ndarray:
    t = torch.from_numpy(rgb.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)
    if ensemble:
        acc = None
        for i in range(8):
            y = _unaug(runner(_aug(t, i)).clamp(0, 1), i)
            acc = y if acc is None else acc + y
        out = acc / 8.0
    else:
        out = runner(t).clamp(0, 1)
    out = out.squeeze(0).permute(1, 2, 0).cpu().numpy()
    return (out * 255.0).round().astype(np.uint8)


def main() -> None:
    ap = argparse.ArgumentParser(description="Universal AI upscaler (RGBA-aware)")
    ap.add_argument("input")
    ap.add_argument("output")
    ap.add_argument("--weights", default="weights/4x-UltraSharp.pth")
    ap.add_argument("--scale", type=float, default=0.0, help="final output scale (default = model's)")
    ap.add_argument(
        "--ensemble",
        action="store_true",
        help="geometric self-ensemble (8x GPU work, cleaner+sharper, de-speckles flat areas)",
    )
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    gpu = torch.cuda.get_device_name(0) if device == "cuda" else "CPU"
    print(f"device: {device} ({gpu})")

    runner, model_scale = load_model(args.weights, device)

    im = Image.open(args.input)
    has_alpha = im.mode in ("RGBA", "LA") or (im.mode == "P" and "transparency" in im.info)
    im = im.convert("RGBA" if has_alpha else "RGB")
    arr = np.array(im)

    out_rgb = run_model(runner, arr[..., :3], device, ensemble=args.ensemble)
    if has_alpha:
        alpha = Image.fromarray(arr[..., 3]).resize(
            (out_rgb.shape[1], out_rgb.shape[0]), Image.LANCZOS
        )
        res = Image.fromarray(np.dstack([out_rgb, np.array(alpha)]), "RGBA")
    else:
        res = Image.fromarray(out_rgb, "RGB")

    target = args.scale if args.scale > 0 else float(model_scale)
    if abs(res.width - im.width * target) > 0.5:
        res = res.resize((round(im.width * target), round(im.height * target)), Image.LANCZOS)

    res.save(args.output)
    print(f"saved: {args.output}  {res.size}  ({im.size} -> {target}x)")


if __name__ == "__main__":
    main()