| from __future__ import annotations |
| import argparse, json, math |
| import numpy as np |
| import torch |
| from PIL import Image |
| from dit_v6 import MMDiT |
|
|
| def encode(ckpt, png_path, cfg_path, which="ema"): |
| ck = torch.load(ckpt, map_location="cpu") |
| c = ck["cfg"] |
| model = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"], t5_len=c["t5_len"]) |
| 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 = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"], t5_len=c["t5_len"]) |
| 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 -> MMDiT with {sum(p.numel() for p in model.parameters()):,} params", flush=True) |
| if ckpt: |
| ck = torch.load(ckpt, map_location="cpu") |
| c = ck["cfg"] |
| ref = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"], t5_len=c["t5_len"]) |
| 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) |
|
|
