#!/usr/bin/env python3 """단일 업스케일 실행 — spandrel 기반. Mac MPS / Spark CUDA 양쪽 공통. Usage: upscale-run.py \ --model PATH \ --input PATH \ --output PATH \ --target-size 3000 \ [--device auto|mps|cuda|cpu] \ [--dtype float32|float16] \ [--tile 512] # tile=0 은 whole-image 시도 Output (stdout, 마지막 줄 JSON): {"ok": true, "wall_s": 12.3, "load_s": 1.2, "infer_s": 10.1, "save_s": 0.3, "peak_mem_gb": 4.2, "input_wh": [1024,1024], "output_wh": [3000,3000], "scale_native": 4, "sha256": "..."} 실패: {"ok": false, "error": "...", "phase": "load|infer|save"} """ import argparse import hashlib import json import os import sys import time import traceback from pathlib import Path try: import torch from PIL import Image import numpy as np from spandrel import ModelLoader, ImageModelDescriptor except ImportError as e: print(json.dumps({"ok": False, "error": f"import: {e}", "phase": "import"})) sys.exit(1) def pick_device(arg: str) -> str: if arg != "auto": return arg if torch.cuda.is_available(): return "cuda" if torch.backends.mps.is_available(): return "mps" return "cpu" def peak_mem_gb(device: str) -> float | None: try: if device == "cuda": return torch.cuda.max_memory_allocated() / (1024**3) if device == "mps": return torch.mps.driver_allocated_memory() / (1024**3) except Exception: return None return None def reset_mem(device: str) -> None: try: if device == "cuda": torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() elif device == "mps": torch.mps.empty_cache() except Exception: pass def synchronize(device: str) -> None: if device == "cuda": torch.cuda.synchronize() elif device == "mps": torch.mps.synchronize() def load_image(path: Path) -> torch.Tensor: """HWC uint8 RGB PIL → CHW float32 [0,1] tensor (no batch).""" img = Image.open(path).convert("RGB") arr = np.asarray(img, dtype=np.float32) / 255.0 # H W 3 return torch.from_numpy(arr).permute(2, 0, 1).contiguous() # 3 H W def save_image(t: torch.Tensor, path: Path) -> None: """CHW [0,1] float → PNG.""" arr = t.clamp(0, 1).mul(255).round().to(torch.uint8).permute(1, 2, 0).cpu().numpy() Image.fromarray(arr).save(path, format="PNG", optimize=False) def tile_forward( model, x: torch.Tensor, scale: int, tile: int, overlap: int = 16 ) -> torch.Tensor: """Tile-based inference for large inputs. x: 1 C H W.""" if tile <= 0: return model(x) b, c, h, w = x.shape out_h, out_w = h * scale, w * scale out = torch.zeros((b, c, out_h, out_w), dtype=x.dtype, device=x.device) weight = torch.zeros((b, 1, out_h, out_w), dtype=x.dtype, device=x.device) for y0 in range(0, h, tile - overlap): for x0 in range(0, w, tile - overlap): y1 = min(y0 + tile, h) x1 = min(x0 + tile, w) y0c = max(0, y1 - tile) x0c = max(0, x1 - tile) patch = x[:, :, y0c:y1, x0c:x1] with torch.inference_mode(): p_out = model(patch) oy0, oy1 = y0c * scale, y1 * scale ox0, ox1 = x0c * scale, x1 * scale out[:, :, oy0:oy1, ox0:ox1] += p_out weight[:, :, oy0:oy1, ox0:ox1] += 1.0 return out / weight.clamp(min=1.0) def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--model", required=True) ap.add_argument("--input", required=True) ap.add_argument("--output", required=True) ap.add_argument("--target-size", type=int, required=True, help="최종 출력 한 변 (정사각 가정). 모델 native 배율 후 Lanczos 다운샘플.") ap.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"]) ap.add_argument("--dtype", default="float32", choices=["float32", "float16"]) ap.add_argument("--tile", type=int, default=0, help="tile 사이즈 (px). 0 = whole-image. OOM 나면 512/256 시도.") args = ap.parse_args() # ~/ expansion — shell 더블쿼트에선 expand 안 되므로 여기서 처리 args.model = os.path.expanduser(args.model) args.input = os.path.expanduser(args.input) args.output = os.path.expanduser(args.output) t0 = time.perf_counter() result: dict = {"ok": False} try: device = pick_device(args.device) dtype = torch.float16 if args.dtype == "float16" else torch.float32 reset_mem(device) # ---- Load ---- t_load0 = time.perf_counter() descriptor: ImageModelDescriptor = ModelLoader().load_from_file(args.model) model = descriptor.model.eval() if device != "cpu": model = model.to(device) if dtype == torch.float16 and device != "mps": # MPS fp16 often hurts; keep fp32 there model = model.to(dtype) native_scale = int(descriptor.scale) synchronize(device) t_load = time.perf_counter() - t_load0 # ---- Input ---- x = load_image(Path(args.input)).unsqueeze(0).to(device) if dtype == torch.float16 and device != "mps": x = x.to(dtype) ih, iw = x.shape[-2], x.shape[-1] # ---- Infer ---- synchronize(device) t_inf0 = time.perf_counter() with torch.inference_mode(): y = tile_forward(model, x, native_scale, args.tile) synchronize(device) t_inf = time.perf_counter() - t_inf0 # ---- Resize to target ---- y_full = y.squeeze(0).float().cpu() # 3 H W, [0,1] native_wh = (iw * native_scale, ih * native_scale) if args.target_size != native_wh[0] or args.target_size != native_wh[1]: # Lanczos downsample via PIL pil = Image.fromarray( y_full.clamp(0, 1).mul(255).round().to(torch.uint8) .permute(1, 2, 0).numpy() ) pil = pil.resize((args.target_size, args.target_size), Image.Resampling.LANCZOS) pil.save(args.output, format="PNG", optimize=False) else: t_save0 = time.perf_counter() save_image(y_full, Path(args.output)) t_save = time.perf_counter() - (t_save0 if "t_save0" in dir() else t_inf0 + t_inf) # ---- Hash ---- h = hashlib.sha256(Path(args.output).read_bytes()).hexdigest()[:16] wall = time.perf_counter() - t0 result = { "ok": True, "wall_s": round(wall, 3), "load_s": round(t_load, 3), "infer_s": round(t_inf, 3), "peak_mem_gb": round(peak_mem_gb(device) or 0, 3), "input_wh": [iw, ih], "output_wh": [args.target_size, args.target_size], "scale_native": native_scale, "device": device, "dtype": args.dtype, "tile": args.tile, "sha256_16": h, } except Exception as e: result = { "ok": False, "error": f"{type(e).__name__}: {e}"[:300], "phase": "infer", "traceback": traceback.format_exc()[-800:], } print(json.dumps(result, ensure_ascii=False)) return 0 if result.get("ok") else 2 if __name__ == "__main__": sys.exit(main())