[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-06T15:09:35Z

@sayakpaul I've uploaded the latest test log in this comment, and commented the difficulties I faced with certain changes. Some pointers would be appreciated.

src/diffusers/loaders/lora_pipeline.py

src/diffusers/models/transformers/auraflow_transformer_2d.py

src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

tests/lora/test_lora_layers_af.py

sayakpaul

Thanks for your work!

I left some more comments. LMK if they make sense.

hameerabbasi · 2025-01-07T02:58:19Z

@sayakpaul Here's the latest test log: pytest.log

Here's the list of failures:

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_injection - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_low_cpu_mem_usage_with_loading - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_set_adapters_match_attention_kwargs - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_save_pretrained - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_dora - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_partial_text_lora - AssertionError: False is not true
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
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_lora_save_load - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_denoiser_lora_unloaded - AssertionError: False is not true
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
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_and_scale - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_denoiser_fused - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_denoiser_fused_multi - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_fused - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_save_load - AssertionError: False is not true
FAILED tests/lora/test_lora_layers_auraflow.py::AuraFlowLoRATests::test_simple_inference_with_text_lora_unloaded - AssertionError: False is not true
===== 24 failed, 6 passed, 4 skipped, 396 deselected, 1 warning in 24.20s ======

Summary of the latest changes: I copied from Flux mostly as it was the only one with a transformer and one TE. Most of the issues seem to be coming from the line:

self.assertTrue(output_no_lora.shape == self.output_shape)

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 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 mentioned this pull request Jan 6, 2025

[LoRA] Add LoRA support to AuraFlow #9017

Open

6 tasks