scripts/infer.py

#!/usr/bin/env python3
"""Run the released RAESR or RAESR-D 512px inference protocol on an image folder.

Both models use the same input protocol and the paper-reported mean-only color
matching applied to the float decoder output before PNG quantization.
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from raesr.artifacts import (
    DELIVERY_ROOT as PACKAGE_ROOT,
    DINO_S_WEIGHT_NAME,
    DINO_WEIGHT_NAME,
    build_paper_model,
    build_raesr_d_model,
    configure_dinov3,
    file_record,
    runtime_environment,
    runtime_record,
    verify_decoder,
    verify_frontend,
    verify_runtime,
    write_json,
)

MEAN_ONLY_PROTOCOL = "paper_mean_only_color_match_512px_bicubic_input_resize_v1"
COLOR_MATCH_FORMULA = "clamp(sr - mean(sr,H,W) + mean(lr,H,W), 0, 1)"
IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}
MODEL_CHOICES = ("raesr", "raesr-d")


def match_color_mean(sr: torch.Tensor, lr: torch.Tensor) -> torch.Tensor:
    """Shift each SR channel to the LR mean without changing its contrast."""
    sr = sr.float()
    lr = lr.float()
    mean_sr = sr.mean(dim=(-2, -1), keepdim=True)
    mean_lr = lr.mean(dim=(-2, -1), keepdim=True)
    return (sr - mean_sr + mean_lr).clamp(0.0, 1.0)


def collect_images(directory: Path) -> list[Path]:
    if not directory.is_dir():
        raise FileNotFoundError(f"input directory does not exist: {directory}")
    images = sorted(path for path in directory.iterdir() if path.suffix.lower() in IMAGE_SUFFIXES)
    if not images:
        raise ValueError(f"no supported images in {directory}")
    stems: set[str] = set()
    for path in images:
        if path.stem in stems:
            raise ValueError(f"ambiguous duplicate input stem: {path.stem}")
        stems.add(path.stem)
        with Image.open(path) as image:
            image.convert("RGB").load()
    return images


def is_valid_output(path: Path) -> bool:
    try:
        with Image.open(path) as image:
            image.verify()
        with Image.open(path) as image:
            return image.format == "PNG" and image.size == (512, 512) and image.mode in {"RGB", "RGBA"}
    except (OSError, ValueError):
        return False


def load_input(path: Path, device: str) -> torch.Tensor:
    array = np.asarray(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0
    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).contiguous().to(device)
    if tensor.shape[-2:] != (512, 512):
        tensor = F.interpolate(tensor, size=(512, 512), mode="bicubic", align_corners=False).clamp(0, 1)
    return tensor


def save_png(tensor: torch.Tensor, destination: Path) -> None:
    array = tensor[0].detach().float().cpu().numpy().transpose(1, 2, 0)
    image = Image.fromarray((array * 255.0).round().astype(np.uint8))
    temporary = destination.with_name(f".{destination.name}.tmp.{os.getpid()}")
    try:
        image.save(temporary, format="PNG")
        temporary.replace(destination)
    finally:
        if temporary.exists():
            temporary.unlink()


def build_model(args: argparse.Namespace):
    if args.model == "raesr-d":
        return build_raesr_d_model(
            args.frontend, args.decoder, args.dinov3_root, args.dinov3_weights, args.device
        )
    return build_paper_model(args.decoder, args.dinov3_root, args.dinov3_weights, args.device)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--model", choices=MODEL_CHOICES, default="raesr")
    parser.add_argument("--decoder", type=Path, default=PACKAGE_ROOT / "artifacts" / "ema_decoder.pt")
    parser.add_argument("--frontend", type=Path, default=PACKAGE_ROOT / "weights" / "raesr_d_frontend.pt",
                        help="RAESR-D frontend weights (used with --model raesr-d)")
    parser.add_argument("--dinov3-root", type=Path, required=True)
    parser.add_argument("--dinov3-weights", type=Path, required=True)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards:
        parser.error("--shard-index must be in [0, --num-shards)")

    images = collect_images(args.input_dir.expanduser().resolve())
    args.output_dir.mkdir(parents=True, exist_ok=True)
    selected = images[args.shard_index::args.num_shards]
    pending = [
        path for path in selected
        if args.overwrite or not is_valid_output(args.output_dir / f"{path.stem}.png")
    ]
    dinov3_weight_name = DINO_S_WEIGHT_NAME if args.model == "raesr-d" else DINO_WEIGHT_NAME
    base_manifest = {
        "model": args.model,
        "protocol": MEAN_ONLY_PROTOCOL,
        "color_match": "mean-only",
        "color_match_formula": COLOR_MATCH_FORMULA,
        "input_dir": str(args.input_dir.expanduser().resolve()),
        "output_dir": str(args.output_dir.resolve()),
        "shard": {"index": args.shard_index, "count": args.num_shards},
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "input_count": len(images),
        "selected_count": len(selected),
        "pending_count": len(pending),
        "inputs": [file_record(path) for path in selected],
        "decoder": verify_decoder(args.decoder),
        "dinov3": configure_dinov3(args.dinov3_root, args.dinov3_weights, dinov3_weight_name),
        "runtime": runtime_record(),
        "runtime_verification": verify_runtime(),
        "environment": runtime_environment(),
    }
    if args.model == "raesr-d":
        base_manifest["frontend"] = verify_frontend(args.frontend)
    manifest_name = f"run_manifest.shard_{args.shard_index:02d}.json"
    if args.dry_run:
        base_manifest["status"] = "DRY_RUN"
        write_json(args.output_dir / manifest_name, base_manifest)
        print(f"validated {len(selected)} inputs; manifest -> {args.output_dir / manifest_name}")
        return

    model = build_model(args)
    for index, path in enumerate(pending, start=1):
        lr = load_input(path, args.device)
        with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
            output = model(lr).clamp(0.0, 1.0)
        output = match_color_mean(output, lr)
        save_png(output, args.output_dir / f"{path.stem}.png")
        if index % 25 == 0 or index == len(pending):
            print(f"{index}/{len(pending)}")

    outputs = [args.output_dir / f"{path.stem}.png" for path in selected]
    if not all(is_valid_output(path) for path in outputs):
        raise RuntimeError("one or more expected outputs are missing or invalid after inference")
    base_manifest["status"] = "COMPLETE"
    base_manifest["outputs"] = [file_record(path) for path in outputs]
    write_json(args.output_dir / manifest_name, base_manifest)
    print(f"wrote {len(pending)} images; manifest -> {args.output_dir / manifest_name}")


if __name__ == "__main__":
    main()