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...