|
|
| """Standalone inference for the Universal Upscale model.
|
|
|
| This repository intentionally ships no training code. `model_scripted.pt` already
|
| contains the network graph, the trained weights and the tiled inference logic,
|
| so this script only needs torch + Pillow + numpy.
|
| """
|
| import argparse
|
| from pathlib import Path
|
| import warnings
|
|
|
| import numpy as np
|
| import torch
|
| from PIL import Image
|
|
|
| warnings.filterwarnings("ignore", category=FutureWarning, message=".*torch.jit.load.*")
|
|
|
| MAX_OUTPUT_PIXELS = 64_000_000
|
|
|
|
|
| def load_model(path, repo_id=None, revision="main", device=None):
|
| local = Path(path)
|
| if not local.exists():
|
| script_local = Path(__file__).resolve().parent / path
|
| if script_local.exists():
|
| local = script_local
|
| if not local.exists():
|
| from huggingface_hub import hf_hub_download
|
|
|
| if not repo_id:
|
| raise SystemExit("model_scripted.pt not found locally: pass --repo-id <user>/<repo> to download it.")
|
| local = Path(hf_hub_download(repo_id=repo_id, filename="model_scripted.pt", revision=revision))
|
| model = torch.jit.load(str(local), map_location=device if device is not None else "cpu")
|
| model.eval()
|
| return model
|
|
|
|
|
| def pil_to_tensor(img):
|
| arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
|
| return torch.from_numpy(arr).permute(2, 0, 1).contiguous().unsqueeze(0)
|
|
|
|
|
| def tensor_to_pil(t):
|
| arr = t.detach().float().clamp(0.0, 1.0).permute(1, 2, 0).cpu().numpy()
|
| return Image.fromarray((arr * 255.0).round().astype(np.uint8), mode="RGB")
|
|
|
|
|
| def compute_target(img, args):
|
| w, h = img.size
|
| if args.scale is not None:
|
| scale = float(args.scale)
|
| elif args.target_long_side is not None:
|
| scale = float(args.target_long_side) / float(max(w, h))
|
| elif args.target_width or args.target_height:
|
| candidates = []
|
| if args.target_width:
|
| candidates.append(float(args.target_width) / float(w))
|
| if args.target_height:
|
| candidates.append(float(args.target_height) / float(h))
|
| scale = min(candidates)
|
| else:
|
| raise SystemExit("Provide --scale, --target-long-side, --target-width or --target-height.")
|
| if scale <= 0:
|
| raise SystemExit("Target scale must be positive.")
|
|
|
| out_w, out_h = max(1, int(round(w * scale))), max(1, int(round(h * scale)))
|
| canvas_w, canvas_h = out_w, out_h
|
| if args.pad_to_target and args.target_width and args.target_height:
|
| canvas_w, canvas_h = int(args.target_width), int(args.target_height)
|
| if canvas_w < out_w or canvas_h < out_h:
|
| raise SystemExit("--pad-to-target canvas must contain the aspect-preserving output.")
|
| return out_w, out_h, canvas_w, canvas_h, scale
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser(description="Universal image upscaler (self-contained TorchScript model).")
|
| ap.add_argument("--input", required=True, help="Input image (.png/.jpg/.webp/.tif)")
|
| ap.add_argument("--output", required=True, help="Output image (.png/.jpg/.webp/.tif)")
|
| ap.add_argument("--model", default="model_scripted.pt", help="Path to the TorchScript model")
|
| ap.add_argument("--repo-id", default=None, help="Download the model from this Hub repo instead")
|
| ap.add_argument("--scale", type=float, default=None, help="Arbitrary scale factor, e.g. 4 or 10")
|
| ap.add_argument("--target-long-side", type=int, default=None, help="Target long side in pixels, e.g. 3000")
|
| ap.add_argument("--target-width", type=int, default=None)
|
| ap.add_argument("--target-height", type=int, default=None)
|
| ap.add_argument("--pad-to-target", action="store_true", help="Pad the canvas to the exact target size")
|
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
| ap.add_argument("--amp", action="store_true", help="fp16 autocast (CUDA only)")
|
| ap.add_argument("--webp-lossless", action="store_true")
|
| args = ap.parse_args()
|
|
|
| device = torch.device(args.device)
|
| img = Image.open(args.input).convert("RGB")
|
| out_w, out_h, canvas_w, canvas_h, scale = compute_target(img, args)
|
| if canvas_w * canvas_h > MAX_OUTPUT_PIXELS:
|
| raise SystemExit(f"Output {canvas_w}x{canvas_h} is too large ({canvas_w * canvas_h} px). Lower --scale/--target.")
|
|
|
| print(f"Input {img.width}x{img.height} -> output {out_w}x{out_h} (scale {scale:.4f}, device {device})")
|
|
|
| if scale <= 1.0001:
|
| out_img = img.resize((out_w, out_h), Image.LANCZOS)
|
| else:
|
| model = load_model(args.model, args.repo_id, device=device)
|
| lr = pil_to_tensor(img).to(device)
|
| with torch.autocast(device_type=device.type, enabled=bool(args.amp and device.type == "cuda")):
|
| out = model(lr, out_h, out_w)
|
| out_img = tensor_to_pil(out.squeeze(0))
|
|
|
| if canvas_w != out_w or canvas_h != out_h:
|
| canvas = Image.new("RGB", (canvas_w, canvas_h), (0, 0, 0))
|
| canvas.paste(out_img, ((canvas_w - out_w) // 2, (canvas_h - out_h) // 2))
|
| out_img = canvas
|
|
|
| suffix = Path(args.output).suffix.lower()
|
| kwargs = {}
|
| dpi = img.info.get("dpi")
|
| if dpi:
|
| kwargs["dpi"] = dpi
|
| if suffix in (".jpg", ".jpeg"):
|
| kwargs.update(quality=95, subsampling=0)
|
| elif suffix == ".webp":
|
| kwargs.update(quality=100, lossless=bool(args.webp_lossless))
|
| elif suffix in (".tif", ".tiff"):
|
| kwargs.update(compression="tiff_lzw")
|
|
|
| out_img.save(args.output, **kwargs)
|
| print(f"Saved: {args.output}")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
