| |
| |
| |
| |
| |
| |
| |
| |
| """Core AI export vs the original PyTorch weights (float32), tile by tile, on stock photos.""" |
| import asyncio, sys, time |
| import numpy as np, torch |
| from PIL import Image |
| from spandrel import ModelLoader |
| import coreai.runtime as rt |
|
|
| ckpt, path, tile = sys.argv[1], sys.argv[2], int(sys.argv[3]) |
| images = sys.argv[4:] |
| desc = ModelLoader().load_from_file(ckpt) |
| net = desc.model.eval().float() |
| luma = torch.tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1) |
| def reference(x): |
| if desc.input_channels == 1: x = (x * luma).sum(dim=1, keepdim=True) |
| y = net(x) |
| if desc.output_channels == 1: y = y.repeat(1, 3, 1, 1) |
| return y.clamp(0, 1) |
| def psnr(a, b): |
| mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) |
| return 99.0 if mse == 0 else 10 * np.log10(1.0 / mse) |
| async def main(): |
| rng = np.random.default_rng(0); tiles = [] |
| for p in images: |
| im = np.asarray(Image.open(p).convert("RGB")).astype(np.float32) / 255 |
| h, w, _ = im.shape |
| for _ in range(3): |
| y = int(rng.integers(0, max(1, h - tile))); x = int(rng.integers(0, max(1, w - tile))) |
| c = im[y:y+tile, x:x+tile] |
| if c.shape[:2] == (tile, tile): tiles.append(c) |
| m = await rt.AIModel.load(path) |
| fn = m.load_function("main") |
| scores, times = [], [] |
| for c in tiles: |
| x = torch.from_numpy(c).permute(2, 0, 1)[None].contiguous() |
| with torch.no_grad(): ref = reference(x)[0].permute(1, 2, 0).numpy() |
| t = time.time() |
| out = await fn({"x": rt.NDArray(data=x.half().contiguous())}) |
| times.append(time.time() - t) |
| got = out["output"].numpy().astype(np.float32)[0].transpose(1, 2, 0) |
| scores.append(psnr(ref, got)) |
| print(f"{path.split('/')[-1]:34} scale x{desc.scale} PSNR vs PyTorch fp32 over {len(scores)} tiles: " |
| f"min {min(scores):.1f} dB, mean {np.mean(scores):.1f} dB; {np.median(times)*1000:.0f} ms/tile") |
| asyncio.run(main()) |
|
|
