"""
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()