Flux 2 Klein load lora weights LoKr error #13137

Describe the bug

ValueError: original_state_dict should be empty at this point but has original_state_dict.keys()=dict_keys(['double_blocks.0.img_attn.proj.alpha', 'double_blocks.0.img_attn.proj.lokr_w1', 'double_blocks.0.img_attn.proj.lokr_w2', 'double_blocks.0.img_attn.qkv.alpha', 'double_blocks.0.img_attn.qkv.lokr_w1', 'double_blocks.0.img_attn.qkv.lokr_w2', 'double_blocks.0.img_mlp.0.alpha', 'double_blocks.0.img_mlp.0.lokr_w1', 'double_blocks.0.img_mlp.0.lokr_w2', 'double_blocks.0.img_mlp.2.alpha', 'double_blocks.0.img_mlp.2.lokr_w1', 'double_blocks.0.img_mlp.2.lokr_w2', 'double_blocks.0.txt_attn.proj.alpha', 'double_blocks.0.txt_attn.proj.lokr_w1', 'double_blocks.0.txt_attn.proj.lokr_w2', 'double_blocks.0.txt_attn.qkv.alpha', 'double_blocks.0.txt_attn.qkv.lokr_w1', 'double_blocks.0.txt_attn.qkv.lokr_w2', 'double_blocks.0.txt_mlp.0.alpha', 'double_blocks.0.txt_mlp.0.lokr_w1', 'double_blocks.0.txt_mlp.0.lokr_w2', 'double_blocks.0.txt_mlp.2.alpha', 'double_blocks.0.txt_mlp.2.lokr_w1', 'double_blocks.0.txt_mlp.2.lokr_w2', 'double_blocks.1.img_attn.proj.alpha', 'double_blocks.1.img_attn.proj.lokr_w1', 'double_blocks.1.img_attn.proj.lokr_w2', 'double_blocks.1.img_attn.qkv.alpha', 'double_blocks.1.img_attn.qkv.lokr_w1', 'double_blocks.1.img_attn.qkv.lokr_w2', 'double_blocks.1.img_mlp.0.alpha', 'double_blocks.1.img_mlp.0.lokr_w1', 'double_blocks.1.img_mlp.0.lokr_w2', 'double_blocks.1.img_mlp.2.alpha', 'double_blocks.1.img_mlp.2.lokr_w1', 'double_blocks.1.img_mlp.2.lokr_w2', 'double_blocks.1.txt_attn.proj.alpha', 'double_blocks.1.txt_attn.proj.lokr_w1', 'double_blocks.1.txt_attn.proj.lokr_w2', 'double_blocks.1.txt_attn.qkv.alpha', 'double_blocks.1.txt_attn.qkv.lokr_w1', 'double_blocks.1.txt_attn.qkv.lokr_w2', 'double_blocks.1.txt_mlp.0.alpha', 'double_blocks.1.txt_mlp.0.lokr_w1', 'double_blocks.1.txt_mlp.0.lokr_w2', 'double_blocks.1.txt_mlp.2.alpha', 'double_blocks.1.txt_mlp.2.lokr_w1', 'double_blocks.1.txt_mlp.2.lokr_w2', 'double_blocks.2.img_attn.proj.alpha', 'double_blocks.2.img_attn.proj.lokr_w1', 'double_blocks.2.img_attn.proj.lokr_w2', 'double_blocks.2.img_attn.qkv.alpha', 'double_blocks.2.img_attn.qkv.lokr_w1', 'double_blocks.2.img_attn.qkv.lokr_w2', 'double_blocks.2.img_mlp.0.alpha', 'double_blocks.2.img_mlp.0.lokr_w1', 'double_blocks.2.img_mlp.0.lokr_w2', 'double_blocks.2.img_mlp.2.alpha', 'double_blocks.2.img_mlp.2.lokr_w1', 'double_blocks.2.img_mlp.2.lokr_w2', 'double_blocks.2.txt_attn.proj.alpha', 'double_blocks.2.txt_attn.proj.lokr_w1', 'double_blocks.2.txt_attn.proj.lokr_w2', 'double_blocks.2.txt_attn.qkv.alpha', 'double_blocks.2.txt_attn.qkv.lokr_w1', 'double_blocks.2.txt_attn.qkv.lokr_w2', 'double_blocks.2.txt_mlp.0.alpha', 'double_blocks.2.txt_mlp.0.lokr_w1', 'double_blocks.2.txt_mlp.0.lokr_w2', 'double_blocks.2.txt_mlp.2.alpha', 'double_blocks.2.txt_mlp.2.lokr_w1', 'double_blocks.2.txt_mlp.2.lokr_w2', 'double_blocks.3.img_attn.proj.alpha', 'double_blocks.3.img_attn.proj.lokr_w1', 'double_blocks.3.img_attn.proj.lokr_w2', 'double_blocks.3.img_attn.qkv.alpha', 'double_blocks.3.img_attn.qkv.lokr_w1', 'double_blocks.3.img_attn.qkv.lokr_w2', 'double_blocks.3.img_mlp.0.alpha', 'double_blocks.3.img_mlp.0.lokr_w1', 'double_blocks.3.img_mlp.0.lokr_w2', 'double_blocks.3.img_mlp.2.alpha', 'double_blocks.3.img_mlp.2.lokr_w1', 'double_blocks.3.img_mlp.2.lokr_w2', 'double_blocks.3.txt_attn.proj.alpha', 'double_blocks.3.txt_attn.proj.lokr_w1', 'double_blocks.3.txt_attn.proj.lokr_w2', 'double_blocks.3.txt_attn.qkv.alpha', 'double_blocks.3.txt_attn.qkv.lokr_w1', 'double_blocks.3.txt_attn.qkv.lokr_w2', 'double_blocks.3.txt_mlp.0.alpha', 'double_blocks.3.txt_mlp.0.lokr_w1', 'double_blocks.3.txt_mlp.0.lokr_w2', 'double_blocks.3.txt_mlp.2.alpha', 'double_blocks.3.txt_mlp.2.lokr_w1', 'double_blocks.3.txt_mlp.2.lokr_w2', 'double_blocks.4.img_attn.proj.alpha', 'double_blocks.4.img_attn.proj.lokr_w1', 'double_blocks.4.img_attn.proj.lokr_w2', 'double_blocks.4.img_attn.qkv.alpha', 'double_blocks.4.img_attn.qkv.lokr_w1', 'double_blocks.4.img_attn.qkv.lokr_w2', 'double_blocks.4.img_mlp.0.alpha', 'double_blocks.4.img_mlp.0.lokr_w1', 'double_blocks.4.img_mlp.0.lokr_w2', 'double_blocks.4.img_mlp.2.alpha', 'double_blocks.4.img_mlp.2.lokr_w1', 'double_blocks.4.img_mlp.2.lokr_w2', 'double_blocks.4.txt_attn.proj.alpha', 'double_blocks.4.txt_attn.proj.lokr_w1', 'double_blocks.4.txt_attn.proj.lokr_w2', 'double_blocks.4.txt_attn.qkv.alpha', 'double_blocks.4.txt_attn.qkv.lokr_w1', 'double_blocks.4.txt_attn.qkv.lokr_w2', 'double_blocks.4.txt_mlp.0.alpha', 'double_blocks.4.txt_mlp.0.lokr_w1', 'double_blocks.4.txt_mlp.0.lokr_w2', 'double_blocks.4.txt_mlp.2.alpha', 'double_blocks.4.txt_mlp.2.lokr_w1', 'double_blocks.4.txt_mlp.2.lokr_w2', 'double_blocks.5.img_attn.proj.alpha', 'double_blocks.5.img_attn.proj.lokr_w1', 'double_blocks.5.img_attn.proj.lokr_w2', 'double_blocks.5.img_attn.qkv.alpha', 'double_blocks.5.img_attn.qkv.lokr_w1', 'double_blocks.5.img_attn.qkv.lokr_w2', 'double_blocks.5.img_mlp.0.alpha', 'double_blocks.5.img_mlp.0.lokr_w1', 'double_blocks.5.img_mlp.0.lokr_w2', 'double_blocks.5.img_mlp.2.alpha', 'double_blocks.5.img_mlp.2.lokr_w1', 'double_blocks.5.img_mlp.2.lokr_w2', 'double_blocks.5.txt_attn.proj.alpha', 'double_blocks.5.txt_attn.proj.lokr_w1', 'double_blocks.5.txt_attn.proj.lokr_w2', 'double_blocks.5.txt_attn.qkv.alpha', 'double_blocks.5.txt_attn.qkv.lokr_w1', 'double_blocks.5.txt_attn.qkv.lokr_w2', 'double_blocks.5.txt_mlp.0.alpha', 'double_blocks.5.txt_mlp.0.lokr_w1', 'double_blocks.5.txt_mlp.0.lokr_w2', 'double_blocks.5.txt_mlp.2.alpha', 'double_blocks.5.txt_mlp.2.lokr_w1', 'double_blocks.5.txt_mlp.2.lokr_w2', 'double_blocks.6.img_attn.proj.alpha', 'double_blocks.6.img_attn.proj.lokr_w1', 'double_blocks.6.img_attn.proj.lokr_w2', 'double_blocks.6.img_attn.qkv.alpha', 'double_blocks.6.img_attn.qkv.lokr_w1', 'double_blocks.6.img_attn.qkv.lokr_w2', 'double_blocks.6.img_mlp.0.alpha', 'double_blocks.6.img_mlp.0.lokr_w1', 'double_blocks.6.img_mlp.0.lokr_w2', 'double_blocks.6.img_mlp.2.alpha', 'double_blocks.6.img_mlp.2.lokr_w1', 'double_blocks.6.img_mlp.2.lokr_w2', 'double_blocks.6.txt_attn.proj.alpha', 'double_blocks.6.txt_attn.proj.lokr_w1', 'double_blocks.6.txt_attn.proj.lokr_w2', 'double_blocks.6.txt_attn.qkv.alpha', 'double_blocks.6.txt_attn.qkv.lokr_w1', 'double_blocks.6.txt_attn.qkv.lokr_w2', 'double_blocks.6.txt_mlp.0.alpha', 'double_blocks.6.txt_mlp.0.lokr_w1', 'double_blocks.6.txt_mlp.0.lokr_w2', 'double_blocks.6.txt_mlp.2.alpha', 'double_blocks.6.txt_mlp.2.lokr_w1', 'double_blocks.6.txt_mlp.2.lokr_w2', 'double_blocks.7.img_attn.proj.alpha', 'double_blocks.7.img_attn.proj.lokr_w1', 'double_blocks.7.img_attn.proj.lokr_w2', 'double_blocks.7.img_attn.qkv.alpha', 'double_blocks.7.img_attn.qkv.lokr_w1', 'double_blocks.7.img_attn.qkv.lokr_w2', 'double_blocks.7.img_mlp.0.alpha', 'double_blocks.7.img_mlp.0.lokr_w1', 'double_blocks.7.img_mlp.0.lokr_w2', 'double_blocks.7.img_mlp.2.alpha', 'double_blocks.7.img_mlp.2.lokr_w1', 'double_blocks.7.img_mlp.2.lokr_w2', 'double_blocks.7.txt_attn.proj.alpha', 'double_blocks.7.txt_attn.proj.lokr_w1', 'double_blocks.7.txt_attn.proj.lokr_w2', 'double_blocks.7.txt_attn.qkv.alpha', 'double_blocks.7.txt_attn.qkv.lokr_w1', 'double_blocks.7.txt_attn.qkv.lokr_w2', 'double_blocks.7.txt_mlp.0.alpha', 'double_blocks.7.txt_mlp.0.lokr_w1', 'double_blocks.7.txt_mlp.0.lokr_w2', 'double_blocks.7.txt_mlp.2.alpha', 'double_blocks.7.txt_mlp.2.lokr_w1', 'double_blocks.7.txt_mlp.2.lokr_w2', 'single_blocks.0.linear1.alpha', 'single_blocks.0.linear1.lokr_w1', 'single_blocks.0.linear1.lokr_w2', 'single_blocks.0.linear2.alpha', 'single_blocks.0.linear2.lokr_w1', 'single_blocks.0.linear2.lokr_w2', 'single_blocks.1.linear1.alpha', 'single_blocks.1.linear1.lokr_w1', 'single_blocks.1.linear1.lokr_w2', 'single_blocks.1.linear2.alpha', 'single_blocks.1.linear2.lokr_w1', 'single_blocks.1.linear2.lokr_w2', 'single_blocks.10.linear1.alpha', 'single_blocks.10.linear1.lokr_w1', 'single_blocks.10.linear1.lokr_w2', 'single_blocks.10.linear2.alpha', 'single_blocks.10.linear2.lokr_w1', 'single_blocks.10.linear2.lokr_w2', 'single_blocks.11.linear1.alpha', 'single_blocks.11.linear1.lokr_w1', 'single_blocks.11.linear1.lokr_w2', 'single_blocks.11.linear2.alpha', 'single_blocks.11.linear2.lokr_w1', 'single_blocks.11.linear2.lokr_w2', 'single_blocks.12.linear1.alpha', 'single_blocks.12.linear1.lokr_w1', 'single_blocks.12.linear1.lokr_w2', 'single_blocks.12.linear2.alpha', 'single_blocks.12.linear2.lokr_w1', 'single_blocks.12.linear2.lokr_w2', 'single_blocks.13.linear1.alpha', 'single_blocks.13.linear1.lokr_w1', 'single_blocks.13.linear1.lokr_w2', 'single_blocks.13.linear2.alpha', 'single_blocks.13.linear2.lokr_w1', 'single_blocks.13.linear2.lokr_w2', 'single_blocks.14.linear1.alpha', 'single_blocks.14.linear1.lokr_w1', 'single_blocks.14.linear1.lokr_w2', 'single_blocks.14.linear2.alpha', 'single_blocks.14.linear2.lokr_w1', 'single_blocks.14.linear2.lokr_w2', 'single_blocks.15.linear1.alpha', 'single_blocks.15.linear1.lokr_w1', 'single_blocks.15.linear1.lokr_w2', 'single_blocks.15.linear2.alpha', 'single_blocks.15.linear2.lokr_w1', 'single_blocks.15.linear2.lokr_w2', 'single_blocks.16.linear1.alpha', 'single_blocks.16.linear1.lokr_w1', 'single_blocks.16.linear1.lokr_w2', 'single_blocks.16.linear2.alpha', 'single_blocks.16.linear2.lokr_w1', 'single_blocks.16.linear2.lokr_w2', 'single_blocks.17.linear1.alpha', 'single_blocks.17.linear1.lokr_w1', 'single_blocks.17.linear1.lokr_w2', 'single_blocks.17.linear2.alpha', 'single_blocks.17.linear2.lokr_w1', 'single_blocks.17.linear2.lokr_w2', 'single_blocks.18.linear1.alpha', 'single_blocks.18.linear1.lokr_w1', 'single_blocks.18.linear1.lokr_w2', 'single_blocks.18.linear2.alpha', 'single_blocks.18.linear2.lokr_w1', 'single_blocks.18.linear2.lokr_w2', 'single_blocks.19.linear1.alpha', 'single_blocks.19.linear1.lokr_w1', 'single_blocks.19.linear1.lokr_w2', 'single_blocks.19.linear2.alpha', 'single_blocks.19.linear2.lokr_w1', 'single_blocks.19.linear2.lokr_w2', 'single_blocks.2.linear1.alpha', 'single_blocks.2.linear1.lokr_w1', 'single_blocks.2.linear1.lokr_w2', 'single_blocks.2.linear2.alpha', 'single_blocks.2.linear2.lokr_w1', 'single_blocks.2.linear2.lokr_w2', 'single_blocks.20.linear1.alpha', 'single_blocks.20.linear1.lokr_w1', 'single_blocks.20.linear1.lokr_w2', 'single_blocks.20.linear2.alpha', 'single_blocks.20.linear2.lokr_w1', 'single_blocks.20.linear2.lokr_w2', 'single_blocks.21.linear1.alpha', 'single_blocks.21.linear1.lokr_w1', 'single_blocks.21.linear1.lokr_w2', 'single_blocks.21.linear2.alpha', 'single_blocks.21.linear2.lokr_w1', 'single_blocks.21.linear2.lokr_w2', 'single_blocks.22.linear1.alpha', 'single_blocks.22.linear1.lokr_w1', 'single_blocks.22.linear1.lokr_w2', 'single_blocks.22.linear2.alpha', 'single_blocks.22.linear2.lokr_w1', 'single_blocks.22.linear2.lokr_w2', 'single_blocks.23.linear1.alpha', 'single_blocks.23.linear1.lokr_w1', 'single_blocks.23.linear1.lokr_w2', 'single_blocks.23.linear2.alpha', 'single_blocks.23.linear2.lokr_w1', 'single_blocks.23.linear2.lokr_w2', 'single_blocks.3.linear1.alpha', 'single_blocks.3.linear1.lokr_w1', 'single_blocks.3.linear1.lokr_w2', 'single_blocks.3.linear2.alpha', 'single_blocks.3.linear2.lokr_w1', 'single_blocks.3.linear2.lokr_w2', 'single_blocks.4.linear1.alpha', 'single_blocks.4.linear1.lokr_w1', 'single_blocks.4.linear1.lokr_w2', 'single_blocks.4.linear2.alpha', 'single_blocks.4.linear2.lokr_w1', 'single_blocks.4.linear2.lokr_w2', 'single_blocks.5.linear1.alpha', 'single_blocks.5.linear1.lokr_w1', 'single_blocks.5.linear1.lokr_w2', 'single_blocks.5.linear2.alpha', 'single_blocks.5.linear2.lokr_w1', 'single_blocks.5.linear2.lokr_w2', 'single_blocks.6.linear1.alpha', 'single_blocks.6.linear1.lokr_w1', 'single_blocks.6.linear1.lokr_w2', 'single_blocks.6.linear2.alpha', 'single_blocks.6.linear2.lokr_w1', 'single_blocks.6.linear2.lokr_w2', 'single_blocks.7.linear1.alpha', 'single_blocks.7.linear1.lokr_w1', 'single_blocks.7.linear1.lokr_w2', 'single_blocks.7.linear2.alpha', 'single_blocks.7.linear2.lokr_w1', 'single_blocks.7.linear2.lokr_w2', 'single_blocks.8.linear1.alpha', 'single_blocks.8.linear1.lokr_w1', 'single_blocks.8.linear1.lokr_w2', 'single_blocks.8.linear2.alpha', 'single_blocks.8.linear2.lokr_w1', 'single_blocks.8.linear2.lokr_w2', 'single_blocks.9.linear1.alpha', 'single_blocks.9.linear1.lokr_w1', 'single_blocks.9.linear1.lokr_w2', 'single_blocks.9.linear2.alpha', 'single_blocks.9.linear2.lokr_w1', 'single_blocks.9.linear2.lokr_w2']).

Trying to load lokr weights with load_lora_weights

Reproduction

loading LoKr weights from CivitAI

Logs

System Info

Python 3.11
diffusers 0.37.0-dev

Who can help?

No response