File size: 2,214 Bytes
57e8f83 | 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 | # /// script
# requires-python = ">=3.11"
# dependencies = ["coreai-core==1.0.0b3", "coreai-torch==0.4.3", "spandrel==0.4.1", "numpy", "pillow"]
# [tool.uv]
# index-url = "https://pypi.org/simple"
# prerelease = "allow"
# index-strategy = "unsafe-best-match"
# ///
"""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())
|