#!/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()