File size: 3,426 Bytes
6c9c825 | 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 | from __future__ import annotations
import argparse, json, math
import numpy as np
import torch
from PIL import Image
from dit import DiT
def encode(ckpt, png_path, cfg_path, which="ema"):
ck = torch.load(ckpt, map_location="cpu")
c = ck["cfg"]
model = DiT(dim=c["dim"], depth=c["depth"], heads=c["heads"])
model.load_state_dict(ck[which])
parts, manifest = [], []
for name, p in model.named_parameters():
a = p.detach().to(torch.float16).contiguous().view(-1).numpy()
parts.append(a)
manifest.append({"name": name, "shape": list(p.shape), "numel": int(a.size)})
flat = np.concatenate(parts)
N = flat.size
side = math.ceil(math.sqrt(N))
u16 = flat.view(np.uint16)
img = np.zeros((side * side, 3), dtype=np.uint8)
img[:N, 0] = (u16 >> 8).astype(np.uint8)
img[:N, 1] = (u16 & 0xFF).astype(np.uint8)
Image.fromarray(img.reshape(side, side, 3), "RGB").save(png_path)
total = sum(m["numel"] for m in manifest)
with open(cfg_path, "w") as f:
json.dump({"cfg": c, "params": manifest, "total_parameters": total,
"side": side, "dtype": "float16", "channels": "R=hi,G=lo,B=unused"}, f)
import os
mb = os.path.getsize(png_path) / 1e6
print(f"[png] encoded {total:,} params -> {side}x{side} PNG ({mb:.1f} MB)", flush=True)
return total, side
def load_model_png(png_path, cfg_path, device="cpu"):
meta = json.load(open(cfg_path))
c = meta["cfg"]
model = DiT(dim=c["dim"], depth=c["depth"], heads=c["heads"])
arr = np.asarray(Image.open(png_path).convert("RGB")).reshape(-1, 3)
total = meta["total_parameters"]
hi = arr[:total, 0].astype(np.uint16)
lo = arr[:total, 1].astype(np.uint16)
flat = ((hi << 8) | lo).astype(np.uint16).view(np.float16)
sd = dict(model.named_parameters())
off = 0
with torch.no_grad():
for m in meta["params"]:
n = m["numel"]
chunk = flat[off:off + n].astype(np.float16)
t = torch.from_numpy(chunk.copy()).view(*m["shape"]).to(torch.float32)
sd[m["name"]].copy_(t)
off += n
return model.to(device).eval()
def decode_and_verify(png_path, cfg_path, ckpt=None, which="ema"):
model = load_model_png(png_path, cfg_path)
print(f"[png] decoded -> DiT with {sum(p.numel() for p in model.parameters()):,} params", flush=True)
if ckpt:
ck = torch.load(ckpt, map_location="cpu")
ref = DiT(dim=ck["cfg"]["dim"], depth=ck["cfg"]["depth"], heads=ck["cfg"]["heads"])
ref.load_state_dict(ck[which])
maxdiff = 0.0
for (n1, p1), (n2, p2) in zip(model.named_parameters(), ref.named_parameters()):
maxdiff = max(maxdiff, (p1.float() - p2.half().float()).abs().max().item())
print(f"[png] max |decoded - original(fp16)| = {maxdiff:.2e} (0 == lossless)", flush=True)
return model
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("mode", choices=["encode", "decode"])
ap.add_argument("--ckpt", default="ckpt/final.pt")
ap.add_argument("--png", default="model.png")
ap.add_argument("--config", default="model_png.json")
ap.add_argument("--which", default="ema")
args = ap.parse_args()
if args.mode == "encode":
encode(args.ckpt, args.png, args.config, args.which)
else:
decode_and_verify(args.png, args.config, args.ckpt, args.which)
|