Flux merge tool does not properly handle LoRA keys from AI-Toolkit when merging w/ Flux.dev model - All keys "Unused" / Unchanged model file outputted

[CCpt5]

Hello,

Currently the featuring to merge Flux LoRA into the base model does not function properly with LoRA trained by Ostris' AI-Toolkit. This appears to be due to a difference in the way the keys are labeled/named. The tool does seem able to merge AI-Toolkit LoRA into other Kohya/AI-Toolkit trained LoRA, but merging into the base model fails at the moment.

Thanks for any assistance!

Console example:

2024-09-18 00:30:17 INFO     loading keys from FLUX.1 model: D:/FLUX.1-dev/flux1-dev.safetensors                                                                                                                                                                                                        flux_merge_lora.py:55
                    INFO     loading: D:/ai-toolkit/output-mthrbrd/motherboard/motherboard_000005000.safetensors                                                                                                                                                                                        flux_merge_lora.py:74
                    INFO     merging...                                                                                                                                                                                                                                                                 flux_merge_lora.py:77
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 988/988 [00:00<?, ?it/s]
                    WARNING  Unused keys in LoRA model: ['transformer.single_transformer_blocks.0.attn.to_k.lora_A.weight', 'transformer.single_transformer_blocks.0.attn.to_k.lora_B.weight', 'transformer.single_transformer_blocks.0.attn.to_q.lora_A.weight',                                      flux_merge_lora.py:130
                             'transformer.single_transformer_blocks.0.attn.to_q.lora_B.weight', 'transformer.single_transformer_blocks.0.attn.to_v.lora_A.weight', 'transformer.single_transformer_blocks.0.attn.to_v.lora_B.weight',
                             'transformer.single_transformer_blocks.0.norm.linear.lora_A.weight', 'transformer.single_transformer_blocks.0.norm.linear.lora_B.weight', 'transformer.single_transformer_blocks.0.proj_mlp.lora_A.weight',
                             'transformer.single_transformer_blocks.0.proj_mlp.lora_B.weight', 'transformer.single_transformer_blocks.0.proj_out.lora_A.weight', 'transformer.single_transformer_blocks.0.proj_out.lora_B.weight', 'transformer.single_transformer_blocks.1.attn.to_k.lora_A.weight',
                             'transformer.single_transformer_blocks.1.attn.to_k.lora_B.weight', 'transformer.single_transformer_blocks.1.attn.to_q.lora_A.weight', 'transformer.single_transformer_blocks.1.attn.to_q.lora_B.weight',
                             'transformer.single_transformer_blocks.1.attn.to_v.lora_A.weight', 'transformer.single_transformer_blocks.1.attn.to_v.lora_B.weight', 'transformer.single_transformer_blocks.1.norm.linear.lora_A.weight',
                             'transformer.single_transformer_blocks.1.norm.linear.lora_B.weight', 'transformer.single_transformer_blocks.1.proj_mlp.lora_A.weight', 'transformer.single_transformer_blocks.1.proj_mlp.lora_B.weight',
                             'transformer.single_transformer_blocks.1.proj_out.lora_A.weight', 'transformer.single_transformer_blocks.1.proj_out.lora_B.weight', 'transformer.single_transformer_blocks.10.attn.to_k.lora_A.weight',
                             'transformer.single_transformer_blocks.10.attn.to_k.lora_B.weight', 'transformer.single_transformer_blocks.10.attn.to_q.lora_A.weight', 'transformer.single_transformer_blocks.10.attn.to_q.lora_B.weight',
                             'transformer.single_transformer_blocks.10.attn.to_v.lora_A.weight', 'transformer.single_transformer_blocks.10.attn.to_v.lora_B.weight', 'transformer.single_transformer_blocks.10.norm.linear.lora_A.weight',
                             'transformer.single_transformer_blocks.10.norm.linear.lora_B.weight', 'transformer.single_transformer_blocks.10.proj_mlp.lora_A.weight', 'transformer.single_transformer_blocks.10.proj_mlp.lora_B.weight',
                             'transformer.single_transformer_blocks.10.proj_out.lora_A.weight', 'transformer.single_transformer_blocks.10.proj_out.lora_B.weight', 'transformer.single_transformer_blocks.11.attn.to_k.lora_A.weight',
                             'transformer.single_transformer_blocks.11.attn.to_k.lora_B.weight', 'transformer.single_transformer_blocks.11.attn.to_q.lora_A.weight', 'transformer.single_transformer_blocks.11.attn.to_q.lora_B.weight',
                             'transformer.single_transformer_blocks.11.attn.to_v.lora_A.weight', 'transformer.single_transformer_blocks.11.attn.to_v.lora_B.weight', 'transformer.single_transformer_blocks.11.norm.linear.lora_A.weight',
                             'transformer.single_transformer_blocks.11.norm.linear.lora_B.weight', 'transformer.single_transformer_blocks.11.proj_mlp.lora_A.weight', 'transformer.single_transformer_blocks.11.proj_mlp.lora_B.weight',
                             'transformer.single_transformer_blocks.11.proj_out.lora_A.weight', 'transformer.single_transformer_blocks.11.proj_out.lora_B.weight', 'transformer.single_transformer_blocks.12.attn.to_k.lora_A.weight',
                             'transformer.single_transformer_blocks.12.attn.to_k.lora_B.weight', 'transformer.single_transformer_blocks.12.attn.to_q.lora_A.weight', 'transformer.single_transformer_blocks.12.attn.to_q.lora_B.weight',

[Have to cut some out due to txt limits]

                             'transformer.transformer_blocks.9.norm1_context.linear.lora_B.weight']
                    INFO     saving FLUX model to: H:/sd-checkpoints-fast/2kpr-mthrbrd-ostris-merged-into-flux-dev-for-training.safetensors                                                                                                                                                            flux_merge_lora.py:549
                    INFO     converting to torch.float16...                                                                                                                                                                                                                                             flux_merge_lora.py:39
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 780/780 [01:32<00:00,  8.40it/s]
2024-09-18 00:31:50 INFO     saving to: H:/sd-checkpoints-fast/2kpr-mthrbrd-ostris-merged-into-flux-dev-for-training.safetensors                                                                                                                                                                        flux_merge_lora.py:44
00:32:09-317573 INFO     Done merging...
00:33:29-319169 INFO     Merge FLUX LoRA...
00:33:29-320169 INFO     Executing command: D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\venv\Scripts\python.exe D:/Doctor-SD/Kohya_GUI_Flux_Installer_21/kohya_ss/sd-scripts/networks/flux_merge_lora.py --save_precision fp16 --precision float --save_to H:/sd-checkpoints-fast/lora-merge-test.safetensors
                         --loading_device cpu --working_device cuda --models D:/ai-toolkit/output-mthrbrd/motherboard/motherboard_000005000.safetensors D:/stable-diffusion-webui/models/Lora/lcsrts_000007250.safetensors --ratios 0 0
D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\venv\lib\site-packages\diffusers\utils\outputs.py:63: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
  torch.utils._pytree._register_pytree_node(
D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\venv\lib\site-packages\diffusers\utils\outputs.py:63: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
  torch.utils._pytree._register_pytree_node(
2024-09-18 00:33:31 WARNING  WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:                                                                                                                                                                                              _cpp_lib.py:148
                                 PyTorch 2.4.0+cu121 with CUDA 1201 (you have 2.5.0.dev20240906+cu124)
                                 Python  3.10.11 (you have 3.10.11)
                               Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)
                               Memory-efficient attention, SwiGLU, sparse and more won't be available.
                               Set XFORMERS_MORE_DETAILS=1 for more details
D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\venv\lib\site-packages\diffusers\utils\outputs.py:63: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
  torch.utils._pytree._register_pytree_node(
D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\venv\lib\site-packages\diffusers\utils\outputs.py:63: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
  torch.utils._pytree._register_pytree_node(
2024-09-18 00:33:32 INFO     loading: D:/ai-toolkit/output-mthrbrd/motherboard/motherboard_000005000.safetensors                                                                                                                                                                                       flux_merge_lora.py:394
                    INFO     dim: [], alpha: []                                                                                                                                                                                                                                                        flux_merge_lora.py:425
                    INFO     merging...                                                                                                                                                                                                                                                                flux_merge_lora.py:428
  0%|                                                                                                                                                                                                                                                                                               | 0/988 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\sd-scripts\networks\flux_merge_lora.py", line 657, in <module>
    merge(args)
  File "D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\sd-scripts\networks\flux_merge_lora.py", line 553, in merge
    state_dict, metadata = merge_lora_models(args.models, args.ratios, merge_dtype, args.concat, args.shuffle)
  File "D:\Doctor-SD\Kohya_GUI_Flux_Installer_21\kohya_ss\sd-scripts\networks\flux_merge_lora.py", line 442, in merge_lora_models
    base_alpha = base_alphas[lora_module_name]
KeyError: 'transformer.single_transformer_blocks.0.attn.to_k'
00:33:33-203707 INFO     Done merging...

[kohya-ss]

Please add --diffusers option when merging AI-Toolkit LoRA. Unfortunately auto-detection is not yet implemented.

[Jonseed]

@CCpt5 in your output above, how can you tell which LoRA merged ok, and which didn't? Are "unused keys" list a bad merge or empty dim and alpha, or KeyError? It looks like they both failed with errors.

[CCpt5]

@CCpt5 in your output above, how can you tell which LoRA merged ok, and which didn't? Are "unused keys" list a bad merge or empty dim and alpha, or KeyError? It looks like they both failed with errors.

Yea, all of those failed. I did it properly through command line w/ the --diffusers flag as Kohya advised and it worked properly there (sorry don't have a screen capture of that). I was just hoping a toggle to add that flag could be added to the GUI to make it easier.

I think the one caveat at the moment is the LoRA w/ TXT encoder changes won't have those properly merged but there's a warning to that effect indicating they will be in the near future.

[Jonseed]

@CCpt5 How can we tell if a LoRA is AI-Toolkit or Kohya besides getting an error (or is that the only way)? Do we need to merge these LoRAs separately if we have both kinds, using the --diffusers flag only with the AI-Toolkit LoRAs?

[Mithadon]

When I tried yesterday, the sd3-flux.1 branch, there was a checkbox for diffusers lora which added --diffusers to the command. I managed to merge ai-toolkit Loras into a unet with it. But it still doesn't work for merging Loras together.

[CCpt5]

When I tried yesterday, the sd3-flux.1 branch, there was a checkbox for diffusers lora which added --diffusers to the command. I managed to merge ai-toolkit Loras into a unet with it. But it still doesn't work for merging Loras together.

Ahh interesting! I'll double check - maybe I didn't update as I should have.

Not sure on merging LoRA tbh - I haven't tried that to this point.

[Jonseed]

If there are a lot of "unused keys," does that mean it is AI-Toolkit format?

[CCpt5]

If there are a lot of "unused keys," does that mean it is AI-Toolkit format?

It may, or it could be that the TXT encoder(s) were trained as well for those LoRA. The console error should list that though if that's the cause.

[Jonseed]

When it said it had "unused keys" I added the --diffusers and it looks like it searched for, found, and merged the keys properly. One of my LoRAs did not have any "unused keys" so was probably already in Kohya format, and I didn't need to use the --diffusers.