"""Symbioza example job: upscale. Upscales a photograph 2x with Real-ESRGAN (RealESRGAN_x2plus). The photo is a public-domain rocket launch from scikit-image's sample data. Every input is downloaded from a pinned URL (3 attempts) and refused unless its sha256 matches. The network is BasicSR v1.4.2's RRDBNet, taken unmodified from its two sha256-verified source files: only the classes and functions it needs are compiled, without the registry decorator, so basicsr is not pip-installed. Inference follows Real-ESRGAN's own x2 path: reflection-pad to a multiple of 2, run, crop, clamp, round. numpy and Pillow are pip-installed (pinned) only if the image lacks them. On any failure it writes error.json; SHA256SUMS is written on every exit path. CPU_OK=1 lets it run without CUDA, for local testing only; the published spec never sets it. """ import ast import hashlib import importlib import json import math import os import re import shutil import signal import subprocess import sys import time import traceback import urllib.request import warnings warnings.filterwarnings("ignore", message=".*cuda capability") # torch names the card in this warning import torch # noqa: E402 import torch.nn.functional as F # noqa: E402 OUT = os.environ.get("ARTIFACTS_DIR", "/workspace/artifacts") WORK = os.environ.get("WORK_DIR", "/tmp/upscale-inputs") BASICSR = "https://raw.githubusercontent.com/XPixelGroup/BasicSR/v1.4.2/" WEIGHTS_LICENSE = ("released as an asset of the BSD-3-Clause Real-ESRGAN repository; " "the weights file carries no separate license") INPUTS = { # name: (pinned url, sha256, license) "rocket.jpg": ( "https://raw.githubusercontent.com/scikit-image/scikit-image/v0.24.0/skimage/data/rocket.jpg", "c2dd0de7c538df8d111e479619b129464d0269d0ae5fd18ca91d33a7fdfea95c", "public domain: captured by SpaceX and released in the public domain " "(scikit-image v0.24.0, skimage/data/_fetchers.py, rocket())"), "RealESRGAN_x2plus.pth": ( "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth", "49fafd45f8fd7aa8d31ab2a22d14d91b536c34494a5cfe31eb5d89c2fa266abb", WEIGHTS_LICENSE), "rrdbnet_arch.py": (BASICSR + "basicsr/archs/rrdbnet_arch.py", "6aa8eb3f77e918b90eeff4ebfaf8745abe0c40eb8fcc03abbc1928693d70ab58", "Apache-2.0 (BasicSR)"), "arch_util.py": (BASICSR + "basicsr/archs/arch_util.py", "45b322b874057c4bff382cc8f1b64554857b9aa4e0c47e73fe67e4afc0b28ac1", "Apache-2.0 (BasicSR)"), "LICENSE-BasicSR.txt": (BASICSR + "LICENSE.txt", "5447872f21ddc29dc0dc060d5c885513c01c904e9da9931bba69898a6fb4d415", "license text"), "LICENSE-Real-ESRGAN.txt": ("https://raw.githubusercontent.com/xinntao/Real-ESRGAN/v0.2.1/LICENSE", "4a699ec4863d96a91fc265948a0c90033f7e8735d515524dcf3444736406e0c2", "license text"), } CROP = (192, 85, 256) # x, y, size of the sanity-check crop MIN_MODEL_PSNR_DB = 20.0 # a wrong network or wrong weights lands far below this STAGE = "start" np = Image = None # imported in main(), after the dependency check class Fail(Exception): """A failure the job names itself: an exit code and the reason.""" def __init__(self, code, msg): super().__init__(msg) self.code = code def stage(name): global STAGE STAGE = name print(f"== {name}", flush=True) def sha256(path): h = hashlib.sha256() with open(path, "rb") as fh: for block in iter(lambda: fh.read(1 << 20), b""): h.update(block) return h.hexdigest() def write_sha256sums(): with open(os.path.join(OUT, "SHA256SUMS"), "w") as fh: for name in sorted(os.listdir(OUT)): path = os.path.join(OUT, name) if name != "SHA256SUMS" and os.path.isfile(path): fh.write(f"{sha256(path)} {name}\n") def redact(text): return re.sub(r"\d+(\.\d+)? ?[KMGT]i?B\b", "", text) # an out-of-memory message states the card's memory def run(main, job): """Run main(). On ANY failure leave error.json; write SHA256SUMS on every exit path.""" signal.signal(signal.SIGTERM, lambda *_: sys.exit(143)) code, err = 0, None try: main() except Fail as e: code, err = e.code, str(e) except SystemExit as e: code = e.code if isinstance(e.code, int) else 1 err = "stopped by SIGTERM" if code == 143 else f"exit {e.code}" except BaseException as e: code, err = 1, f"{type(e).__name__}: {e}" print(redact(traceback.format_exc()), flush=True) os.makedirs(OUT, exist_ok=True) for name in os.listdir(OUT): if name.endswith(".part"): # a save the failure interrupted os.remove(os.path.join(OUT, name)) if code: err = redact(err)[:2000] print(f"{job} FAILED during '{STAGE}' (exit {code}): {err}", flush=True) with open(os.path.join(OUT, "error.json"), "w") as fh: json.dump({"ok": False, "job": job, "stage": STAGE, "error": err, "exit_code": code}, fh, indent=2) write_sha256sums() sys.exit(code) def fetch(name, url, digest): """Download into WORK (3 attempts, 60 s timeout each); refuse bytes whose sha256 is not the pinned one.""" path, err = os.path.join(WORK, name), None for attempt in range(3): time.sleep(10 * attempt) try: with urllib.request.urlopen(url, timeout=60) as r, open(path, "wb") as fh: shutil.copyfileobj(r, fh) except Exception as e: err = f"{type(e).__name__}: {e}" continue got = sha256(path) if got == digest: return path err = f"sha256 {got} is not the pinned {digest}" raise Fail(2, f"{name} from {url}: {err} (3 attempts)") def need(module, pin): """Import a module; pip-install the pinned version first only if the image lacks it.""" try: return importlib.import_module(module) except ImportError: stage(f"pip install {pin}") subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir", pin], check=True) importlib.invalidate_caches() return importlib.import_module(module) def build_rrdbnet(arch_util_src, rrdbnet_src): keep = {"default_init_weights", "make_layer", "pixel_unshuffle", "ResidualDenseBlock", "RRDB", "RRDBNet"} body = [] for src in (arch_util_src, rrdbnet_src): for node in ast.parse(src).body: if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in keep: node.decorator_list = [d for d in node.decorator_list if "ARCH_REGISTRY" not in ast.dump(d)] body.append(node) ns = {"torch": torch, "nn": torch.nn, "F": F, "init": torch.nn.init, "_BatchNorm": torch.nn.modules.batchnorm._BatchNorm} exec(compile(ast.Module(body=body, type_ignores=[]), "basicsr_v1.4.2_rrdbnet", "exec"), ns) missing = keep - ns.keys() if missing: raise Fail(2, f"BasicSR source is missing {sorted(missing)}") return ns["RRDBNet"] @torch.no_grad() def upscale(model, rgb, dev): """HxWx3 uint8 RGB -> 2Hx2Wx3 uint8, Real-ESRGAN's x2 pre/post-processing.""" x = torch.from_numpy(np.array(rgb)).permute(2, 0, 1).float().div(255).unsqueeze(0).to(dev) h, w = x.shape[2:] ph, pw = h % 2, w % 2 # pad to a multiple of 2 by reflection, as RealESRGANer does for scale 2 if ph or pw: x = F.pad(x, (0, pw, 0, ph), "reflect") y = model(x)[:, :, :2 * h, :2 * w].clamp_(0, 1) return (y[0].permute(1, 2, 0).cpu().numpy() * 255.0).round().astype(np.uint8) def resize(rgb, w, h, method): return np.asarray(Image.fromarray(rgb).resize((w, h), method)) def psnr(a, b): mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2) return round(10 * math.log10(255.0 ** 2 / mse), 2) def main(): global np, Image stage("check CUDA") cuda = torch.cuda.is_available() if not cuda and os.environ.get("CPU_OK") != "1": raise Fail(3, "torch.cuda.is_available() is False; this job requires a GPU") dev = torch.device("cuda" if cuda else "cpu") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True stage("python dependencies") np = need("numpy", "numpy==1.26.4") PIL = need("PIL", "pillow==10.4.0") from PIL import Image as PILImage Image = PILImage os.makedirs(OUT, exist_ok=True) os.makedirs(WORK, exist_ok=True) stage("download and verify inputs") paths = {n: fetch(n, url, digest) for n, (url, digest, _) in INPUTS.items()} stage("build the network and load the weights") RRDBNet = build_rrdbnet(open(paths["arch_util.py"]).read(), open(paths["rrdbnet_arch.py"]).read()) model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) weights = torch.load(paths["RealESRGAN_x2plus.pth"], map_location="cpu", weights_only=True) key = "params_ema" if "params_ema" in weights else "params" model.load_state_dict(weights[key], strict=True) model.eval().to(dev) stage("upscale the photo") photo = np.asarray(Image.open(paths["rocket.jpg"]).convert("RGB")) h, w = photo.shape[:2] up = upscale(model, photo, dev) shutil.copyfile(paths["rocket.jpg"], os.path.join(OUT, "original.jpg")) Image.fromarray(up).save(os.path.join(OUT, "upscaled-x2.png")) nearest = resize(photo, 2 * w, 2 * h, Image.NEAREST) Image.fromarray(np.concatenate([nearest, up], axis=1)).save(os.path.join(OUT, "comparison.png")) stage("sanity check on a crop with a known answer") x0, y0, s = CROP hr = photo[y0:y0 + s, x0:x0 + s] lr = resize(hr, s // 2, s // 2, Image.BICUBIC) bic = resize(lr, s, s, Image.BICUBIC) sr = upscale(model, lr, dev) strip = [resize(lr, s, s, Image.NEAREST), bic, sr, hr] Image.fromarray(np.concatenate(strip, axis=1)).save(os.path.join(OUT, "crop-comparison.png")) scores = {"bicubic": psnr(bic, hr), "model": psnr(sr, hr)} ok = scores["model"] >= MIN_MODEL_PSNR_DB stage("write report") for n in ("LICENSE-BasicSR.txt", "LICENSE-Real-ESRGAN.txt"): shutil.copyfile(paths[n], os.path.join(OUT, n)) report = { "ok": ok, "task": "2x super-resolution", "model": {"weights": "RealESRGAN_x2plus", "weights_license": WEIGHTS_LICENSE, "arch": "BasicSR v1.4.2 RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, " "num_grow_ch=32, scale=2)", "arch_license": "Apache-2.0", "state_dict_key": key, "strict_load": True, "parameters": sum(p.numel() for p in model.parameters())}, "inputs": {n: {"url": u, "sha256": d, "license": lic, "sha256_verified": True} for n, (u, d, lic) in INPUTS.items()}, "input_size": [w, h], "output_size": [up.shape[1], up.shape[0]], "sanity_check": { "what": "a crop is shrunk 2x (bicubic) and enlarged back by bicubic and by the model, then " "scored by PSNR against the real crop; the job fails if the model scores under " "min_model_psnr_db, which means a broken network or broken weights. Real-ESRGAN aims at " "perceived sharpness, not PSNR, so bicubic can score higher.", "crop_xywh": [x0, y0, s, s], "psnr_db": scores, "min_model_psnr_db": MIN_MODEL_PSNR_DB}, "device": dev.type, "determinism": {"tf32": False, "cudnn_deterministic": True}, "versions": {"torch": torch.__version__, "numpy": np.__version__, "pillow": PIL.__version__}, } with open(os.path.join(OUT, "report.json"), "w") as fh: json.dump(report, fh, indent=2) md = [ "# 2x upscale", "", f"- `original.jpg`: the input, {w}x{h}, a public-domain photo.", f"- `upscaled-x2.png`: the model's 2x output, {up.shape[1]}x{up.shape[0]}.", "- `comparison.png`: left, the original enlarged 2x by repeating pixels; right, the model's 2x.", "", f"Model: Real-ESRGAN `RealESRGAN_x2plus`; the weights are {WEIGHTS_LICENSE}. " "Network code: BasicSR v1.4.2 `RRDBNet` (Apache-2.0).", "", "## Inputs (sha256 verified before use)", "", "| file | license | sha256 |", "|---|---|---|", *[f"| [{n}]({u}) | {lic} | `{d}` |" for n, (u, d, lic) in INPUTS.items()], "", "License texts: `LICENSE-Real-ESRGAN.txt`, `LICENSE-BasicSR.txt`. " "A sanity check on a crop with a known answer is in `report.json` and `crop-comparison.png`.", ] with open(os.path.join(OUT, "report.md"), "w") as fh: fh.write("\n".join(md) + "\n") print(json.dumps(report, indent=2), flush=True) if not ok: raise Fail(4, f"sanity check failed: model PSNR {scores['model']} dB is below {MIN_MODEL_PSNR_DB} dB") if __name__ == "__main__": run(main, "upscale")