| """ | |
| convert_to_safetensors.py - model.png -> model.safetensors. | |
| model.png stays the canonical model; this exports the same weights in the | |
| standard format, with parameter-count metadata embedded in the safetensors | |
| header so the count is verifiable without running any code. | |
| Usage: | |
| python convert_to_safetensors.py | |
| python convert_to_safetensors.py --model model.png --out model.safetensors | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| from safetensors.torch import save_file | |
| from model import MODEL_H, MODEL_W, N_PARAMS, NATIVE_RES, PARAM_SPECS, load_model | |
| def main(): | |
| p = argparse.ArgumentParser() | |
| p.add_argument("--model", default="model.png") | |
| p.add_argument("--out", default="model.safetensors") | |
| args = p.parse_args() | |
| weights = load_model(args.model) | |
| counts = {name: int(weights[name].numel()) for name, _ in PARAM_SPECS} | |
| assert sum(counts.values()) == N_PARAMS | |
| metadata = { | |
| "model_type": "pixelmodel-v2", | |
| "total_parameters": str(N_PARAMS), | |
| "param_breakdown": json.dumps(counts), | |
| "text_encoder_parameters": "0", | |
| "vae_parameters": "0", | |
| "has_bias": "true", | |
| "native_resolution": f"{NATIVE_RES}x{NATIVE_RES}", | |
| "source_png": f"{MODEL_W}x{MODEL_H} px, 16-bit codec (R=high byte, G=low byte)", | |
| } | |
| save_file(weights, args.out, metadata=metadata) | |
| print(f"{args.model} -> {args.out} ({os.path.getsize(args.out)} bytes, " | |
| f"{N_PARAMS} parameters)") | |
| if __name__ == "__main__": | |
| main() | |