import torch | |
from safetensors.torch import load_file, save_file | |
# Load all three models | |
model_1 = load_file("diffusion_pytorch_model-00001-of-00003.safetensors") | |
model_2 = load_file("diffusion_pytorch_model-00002-of-00003.safetensors") | |
model_3 = load_file("diffusion_pytorch_model-00003-of-00003.safetensors") | |
# Combine the models into a single dictionary | |
combined_model = {**model_1, **model_2, **model_3} | |
# Save the combined model as flowgram.safetensors | |
save_file(combined_model, "midflow_high.safetensors") | |
print("Successfully merged the models into flowgram.safetensors") |