Some trained Flux models have state_dict keys starting with `model.diffusion_model.` causing error in loading

[hablaba]

I noticed this when trying to load a non base flux dev model from CivitAI to train a Lora on top of. This basic script seemed to help convert it:

from safetensors.torch import load_file, save_file

dict = load_file("<model_path>", device="cuda")
keys = dict.keys()
fixed= {k.replace("model.diffusion_model.", ""): dict[k] for k in keys}

print(fixed.keys())
save_file(fixed, "<output_model_path")

Perhaps you have a better way of fixing this but figured I'd at least provide the script for posterity :)

[kohya-ss]

⛔⛔⛔⛔⛔⛔⛔⛔ THIS LOOKS DANGEROUS

@notdanilo Thanks! I deleted your reply including the link.

[kohya-ss]

It looks like the model is using different keys than the official BFL model. If this is the only model that gives you the error, it would be best to ask the model author to fix it.

If there are multiple models, it's likely that there are tools out there that create incompatible models. I don't think this is a good thing... The loading script can handle this, but first, please let me know what's going on.

[hablaba]

Sure, makes sense. I was just mostly trying to post this script in case anyone else found it useful. I could potentially add a section to the readme for Flux with this as a utility method in case anyone needs to make the conversion too / runs into the same issue. Or can just leave it here for reference, that's fine as well!