Cannot Merge Lora weights back to the base model

Intermediate
1

I have a Flux-Dev base model which has been trained with the lora technique using the SimpleTuner framework (SimpleTuner/documentation/quickstart/FLUX.md at main · bghira/SimpleTuner · GitHub). The training process has generated checkpoint folders containing the following files:

  • optimizer.bin
  • pytorch_lora_weights.safetensors
  • scheduler.bin
  • random_states_0.pkl
  • training_state-mydataset.json
  • training_state.json

I want to load the Flux base model and the lora weights from one of the checkpoints, and merge the lora weights into the base model. Following the Lora documentation here LoRA I tried the following:

`
from peft import PeftModel
from diffusers import FluxTransformer2DModel

lora_path =<path_to_checkpoint_folder>
base_model = FluxTransformer2DModel.from_pretrained(<local_path_to_flux_model>,
subfolder=‘transformer’,
)
model = PeftModel.from_pretrained(base_model, lora_path)
model.merge_and_unload()
`

I get the following error:
ValueError: Can’t find ‘adapter_config.json’ at <path_to_checkpoint_folder>’

There is indeed no such file in there. How can I load the lora weights and merge them to the model in this case? Is there another way?

1 Like
2

Try this.

#lora_path =<path_to_checkpoint_folder>
lora_path =<path_to_checkpoint_file>
# lora_path ="pytorch_lora_weights.safetensors" # for example
3

Thanks for the answer John. I tried your suggestion and I get the same error again:

ValueError: Can’t find ‘adapter_config.json’ at <path_to_checkpoint_file>

4

from_pretrained() is a function that reads the HF repo format. For single files, use LoraModel class or load_lora_weights() from pipeline.

pipe.load_lora_weights(
    pretrained_model_name_or_path_or_dict="loras",
    weight_name="lora.safetensors"
)
5

I have tried using the pipeline load_lora_weights() as well and that works great for running inference with the pipeline.
However, I do not see an obvious way how to merge the LoRA weights into the base model to be able to use just the merged model for example for further training on different datasets.

6 
from pathlib import Path
a_list = []
w_list = []
for lora in lora_file_list:
    new_lora_file = lora
    w_name = Path(new_lora_file).name
    a_name = Path(new_lora_file).stem
    pipe.load_lora_weights(new_lora_file, weight_name=w_name, adapter_name=a_name, low_cpu_mem_usage=False)
    a_list.append(a_name)
    w_list.append(1.0)

pipe.set_adapters(a_list, adapter_weights=w_list)
pipe.fuse_lora(adapter_names=a_list, lora_scale=1.0)
pipe.unload_lora_weights()
pipe.save_pretrained("lora_merged_model")

If you want to apply LoRA only to the transformer, I think you can just torch.load_file with the safetensors library as just a torch model, then apply LoRA with PEFT and save_file.

7

I do indeed want to apply LoRA only to the transformer. Here is what I got so far:

import safetensors

lora_weights_file = <path_to_checkpoint_file>
lora_weights_dict = safetensors.torch.load_file(lora_weights_file)

Now I have a dictionary of the LoRA weights. How can I apply the LoRA weights with PEFT to the model?

1 Like
8

Sorry, I made a mistake. Apparently you don’t have to load up to state_dict, just the normal model.

base_model = FluxTransformer2DModel.from_pretrained(<local_path_to_flux_model>,
subfolder=‘transformer’,
)

Now all you should have to do is set up LoraConfig and do get_peft_model(), but I don’t know the proper contents of LoraConfig in this case.
Usually Pipeline internal does it on its own…

Edit:
I found it.