File size: 5,671 Bytes
40adac6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | #!/usr/bin/env python3
"""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()
|