| import torch | |
| from safetensors.torch import load_file, save_file | |
| # Load both sharded safetensors | |
| shard_1 = load_file("diffusion_pytorch_model-00001-of-00002.safetensors") | |
| shard_2 = load_file("diffusion_pytorch_model-00002-of-00002.safetensors") | |
| # Combine both shards into a single dictionary | |
| combined_model = {**shard_1, **shard_2} | |
| # Save the combined model as output.safetensors | |
| save_file(combined_model, "flowgram.safetensors") | |
| print("Successfully merged the files into output.safetensors") |