Flux Merged diffusion .safetensors files not readable

[BenDes21]

Hi, I successfully full finetuned flux with ostris ai toolkit and I got theses 3 files at the end of the training ( diffusion model files ) :

diffusion_pytorch_model-00001-of-00003.safetensors diffusion_pytorch_model-00002-of-00003.safetensors diffusion_pytorch_model-00003-of-00003.safetensors

Im not sure how theses 3 " shards " should be used, I tried the nodes Diffusers loader / Combined Diffusers loader for load theses 3 files and save them into a checkpoint ( for be readable by a1111 ) but unfortunately comfyUI dont detect theses files as diffusers but as diffusion_models ( but they're is not nodes for combine them )

I also try the script :

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) 

It's merged the 3 parts .safetensors but the ouput .safetensors is not readable by comfyUI / a1111

Anyone have a solution be read theses 3 files or fusion them and make them usable into comfyUI/a1111 ? Thanks a lot

[bash-j]

The output is a diffusers model.

Try this code. You need to point to the location of the transformers folder which contains the shards. Then change the output file path to whatever you want it to be named. The conversion to fp8 is optional, but it's what most people use.

import torch
from diffusers import FluxTransformer2DModel
from safetensors.torch import save_file

dtype = torch.bfloat16

transformer = FluxTransformer2DModel.from_pretrained("//path_to/transformers", torch_dtype=dtype)
state_dict = transformer.state_dict()

st = dict()

for key in state_dict:
    st[key] = state_dict[key].to(torch.float8_e4m3fn) # convert to fp8. put a # before .to to keep the original dtype

save_file(st, "//path_to/model_name.safetensors", metadata={"format": "pt", **{}})

[github-actions[bot]]

This issue is being marked stale because it has not had any activity for 30 days. Reply below within 7 days if your issue still isn't solved, and it will be left open. Otherwise, the issue will be closed automatically.