kohya-ss / kohya-ss/sd-scripts

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

Open
#1,612 9 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

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:

```shell

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

Contributor guide

No contributing guide indexed for this repository

Research direction

The failure is in sd-scripts/networks/flux_merge_lora.py, in merge_lora_models and merge; start by reproducing the logged command and tracing the key and alpha lookup around the reported KeyError. Done means AI-Toolkit LoRA keys are recognized when merging into a Flux.dev model without the unused-key warning or KeyError, and the output model is produced.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.