90 lines
3.7 KiB
Python
90 lines
3.7 KiB
Python
|
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")
|