Fix Flux multiple Lora loading bug #10388

maxs-kan · 2024-12-26T10:35:40Z

What does this PR do?

The current approach of checking for a key with a base_layer suffix may lead to a bug when multiple Lora models are loaded. If the first loaded Lora model does not have weights for layer n, and the second one does, loading the second model will lead to an error since the transformer state dict currently does not have key n.base_layer.weight. So I explicitly check for the presence of a key with the base_layer suffix.

@yiyixuxu

HuggingFaceDocBuilderDev · 2024-12-26T13:59:19Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

hlky · 2024-12-26T14:25:40Z

Hi @maxs-kan, thanks for your contribution, can you share some example lora checkpoints that may lead to a bug?

maxs-kan · 2024-12-26T15:52:36Z

Sure, try in the same order:
pipe.load_lora_weights(hf_hub_download("TTPlanet/Migration_Lora_flux","Migration_Lora_cloth.safetensors"), adapter_name="cloth")
pipe.load_lora_weights("alimama-creative/FLUX.1-Turbo-Alpha", adapter_name="turbo")

hlky

Code

from diffusers import FluxPipeline
from huggingface_hub import hf_hub_download
import torch

pipe = FluxPipeline.from_pretrained(
  "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
)
pipe.load_lora_weights(
  hf_hub_download("TTPlanet/Migration_Lora_flux", "Migration_Lora_cloth.safetensors"),
  adapter_name="cloth",
)
pipe.load_lora_weights("alimama-creative/FLUX.1-Turbo-Alpha", adapter_name="turbo")

sayakpaul · 2024-12-27T14:58:47Z

@maxs-kan thanks for this PR. Do you want to also propagate the changes from #10396?

sayakpaul · 2024-12-27T15:02:35Z

src/diffusers/loaders/lora_pipeline.py

+        transformer_base_layer_keys = {
+            k[: -len(".base_layer.weight")] for k in transformer_state_dict.keys() if ".base_layer.weight" in k
+        }


Note base_layer substring can only be present when the underlying pipeline has at least one LoRA loaded that affects the layer under consideration. So, perhaps it's better to have an is_peft_loaded check?

In your PR description you mention:

If the first loaded Lora model does not have weights for layer n, and the second one does, loading the second model will lead to an error since the transformer state dict currently does not have key n.base_layer.weight.

Note that we may also have an opposite situation i.e., the first LoRA ckpt may have the params while the second LoRA may not. This is what I considered in #10388.

if is_peft_loaded and ".base_layer.weight" in k might be clearer that this is something when a lora is already loaded.

The case where the first LoRA has extra weights than the second is ok on main

Hyper-FLUX.1-dev-8steps-lora.safetensors

Purz/choose-your-own-adventure

or

alimama-creative/FLUX.1-Turbo-Alpha

TTPlanet/Migration_Lora_flux

In this case base_param_name is set to f"{k.replace(prefix, '')}.base_layer.weight" for the 2nd LoRA and all keys exist.

If loaded in the reverse order f"{k.replace(prefix, '')}.base_layer.weight" doesn't exist for the extra weights.

Purz/choose-your-own-adventure

Hyper-FLUX.1-dev-8steps-lora.safetensors

or

TTPlanet/Migration_Lora_flux

alimama-creative/FLUX.1-Turbo-Alpha

KeyError context_embedder.base_layer.weight

So for the extra weights we use f"{k.replace(prefix, '')}.weight". If another LoRA were loaded with context_embedder it would then use context_embedder.base_layer.weight.

We could continue if f"{k.replace(prefix, '')}.base_layer.weight" is not found but the extra weights may need to be expanded.

In this case, we are considering that LoRA params for certain modules exist in the first checkpoint while they don't exist in the second checkpoint (or any other subsequent checkpoint).

In this case, we don't want to expand no? Or am I missing something? Perhaps better expressed through a short test case like the one I added here?

The test case passes on main, the test case should be in the reverse order:

with tempfile.TemporaryDirectory() as tmpdirname: denoiser_state_dict = get_peft_model_state_dict(pipe.transformer) self.pipeline_class.save_lora_weights(tmpdirname, transformer_lora_layers=denoiser_state_dict) self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) pipe.unload_lora_weights() # Modify the state dict to exclude "x_embedder" related LoRA params. lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) lora_state_dict_without_xembedder = {k: v for k, v in lora_state_dict.items() if "x_embedder" not in k} pipe.load_lora_weights(lora_state_dict_without_xembedder, adapter_name="two") # Load state dict with `x_embedder`. pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"), adapter_name="one")

base_param_name = ( f"{k.replace(prefix, '')}.base_layer.weight" if is_peft_loaded else f"{k.replace(prefix, '')}.weight" ) > base_weight_param = transformer_state_dict[base_param_name] E KeyError: 'x_embedder.base_layer.weight' src\diffusers\loaders\lora_pipeline.py:2471: KeyError

I think we still want to check whether the param needs to be expanded

Cool, I understand it better now. Thanks!

Might be better to ship this PR with proper testing then. Okay with me.

sayakpaul · 2024-12-28T02:12:48Z

Also, I gave @hlky's code snippet here a try in #10396 branch and it seems to work.

hlky · 2024-12-29T11:49:51Z

src/diffusers/loaders/lora_pipeline.py

+        transformer_base_layer_keys = {
+            k[: -len(".base_layer.weight")] for k in transformer_state_dict.keys() if ".base_layer.weight" in k
+        }


if is_peft_loaded and ".base_layer.weight" in k might be clearer that this is something when a lora is already loaded.

hlky · 2024-12-29T11:59:50Z

src/diffusers/loaders/lora_pipeline.py

+        transformer_base_layer_keys = {
+            k[: -len(".base_layer.weight")] for k in transformer_state_dict.keys() if ".base_layer.weight" in k
+        }


The case where the first LoRA has extra weights than the second is ok on main

Hyper-FLUX.1-dev-8steps-lora.safetensors

Purz/choose-your-own-adventure

or

alimama-creative/FLUX.1-Turbo-Alpha

TTPlanet/Migration_Lora_flux

In this case base_param_name is set to f"{k.replace(prefix, '')}.base_layer.weight" for the 2nd LoRA and all keys exist.

If loaded in the reverse order f"{k.replace(prefix, '')}.base_layer.weight" doesn't exist for the extra weights.

Purz/choose-your-own-adventure

Hyper-FLUX.1-dev-8steps-lora.safetensors

or

TTPlanet/Migration_Lora_flux

alimama-creative/FLUX.1-Turbo-Alpha

KeyError context_embedder.base_layer.weight

So for the extra weights we use f"{k.replace(prefix, '')}.weight". If another LoRA were loaded with context_embedder it would then use context_embedder.base_layer.weight.

We could continue if f"{k.replace(prefix, '')}.base_layer.weight" is not found but the extra weights may need to be expanded.

hlky · 2024-12-29T12:03:12Z

src/diffusers/loaders/lora_pipeline.py

            base_param_name = (
-                f"{k.replace(prefix, '')}.base_layer.weight" if is_peft_loaded else f"{k.replace(prefix, '')}.weight"
+                f"{k.replace(prefix, '')}.base_layer.weight"
+                if k in transformer_base_layer_keys
+                else f"{k.replace(prefix, '')}.weight"
            )


base_param_name = f"{k.replace(prefix, '')}.weight" base_layer_name = f"{k.replace(prefix, '')}.base_layer.weight" if is_peft_loaded and base_layer_name in transformer_state_dict: base_param_name = base_layer_name

Something like this might be better.

sayakpaul

@hlky thanks!

Do you wanna propagate your suggestions, too?

hlky

Thanks @sayakpaul. I've left some other comments but should be good to go

Thanks for the PR @maxs-kan

src/diffusers/loaders/lora_pipeline.py

tests/lora/test_lora_layers_flux.py

sayakpaul

Thanks.

I think the test mimics the code here that was producing the error. So, I think we should be good to go.

tests/lora/test_lora_layers_flux.py

Co-authored-by: Sayak Paul <[email protected]>

…or_extra_keys

yiyixuxu · 2025-01-02T18:35:24Z

thank you all! @maxs-kan @hlky @sayakpaul

bghira · 2025-01-12T21:24:08Z

this is a pretty big regression that caused some consumers to need to pull a custom build of diffusers with this patch included. can there perhaps be a hotfix pushed for v0.32.2?

sayakpaul · 2025-01-13T01:49:47Z

Sorry this should go into a patch release. @yiyixuxu I am happy to do the patch release if you're okay.

* check for base_layer key in transformer state dict * test_lora_expansion_works_for_absent_keys * check * Update tests/lora/test_lora_layers_flux.py Co-authored-by: Sayak Paul <[email protected]> * check * test_lora_expansion_works_for_absent_keys/test_lora_expansion_works_for_extra_keys * absent->extra --------- Co-authored-by: hlky <[email protected]> Co-authored-by: Sayak Paul <[email protected]>

check for base_layer key in transformer state dict

f099b2f

a-r-r-o-w requested a review from hlky December 26, 2024 12:25

Merge branch 'main' into flux-lora-base_layer-check

db010ae

hlky approved these changes Dec 27, 2024

View reviewed changes

sayakpaul reviewed Dec 27, 2024

View reviewed changes

Merge branch 'main' into flux-lora-base_layer-check

bdc5de5

hlky reviewed Dec 29, 2024

View reviewed changes

sayakpaul mentioned this pull request Dec 29, 2024

[LoRA] fix indexing in LoRA state dict expansion utils #10396

Closed

hlky added 2 commits December 29, 2024 15:34

test_lora_expansion_works_for_absent_keys

3a4f8a4

Merge branch 'main' into flux-lora-base_layer-check

da00c8d

sayakpaul reviewed Dec 30, 2024

View reviewed changes

check

a2cdcda

hlky approved these changes Dec 30, 2024

View reviewed changes

src/diffusers/loaders/lora_pipeline.py Outdated Show resolved Hide resolved

tests/lora/test_lora_layers_flux.py Show resolved Hide resolved

sayakpaul approved these changes Dec 30, 2024

View reviewed changes

tests/lora/test_lora_layers_flux.py Outdated Show resolved Hide resolved

hlky and others added 4 commits December 30, 2024 07:49

Update tests/lora/test_lora_layers_flux.py

c8d4a1c

Co-authored-by: Sayak Paul <[email protected]>

check

75268c0

test_lora_expansion_works_for_absent_keys/test_lora_expansion_works_f…

08ea124

…or_extra_keys

absent->extra

5a7997b

yiyixuxu merged commit 44640c8 into huggingface:main Jan 2, 2025
12 checks passed

maxs-kan deleted the flux-lora-base_layer-check branch January 13, 2025 10:57

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix Flux multiple Lora loading bug #10388

Fix Flux multiple Lora loading bug #10388

maxs-kan commented Dec 26, 2024

HuggingFaceDocBuilderDev commented Dec 26, 2024

hlky commented Dec 26, 2024

maxs-kan commented Dec 26, 2024 •

edited

Loading

hlky left a comment

sayakpaul commented Dec 27, 2024

sayakpaul Dec 27, 2024

sayakpaul Dec 27, 2024

hlky Dec 29, 2024

hlky Dec 29, 2024

sayakpaul Dec 29, 2024

hlky Dec 29, 2024

sayakpaul Dec 29, 2024

sayakpaul commented Dec 28, 2024

hlky Dec 29, 2024

hlky Dec 29, 2024

hlky Dec 29, 2024

sayakpaul left a comment

hlky left a comment

sayakpaul left a comment

yiyixuxu commented Jan 2, 2025

bghira commented Jan 12, 2025

sayakpaul commented Jan 13, 2025

Fix Flux multiple Lora loading bug #10388

Fix Flux multiple Lora loading bug #10388

Conversation

maxs-kan commented Dec 26, 2024

What does this PR do?

HuggingFaceDocBuilderDev commented Dec 26, 2024

hlky commented Dec 26, 2024

maxs-kan commented Dec 26, 2024 • edited Loading

hlky left a comment

Choose a reason for hiding this comment

sayakpaul commented Dec 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Dec 28, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul left a comment

Choose a reason for hiding this comment

hlky left a comment

Choose a reason for hiding this comment

sayakpaul left a comment

Choose a reason for hiding this comment

yiyixuxu commented Jan 2, 2025

bghira commented Jan 12, 2025

sayakpaul commented Jan 13, 2025

maxs-kan commented Dec 26, 2024 •

edited

Loading