[LoRA] Add LoRA support to AuraFlow #9017

Warlord-K · 2024-07-30T12:53:36Z

What does this PR do?

Adds LoRA support to AuraFlow

from diffusers import AuraFlowPipeline
import torch 

pipe = AuraFlowPipeline.from_pretrained("fal/AuraFlow-v0.2", torch_dtype = torch.float16).to("cuda")
pipe.load_lora_weights("Warlord-K/gorkem-auraflow-lora", weight_name="pytorch_lora_weights.safetensors") # Set weight_name = "lora_peft_format.safetensors" to test loading from peft format
image = pipe("gorkem in a black tuxedo", generator = torch.Generator().manual_seed(2347862)).images[0]
image.save("test.png")

Following functions have also been tested taking SD3 LoRA Tests as reference:

pipe.load_lora_weights()
pipe.unload_lora_weights()
pipe.fuse_lora()
pipe.unfuse_lora()

Fusing lora decreases inference time by ~1.5s and unfusing it increases it again.

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?

@sayakpaul Please review

P.S. make style && make quality fails on some other file hence I wasnt able to run it

sayakpaul

Thanks! Left some comments.

I think the TODOs are:

Tests (like this)
Docs (like this)

LMK if you have questions.

src/diffusers/loaders/lora_pipeline.py

sayakpaul · 2024-07-30T12:58:47Z

src/diffusers/models/transformers/auraflow_transformer_2d.py

+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedJointAttnProcessor2_0
+    def fuse_qkv_projections(self):


This is not a part of the PR. Let's tackle this separately. Also #8952.

Sorry, I have removed it now.

sayakpaul · 2024-07-30T12:59:47Z

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -329,6 +371,7 @@ def forward(
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor = None,
        timestep: torch.LongTensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,


Hmm, AuraFlow has two kinds of attention, right? MMDiT blocks have joint attention and Single DiT blocks have regular attention. So, wondering if it's right to call joint_attention_kwargs.

Yes, attention_kwargs should be more appropriate, I have replaced with that.

sayakpaul · 2024-07-30T13:00:26Z

src/diffusers/loaders/lora_pipeline.py

+
+    _lora_loadable_modules = ["transformer"]
+    transformer_name = TRANSFORMER_NAME
+    text_encoder_name = TEXT_ENCODER_NAME


Do we need the text_encoder_name then?

Sorry I missed that, has been removed.

HuggingFaceDocBuilderDev · 2024-07-30T13:09:07Z

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.

Warlord-K · 2024-07-30T13:28:14Z

I have added the tests using the lora I tested, and am building docs via https://github.com/huggingface/diffusers/tree/main/docs. Just to confirm I need to build the docs for the AuraFlowLoRALoaderMixin locally and then push? Please let me know if anything is incorrect with this

sayakpaul · 2024-07-30T13:30:28Z

@Warlord-K you don't have build the docs locally. You just have to add the entry to corresponding loaders doc: https://huggingface.co/docs/diffusers/main/en/api/loaders/lora.

Warlord-K · 2024-07-30T13:31:34Z

Ah, I think then it should be done. Please check

sayakpaul

Some more comments.

sayakpaul · 2024-07-30T13:31:03Z

docs/source/en/api/loaders/lora.md

@@ -17,6 +17,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi
 - [`StableDiffusionLoraLoaderMixin`] provides functions for loading and unloading, fusing and unfusing, enabling and disabling, and more functions for managing LoRA weights. This class can be used with any model.
 - [`StableDiffusionXLLoraLoaderMixin`] is a [Stable Diffusion (SDXL)](../../api/pipelines/stable_diffusion/stable_diffusion_xl) version of the [`StableDiffusionLoraLoaderMixin`] class for loading and saving LoRA weights. It can only be used with the SDXL model.
 - [`SD3LoraLoaderMixin`] provides similar functions for [Stable Diffusion 3](https://huggingface.co/blog/sd3).
+- [`AuraFlowLoraLoaderMixin`] provides similar functions for [AuraFlow](https://huggingface.co/fal/AuraFlow).


This should suffice for the docs.

sayakpaul · 2024-07-30T13:32:00Z

src/diffusers/loaders/lora_pipeline.py

+            safe_serialization=safe_serialization,
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.fuse_lora


Hmm, this wouldn't be a copy because SD3 has more components.

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1416 in e5b94b4

components: List[str] = ["transformer", "text_encoder", "text_encoder_2"],

Oh ok, I am removing the copy line from other functions too which have small changes

sayakpaul · 2024-07-30T13:32:09Z

src/diffusers/loaders/lora_pipeline.py

+            components=components, lora_scale=lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.unfuse_lora


Same as above.

sayakpaul · 2024-07-30T13:33:14Z

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -232,7 +233,7 @@ def forward(
        return encoder_hidden_states, hidden_states


-class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin):
+class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):


I don't think we need FromOriginalModelMixin. No?

Nope Sorry, Removed.

sayakpaul · 2024-07-30T13:33:42Z

tests/lora/test_lora_layers_af.py

+
+
+@require_peft_backend
+class AFLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):


Suggested change

class AFLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):

class AuraFlowLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):

sayakpaul · 2024-07-30T13:34:26Z

tests/lora/test_lora_layers_af.py

+        "sample_size": 64,
+        "patch_size": 2,
+        "in_channels": 4,
+        "num_mmdit_layers": 4,
+        "num_single_dit_layers": 32,
+        "attention_head_dim": 256,
+        "num_attention_heads": 12,
+        "joint_attention_dim": 2048,


These are very big numbers of for fasts. Please consider using significantly smaller numbers as done in SD3 and others.

sayakpaul · 2024-07-30T13:34:41Z

tests/lora/test_lora_layers_af.py

+    vae_kwargs = {
+        "sample_size": 1024,
+        "in_channels": 3,
+        "out_channels": 3,
+        "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+        "layers_per_block": 2,
+        "latent_channels": 4,
+        "norm_num_groups": 32,
+        "use_quant_conv": True,
+        "use_post_quant_conv": True,
+        "shift_factor": None,
+        "scaling_factor": 0.13025,
+    }
+    has_three_text_encoders = False


sayakpaul · 2024-07-30T13:35:26Z

tests/lora/test_lora_layers_af.py

+    vae_kwargs = {
+        "sample_size": 1024,
+        "in_channels": 3,
+        "out_channels": 3,
+        "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+        "layers_per_block": 2,
+        "latent_channels": 4,
+        "norm_num_groups": 32,
+        "use_quant_conv": True,
+        "use_post_quant_conv": True,
+        "shift_factor": None,
+        "scaling_factor": 0.13025,
+    }
+    has_three_text_encoders = False


We don't have to explicitly specify has_three_text_encoders=False as

diffusers/tests/lora/utils.py

Line 75 in e5b94b4

has_two_text_encoders = False

diffusers/tests/lora/utils.py

Line 76 in e5b94b4

has_three_text_encoders = False

sayakpaul · 2024-07-30T13:35:47Z

tests/lora/test_lora_layers_af.py

+    def test_af_lora(self):
+        """
+        Test loading the loras that are saved with the diffusers and peft formats.
+        Related PR: https://github.com/huggingface/diffusers/pull/8584


How is that PR related?

sayakpaul · 2024-07-30T13:36:24Z

tests/lora/test_lora_layers_af.py

+    has_three_text_encoders = False
+
+    @require_torch_gpu
+    def test_af_lora(self):


I think we can safely remove this test.

Warlord-K · 2024-08-11T20:20:59Z

I have made the required changes and added tests and doc mentions. I have tried to follow #9057 for the tests since utils.py was significantly changed to accomodate the newer models but I get all 26 tests skipped for both flux and auraflow when I run them on my laptop. @sayakpaul Please review and let me know If am making any mistake while running the tests.

sayakpaul · 2024-08-12T01:53:38Z

but I get all 26 tests skipped for both flux and auraflow when I run them on my laptop.

Do you have peft installed in the env where you are running this from? We have this constraint:

diffusers/tests/lora/test_lora_layers_flux.py

Line 30 in 98930ee

@require_peft_backend

sayakpaul

Thanks much for the changes. I just left my comments.

sayakpaul · 2024-08-12T01:55:18Z

src/diffusers/loaders/lora_pipeline.py

+
+    @classmethod
+    @validate_hf_hub_args
+    def lora_state_dict(


There should be a "Copied from statement ..." here like:

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1492 in 98930ee

# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict

sayakpaul · 2024-08-12T01:56:47Z

src/diffusers/loaders/lora_pipeline.py

+            components=components, lora_scale=lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict with text_encoder removed from components


"Copied from ..." statements won't work with keywords like "removed from ...". Better to remove.

sayakpaul · 2024-08-12T01:57:50Z

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -434,6 +435,7 @@ def forward(
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor = None,
        timestep: torch.LongTensor = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,


Actually, sorry for my oversight here. We can call it joint_attention_kwargs as that is what we call them in Flux as well.

vladmandic · 2024-09-12T21:41:11Z

any updates here? seems like a great progress made and then no updates for the past month

github-actions · 2024-10-07T15:04:28Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

sayakpaul · 2024-10-07T15:05:14Z

@Warlord-K a gentle ping here.

github-actions · 2024-11-01T15:04:12Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

hameerabbasi · 2024-11-18T08:47:46Z

I can take over here if that's okay, @Warlord-K.

sayakpaul · 2024-11-18T08:53:04Z

Feel free to cherry pick commits! Thanks for offering to help :)

github-actions · 2024-12-12T15:05:04Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

sayakpaul · 2024-12-13T06:17:57Z

@hameerabbasi, are you still interested in picking this up?

hameerabbasi · 2024-12-13T06:46:01Z

Right — I didn’t find the bandwidth, but happy to let others take over.

Edit: On second thought; I can spend a few hours today.

github-actions · 2025-01-06T15:05:48Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

hameerabbasi · 2025-01-06T15:08:21Z

I guess this one can be closed, it's superceded by #10216

github-actions · 2025-02-01T15:05:20Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

Warlord-K and others added 2 commits July 30, 2024 17:22

Add AuraFlowLoraLoaderMixin

e2fe0e2

Merge branch 'main' into main

61e7f46

sayakpaul reviewed Jul 30, 2024

View reviewed changes

Warlord-K added 4 commits July 30, 2024 18:46

Add comments, remove qkv fusion

a21bd03

Add Tests

e8e9571

Add AuraFlowLoraLoaderMixin to documentation

39efb80

Merge branch 'main' of https://github.com/Warlord-K/diffusers

29e8e7f

sayakpaul reviewed Jul 30, 2024

View reviewed changes

Warlord-K and others added 3 commits August 7, 2024 23:29

Merge branch 'main' into main

3f9e464

Merge branch 'huggingface:main' into main

c52dc64

Add Suggested changes

fa32398

Merge branch 'main' into main

1d4f5b9

sayakpaul reviewed Aug 12, 2024

View reviewed changes

Warlord-K and others added 3 commits August 12, 2024 11:06

Change attention_kwargs->joint_attention_kwargs

5f56714

Merge branch 'main' of https://github.com/Warlord-K/diffusers

db45595

Merge branch 'huggingface:main' into main

7d20bd1

github-actions bot added the stale Issues that haven't received updates label Oct 7, 2024

sayakpaul removed the stale Issues that haven't received updates label Oct 7, 2024

github-actions bot added the stale Issues that haven't received updates label Nov 1, 2024

yiyixuxu removed the stale Issues that haven't received updates label Nov 6, 2024

github-actions bot added the stale Issues that haven't received updates label Dec 12, 2024

sayakpaul removed the stale Issues that haven't received updates label Dec 13, 2024

hameerabbasi mentioned this pull request Dec 13, 2024

[LoRA] Add LoRA support to AuraFlow #10216

Merged

6 tasks

github-actions bot added the stale Issues that haven't received updates label Jan 6, 2025

github-actions bot removed the stale Issues that haven't received updates label Jan 7, 2025

github-actions bot added the stale Issues that haven't received updates label Feb 1, 2025

		# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedJointAttnProcessor2_0
		def fuse_qkv_projections(self):



		@require_peft_backend
		class AFLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):

	class AFLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
	class AuraFlowLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):

[LoRA] Add LoRA support to AuraFlow #9017

Are you sure you want to change the base?

[LoRA] Add LoRA support to AuraFlow #9017

Uh oh!

Conversation

Warlord-K commented Jul 30, 2024

What does this PR do?

Before submitting

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Jul 30, 2024

Uh oh!

Warlord-K commented Jul 30, 2024

Uh oh!

sayakpaul commented Jul 30, 2024

Uh oh!

Warlord-K commented Jul 30, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Warlord-K Jul 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Warlord-K commented Aug 11, 2024

Uh oh!

sayakpaul commented Aug 12, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vladmandic commented Sep 12, 2024

Uh oh!

github-actions bot commented Oct 7, 2024

Warlord-K Jul 30, 2024 •

edited

Loading

hameerabbasi commented Dec 13, 2024 •

edited

Loading