How merge/fusion several files .safetensors ( diffusers format ) to a whole .safetensors checkpoint ?

[BenDes21]

Hi there, I got diffusion_pytorch_model-00001-of-00003.safetensors, diffusion_pytorch_model-00002-of-00003.safetensors and diffusion_pytorch_model-00003-of-00003.safetensors after full finetuned Flux with Ostris AI tool kit, I guess at this point I will have to fusion/merge theses 3 parts for created the whole checkpoint usable for a1111 Forge . Is anyone know a script or a tool who can allow this last step ?

Thanks a lot

[bmaxdk]

To merge the split .safetensors files into a single chckpt, you can usesafetensors library in huggingface.
pip install safetensors

import safetensors.torch

merge_state_dict ={}
files = ["name1.safetensors", "name2.safetensors"...] #file you want to be merged
merged_file = "merged_file.safetensors"
for file in files:
    load_files_dict = safetensors.torch.load_file(your.safetensors files)
merge_state_dict.update(load_files_dict)

# save
safetensors.torch.save_file(merge_state_dict, merged_file)

this merged_file will be now saved to the combined state dictionary with .safetensors file

[BenDes21]

To merge the split .safetensors files into a single chckpt, you can usesafetensors library in huggingface. pip install safetensors

import safetensors.torch

merge_state_dict ={}
files = ["name1.safetensors", "name2.safetensors"...] #file you want to be merged
merged_file = "merged_file.safetensors"
for file in files:
    load_files_dict = safetensors.torch.load_file(your.safetensors files)
merge_state_dict.update(load_files_dict)

# save
safetensors.torch.save_file(merge_state_dict, merged_file)

this merged_file will be now saved to the combined state dictionary with .safetensors file

Hi there, unfortunately the .safetensors merged cannot be read by a1111, got : ValueError: Failed to recognize model type!

[bmaxdk]

It seem you are having merging issue/. can you make sure your file loaded correct and the your .safetensors are not corrupted.

let's try:

from safetensors.torch import load_file, save_file 
from transformers import AutoTokenizer, AutoModel 
from diffusers import StableDiffusionPipeline 
import torch
load1 = load_file("diffusion_pytorch_model-00001-of-00003.safetensors") 
load2 = load_file("diffusion_pytorch_model-00002-of-00003.safetensors") 
load3 = load_file("diffusion_pytorch_model-00003-of-00003.safetensors") 
# Debug 
print(load1.keys())
print(load2.keys())
print(load3.keys())
# If here you got and issue, you may facing file corrupted issue..

# Once everything looks right lets unpack this 
merged_state_dict = {**load1, **load2, **load3}


save_file(merged_state_dict, "merged_diffusion_torch.safetensors")


your_model_name_path = "merged_diffusion_torch.safetensors" 

 # Now, let's try loading the model to ensure it's valid 
merged_load = load_file(your_model_name_path)
print(merged_load.keys())
try: # Load the tokenizer and model 
    tokenizer = AutoTokenizer.from_pretrained(your_model_name_path) 
    model = AutoModel.from_pretrained(your_model_name_path, from_tf=False, from_safetensors=True)
    
    pipeline = StableDiffusionPipeline.from_pretrained(your_model_name_path, torch_dtype=torch.float16) 
    print("it is successfully loaded!") 
    
except Exception as e: 
    print(f"Error: {e}")

I got issue having a111 when my model architecture mismatch.. I am not sure your case but it's possible when file corrupted it could cuase the issue.

[NotTheStallion]

While finetuning an LLM on a local cluster, I needed to reshard safetensor files, so I created a code that merges them then resplits them into N parts.

This might prove useful to you if you're still looking for a solution, or for other people who might refer to this page when facing the same problem.

Here is the code i made for it : https://github.com/NotTheStallion/reshard-safetensors