from safetensors.torch import load_file
import json

weights = load_file("diffusion_pytorch_model.safetensors")
weight_map = {k: "diffusion_pytorch_model.safetensors" for k in weights.keys()}

index = {
    "metadata": {"total_size": sum(v.numel() * v.element_size() for v in weights.values())},
    "weight_map": weight_map
}

with open("diffusion_pytorch_model.safetensors.index.json", "w") as f:
    json.dump(index, f, indent=2)