[LoRA] Add LoRA support to AuraFlow #10216

hameerabbasi · 2024-12-13T16:03:16Z

What does this PR do?

This PR is a simple rebase of #9017

cc @sayakpaul for review.

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

HuggingFaceDocBuilderDev · 2024-12-13T16:18:44Z

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.

hameerabbasi · 2024-12-13T16:44:56Z

Thanks for the helping hand, @hlky!

hlky · 2024-12-13T16:52:16Z

See https://github.com/huggingface/diffusers/blob/main/tests/lora/test_lora_layers_flux.py https://github.com/huggingface/diffusers/blob/main/tests/lora/test_lora_layers_mochi.py etc as an example for tests. Seems to be missing transformer_cls, get_dummy_inputs and probably others, ping @sayakpaul if you need help with lora tests.

hameerabbasi · 2024-12-15T04:10:41Z

@sayakpaul Okay; I'm at a point where I've got actual, valid test failures but have no idea where to look.
pytest.log

hameerabbasi · 2024-12-15T05:47:42Z

Here's the log after the latest commit: pytest.log

sayakpaul

Thanks for the PR. I have left some comments to fix a couple of things. LMK if they're unclear.

src/diffusers/loaders/lora_pipeline.py

hameerabbasi · 2024-12-19T09:28:25Z

Latest test log.

hameerabbasi · 2025-01-07T04:18:37Z

Current pytest.log

Failures:

=========================== short test summary info ============================
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_correct_lora_configs_with_different_ranks - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_lora_B_bias - AssertionError: True is not false
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_low_cpu_mem_usage_with_loading - AttributeError: 'AuraFlowPipeline' object has no attribute 'lora_scale'
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_set_adapters_match_attention_kwargs - AssertionError: True is not false : Lora + scale should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_dora - AssertionError: True is not false : DoRA lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_partial_text_lora - ValueError: Target modules {'out_proj', 'v_proj', 'q_proj', 'k_proj'} not found in the base model. Please check the target modules and try again.
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_block_scale - AssertionError: True is not false : LoRA weights 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_lora_and_scale - AssertionError: False is not true : Lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_lora_save_load - AttributeError: 'AuraFlowPipeline' object has no attribute 'lora_scale'
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter - AssertionError: True is not false : Adapter outputs should be different.
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_block_lora - AssertionError: True is not false : Adapter 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_delete_adapter - AssertionError: True is not false : Adapter 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_weighted - AssertionError: True is not false : Adapter 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora - AssertionError: False is not true : Lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_and_scale - AssertionError: False is not true : Lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_denoiser_fused - AssertionError: True is not false : Fused lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_fused - AssertionError: True is not false : Fused lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_save_load - AttributeError: 'AuraFlowPipeline' object has no attribute 'lora_scale'
===== 18 failed, 12 passed, 4 skipped, 397 deselected, 1 warning in 23.58s =====

hameerabbasi · 2025-01-07T05:25:52Z

It seems nothing but CLIP is supported in

diffusers/src/diffusers/models/lora.py

Lines 41 to 66 in 532013f

    
           def text_encoder_attn_modules(text_encoder): 
        
               attn_modules = [] 
        
               if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)): 
        
                   for i, layer in enumerate(text_encoder.text_model.encoder.layers): 
        
                       name = f"text_model.encoder.layers.{i}.self_attn" 
        
                       mod = layer.self_attn 
        
                       attn_modules.append((name, mod)) 
        
               else: 
        
                   raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}") 
        
               return attn_modules 
        
           def text_encoder_mlp_modules(text_encoder): 
        
               mlp_modules = [] 
        
               if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)): 
        
                   for i, layer in enumerate(text_encoder.text_model.encoder.layers): 
        
                       mlp_mod = layer.mlp 
        
                       name = f"text_model.encoder.layers.{i}.mlp" 
        
                       mlp_modules.append((name, mlp_mod)) 
        
               else: 
        
                   raise ValueError(f"do not know how to get mlp modules for: {text_encoder.__class__.__name__}") 
        
               return mlp_modules

which is called here

diffusers/src/diffusers/loaders/lora_pipeline.py

Lines 2185 to 2197 in 532013f

    
           for name, _ in text_encoder_attn_modules(text_encoder): 
        
               for module in ("out_proj", "q_proj", "k_proj", "v_proj"): 
        
                   rank_key = f"{name}.{module}.lora_B.weight" 
        
                   if rank_key not in text_encoder_lora_state_dict: 
        
                       continue 
        
                   rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1] 
        
           for name, _ in text_encoder_mlp_modules(text_encoder): 
        
               for module in ("fc1", "fc2"): 
        
                   rank_key = f"{name}.{module}.lora_B.weight" 
        
                   if rank_key not in text_encoder_lora_state_dict: 
        
                       continue 
        
                   rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1]

This means essentially that we need more plumbing to support this for arbitrary text encoders, or to only support the transformer for AuraFlow. This is because AuraFlow only has one UMT5 encoder.

sayakpaul · 2025-01-07T06:23:28Z

Thanks, let's support only transformer for the moment, then? @hameerabbasi also let's try to work on the failed tests.

hameerabbasi · 2025-01-08T14:45:50Z

So I skipped all tests requiring TE in e06d8eb. Latest failures are: pytest.log

=========================== short test summary info ============================
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_correct_lora_configs_with_different_ranks - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_lora_B_bias - AssertionError: True is not false
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_set_adapters_match_attention_kwargs - AssertionError: True is not false : Lora + scale should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_dora - AssertionError: True is not false : DoRA lora should change the output
==== 4 failed, 10 passed, 20 skipped, 398 deselected, 2 warnings in 15.99s =====

I'm not entirely sure how to get past these.

sayakpaul · 2025-01-08T14:48:04Z

Could we try to look into each failure and debug?

hameerabbasi · 2025-01-13T05:45:36Z

src/diffusers/loaders/lora_base.py

-            text_encoder_lora_state_dict = convert_state_dict_to_peft(text_encoder_lora_state_dict)
-
-            for name, _ in text_encoder_attn_modules(text_encoder):
-                for module in ("out_proj", "q_proj", "k_proj", "v_proj"):
-                    rank_key = f"{name}.{module}.lora_B.weight"
-                    if rank_key not in text_encoder_lora_state_dict:
-                        continue
-                    rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1]
-
-            for name, _ in text_encoder_mlp_modules(text_encoder):
-                for module in ("fc1", "fc2"):
-                    rank_key = f"{name}.{module}.lora_B.weight"
-                    if rank_key not in text_encoder_lora_state_dict:
-                        continue
-                    rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1]
+            text_encoder_lora_state_dict = convert_state_dict_to_peft(text_encoder_lora_state_dict, original_type=StateDictType.DIFFUSERS)
+
+            for name, module in text_encoder.named_modules():
+                if "lora_A" not in name and "lora_B" not in name and isinstance(module, (nn.Linear, nn.Conv2d)):
+                    rank_key = f"{name.removesuffix(".base_layer")}.lora_B.weight"
+                    if rank_key in text_encoder_lora_state_dict:
+                        rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1]


These changes support TE LoRA other than CLIP; but they cause one other model to fail, likely an edge case.

I think it's fine to not add these changes in this PR as it seems unrelated.

I'll make a separate PR for this and we can review that one first -- it seems to be the crux of what's failing with the tests IIUC. Then coming back to this one should be easy.

I don't think so. If we don't support loading LoRAs into a module (_lora_loadable_modules) then that shouldn't cause any failures.

It still causes the failures here, and I'm kinda clueless how to fix them: #10216 (comment)

For context, this is the first model we'd support without a TE, the tests aren't written for that.

I don't think so.

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 2560 in edb8c1b

class Mochi1LoraLoaderMixin(LoraBaseMixin):

doesn't have text encoder either.

neither does OmniGen

hameerabbasi · 2025-01-19T14:53:44Z

Okay, I've updated the code to mirror Mochi1 instead of Flux, as that was a closer match (no TE LoRA), and removed the TE from loadable modules. I've also skipped TE-only LoRA tests. Here are the remaining failures:

Failure list

=========================== short test summary info ============================
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_correct_lora_configs_with_different_ranks - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_lora_B_bias - AssertionError: True is not false
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_set_adapters_match_attention_kwargs - AssertionError: True is not false : Lora + scale should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_dora - AssertionError: True is not false : DoRA lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_lora_and_scale - AssertionError: False is not true : Lora should change the output
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter - AssertionError: True is not false : Adapter outputs should be different.
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_block_lora - AssertionError: True is not false : Adapter 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_delete_adapter - AssertionError: True is not false : Adapter 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_multi_adapter_weighted - AssertionError: True is not false : Adapter 1 and 2 should give different results
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_denoiser_fused - AssertionError: True is not false : Fused lora should change the output
==== 10 failed, 14 passed, 10 skipped, 398 deselected, 4 warnings in 20.21s ====

Latest pytest.log

sayakpaul · 2025-01-20T03:24:36Z

I think a better approach to debugging is taking a single test failure and trying to see what's causing them. We could compare the implementation with another model having similarities (Mochi-1 is a good one) and then take things from there.

Have we tried that?

We're grateful for your contributions so far. But it might be even better if we tried going further a bit for debugging. This will help solidify your contributions, too. SANALoRALoaderMixin is another class where we could take references from:

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 3199 in 328e0d2

class SanaLoraLoaderMixin(LoraBaseMixin):

hameerabbasi force-pushed the auraflow-lora branch from da41f32 to c8364bc Compare December 13, 2024 16:36

hameerabbasi force-pushed the auraflow-lora branch from 20b5f2d to 80ac0d4 Compare December 15, 2024 04:09

hameerabbasi force-pushed the auraflow-lora branch from 912fb8d to f5b9f90 Compare December 15, 2024 05:45

sayakpaul requested changes Dec 15, 2024

View reviewed changes

Warlord-K and others added 17 commits December 19, 2024 09:57

Add AuraFlowLoraLoaderMixin

e50a108

Add comments, remove qkv fusion

658d058

Add Tests

4208d09

Add AuraFlowLoraLoaderMixin to documentation

98b19f6

Add Suggested changes

71f8bac

Change attention_kwargs->joint_attention_kwargs

0eee03e

Rebasing derp.

4e4f780

fix

c07d1f5

fix

1b7f99f

Quality fixes.

875a3e0

make style

a242d7a

make fix-copies

a73df6b

ruff check --fix

894eac0

Attept 1 to fix tests.

2b36416

Attept 2 to fix tests.

6b762b8

Attept 3 to fix tests.

bc2a466

Address review comments.

1c79095

hameerabbasi force-pushed the auraflow-lora branch from f5b9f90 to 1c79095 Compare December 19, 2024 09:23

Rebasing derp.

9454e84

Merge branch 'main' into auraflow-lora

5700e52

hameerabbasi added 2 commits January 7, 2025 04:20

Add lora_scale property for te LoRAs.

2d02c2c

Make test better.

2b934b4

Remove useless property.

532013f

hlky and others added 2 commits January 8, 2025 13:36

Merge branch 'main' into auraflow-lora

0ea9ecd

Skip TE-only tests for AuraFlow.

e06d8eb

Support LoRA for non-CLIP TEs.

2b35909

hameerabbasi force-pushed the auraflow-lora branch 3 times, most recently from b59b25e to 00c921e Compare January 10, 2025 07:20

Merge remote-tracking branch 'upstream/main' into auraflow-lora

1ec07a1

hameerabbasi force-pushed the auraflow-lora branch from 00c921e to 1ec07a1 Compare January 10, 2025 07:20

hlky and others added 2 commits January 10, 2025 07:22

Merge branch 'main' into auraflow-lora

077a452

Merge branch 'main' into auraflow-lora

3095644

hameerabbasi commented Jan 13, 2025

View reviewed changes

hameerabbasi added 3 commits January 19, 2025 14:48

Merge remote-tracking branch 'upstream/main' into auraflow-lora

df28362

Restore LoRA tests.

7e63330

Undo adding LoRA support for non-CLIP TEs.

5620384

hameerabbasi force-pushed the auraflow-lora branch from 9a89ed0 to 5620384 Compare January 19, 2025 13:55

hameerabbasi added 4 commits January 19, 2025 14:56

Undo support for TE in AuraFlow LoRA.

cd691d3

make fix-copies

0fa5cd5

Sync with upstream changes.

83e0825

Remove unneeded stuff.

12fbd11

hameerabbasi force-pushed the auraflow-lora branch from 0751c4b to 12fbd11 Compare January 19, 2025 14:48

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[LoRA] Add LoRA support to AuraFlow #10216

[LoRA] Add LoRA support to AuraFlow #10216

hameerabbasi commented Dec 13, 2024

HuggingFaceDocBuilderDev commented Dec 13, 2024

hameerabbasi commented Dec 13, 2024

hlky commented Dec 13, 2024

hameerabbasi commented Dec 15, 2024 •

edited

Loading

hameerabbasi commented Dec 15, 2024

sayakpaul left a comment

hameerabbasi commented Dec 19, 2024

hameerabbasi commented Jan 7, 2025

hameerabbasi commented Jan 7, 2025

sayakpaul commented Jan 7, 2025 •

edited

Loading

hameerabbasi commented Jan 8, 2025

sayakpaul commented Jan 8, 2025

hameerabbasi Jan 13, 2025

sayakpaul Jan 13, 2025

hameerabbasi Jan 13, 2025

sayakpaul Jan 13, 2025

hameerabbasi Jan 13, 2025 •

edited

Loading

sayakpaul Jan 13, 2025

bghira Jan 13, 2025

hameerabbasi commented Jan 19, 2025

sayakpaul commented Jan 20, 2025

[LoRA] Add LoRA support to AuraFlow #10216

Are you sure you want to change the base?

[LoRA] Add LoRA support to AuraFlow #10216

Conversation

hameerabbasi commented Dec 13, 2024

What does this PR do?

Before submitting

Who can review?

HuggingFaceDocBuilderDev commented Dec 13, 2024

hameerabbasi commented Dec 13, 2024

hlky commented Dec 13, 2024

hameerabbasi commented Dec 15, 2024 • edited Loading

hameerabbasi commented Dec 15, 2024

sayakpaul left a comment

Choose a reason for hiding this comment

hameerabbasi commented Dec 19, 2024

hameerabbasi commented Jan 7, 2025

hameerabbasi commented Jan 7, 2025

sayakpaul commented Jan 7, 2025 • edited Loading

hameerabbasi commented Jan 8, 2025

sayakpaul commented Jan 8, 2025

hameerabbasi Jan 13, 2025

Choose a reason for hiding this comment

sayakpaul Jan 13, 2025

Choose a reason for hiding this comment

hameerabbasi Jan 13, 2025

Choose a reason for hiding this comment

sayakpaul Jan 13, 2025

Choose a reason for hiding this comment

hameerabbasi Jan 13, 2025 • edited Loading

Choose a reason for hiding this comment

sayakpaul Jan 13, 2025

Choose a reason for hiding this comment

bghira Jan 13, 2025

Choose a reason for hiding this comment

hameerabbasi commented Jan 19, 2025

sayakpaul commented Jan 20, 2025

hameerabbasi commented Dec 15, 2024 •

edited

Loading

sayakpaul commented Jan 7, 2025 •

edited

Loading

hameerabbasi Jan 13, 2025 •

edited

Loading