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