#!/usr/bin/env python
"""
Read a checkpoint/safetensors file and compare it to a template .json.

Returns True if their metadata match.
"""

import sys
import argparse
import json

from pathlib import Path

from invokeai.backend.model_management.models.base import read_checkpoint_meta

parser = argparse.ArgumentParser(description="Compare a checkpoint/safetensors file to a JSON metadata template.")
parser.add_argument("--checkpoint", "--in", type=Path, help="Path to the input checkpoint/safetensors file")
parser.add_argument("--template", "--out", type=Path, help="Path to the template .json file to match against")

opt = parser.parse_args()
ckpt = read_checkpoint_meta(opt.checkpoint)
while "state_dict" in ckpt:
    ckpt = ckpt["state_dict"]

checkpoint_metadata = {}

for key, tensor in ckpt.items():
    checkpoint_metadata[key] = list(tensor.shape)

with open(opt.template, "r", encoding="utf-8") as f:
    template = json.load(f)

for key in template["template"]:
    val1 = checkpoint_metadata.get(key)
    val2 = template["template"][key]
    if val1 != val2:
        print(f"mismatch: {key}: template={val2} != checkpoint={val1}")
        sys.exit(-1)
print("Match")
sys.exit(0)