| import json |
| import math |
| import os |
| import shutil |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from safetensors.torch import load_file |
| from transformers import LlamaConfig, LlamaForCausalLM |
| from huggingface_hub import hf_hub_download, HfApi |
|
|
| def build_model(c): |
| return LlamaForCausalLM(LlamaConfig(**{k: c[k] for k in [ |
| "vocab_size", "hidden_size", "intermediate_size", "num_hidden_layers", |
| "num_attention_heads", "num_key_value_heads", "max_position_embeddings", |
| "rms_norm_eps", "tie_word_embeddings", |
| ]})) |
|
|
| def encode(weights, png_path, cfg_path, config="config.json"): |
| c = json.load(open(config)) |
| model = build_model(c) |
| model.load_state_dict(load_file(weights), strict=False) |
| 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) |
| json.dump({"cfg": c, "params": manifest, "total_parameters": total, |
| "side": side, "dtype": "float16", "channels": "R=hi,G=lo,B=unused"}, |
| open(cfg_path, "w")) |
| mb = os.path.getsize(png_path) / 1e6 |
| print(f"[png] encoded {total:,} params -> {side}x{side} PNG ({mb:.1f} MB)") |
| return total, side |
|
|
| def load_model_png(png_path, cfg_path, device="cpu"): |
| meta = json.load(open(cfg_path)) |
| model = build_model(meta["cfg"]) |
| 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) |
| sd[m["name"]].copy_(torch.from_numpy(chunk.copy()).view(*m["shape"]).to(torch.float32)) |
| off += n |
| return model.to(device).eval() |
|
|
| def verify(weights, png_path, cfg_path, config="config.json"): |
| model = load_model_png(png_path, cfg_path) |
| ref = load_file(weights) |
| worst, name = 0.0, "" |
| for k, p in model.named_parameters(): |
| d = (p.detach() - ref[k].float()).abs().max().item() |
| if d > worst: |
| worst, name = d, k |
| rel = worst / max(1e-12, ref[name].float().abs().max().item()) |
| print(f"[png] round-trip max abs err {worst:.3e} on {name} (relative {rel:.2e})") |
| return worst |
|
|
| repo = "TobiasLogic/textmodel-gci-scratch-ckpt" |
| rev = "step50000" |
| cfg_path = hf_hub_download(repo_id=repo, revision=rev, filename="config.json") |
| w_path = hf_hub_download(repo_id=repo, revision=rev, filename="model.safetensors") |
| shutil.copy(cfg_path, "config.json") |
| shutil.copy(w_path, "model.safetensors") |
|
|
| total, side = encode("model.safetensors", "model.png", "model_png.json", "config.json") |
| worst = verify("model.safetensors", "model.png", "model_png.json", "config.json") |
| print("total params", total, "side", side, "worst abs err", worst) |
|
|
| api = HfApi() |
| api.upload_file(path_or_fileobj="model.png", path_in_repo="model.png", repo_id="TobiasLogic/TextModel-v1") |
| api.upload_file(path_or_fileobj="model_png.json", path_in_repo="model_png.json", repo_id="TobiasLogic/TextModel-v1") |
| print("uploaded model.png + model_png.json") |
|
|
