[Core] Add AuraFlow #8796

sayakpaul · 2024-07-05T08:33:56Z

What does this PR do?

Adds Aura Flow from Fal.

Test code:

from diffusers import AuraFlowPipeline
from diffusers.utils import make_image_grid
import torch


pipeline = AuraFlowPipeline(
	"AuraDiffusion/auradiffusion-v0.1a0", 
	torch_dtype=torch.float16
).to("cuda")

images = pipeline(
    prompt="a cute cat with tiger like looks",
    height=512,
    width=512,
    num_inference_steps=50, 
    num_images_per_prompt=4,
    generator=torch.Generator().manual_seed(666),
    guidance_scale=3.5,
).images
make_image_grid(images, 1, 4).save("demo_hf.png")

Warning

To download the model you must be a member of the AuraDiffusion org. Follow this (internal) Slack message.

Gives:

TODOS

Docs
Tests
Scheduler (@yiyixuxu would you be able to help out here? I couldn't find a way to use our existing flow matching scheduler in this case)

~~Because of the last point above, the noise scheduling code is taken from the original codebase. But I think this PR is still ready for a first review.~~

HuggingFaceDocBuilderDev · 2024-07-05T08:45:39Z

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.

sayakpaul · 2024-07-05T08:46:55Z

src/diffusers/models/transformers/lavender_transformer_2d.py

@@ -0,0 +1,401 @@
+# Copyright 2024 Stability AI, Lavender Flow, The HuggingFace Team. All rights reserved.


New model because it differs from the SD3 one (non-exhaustive list):

Uses register tokens

Mixes MMDiT and another kind of simple DiT block (that uses a concatenated encoder_hidden_states and hidden_states as its inputs)

The final layer norm is different

Position embeddings are different (uses learned positional embeddings)

The feedforward is different. We only support GeLU and its variants in the feedforward. It uses SwiGLU.

No pooled projections.

sayakpaul · 2024-07-05T08:49:15Z

src/diffusers/models/transformers/lavender_transformer_2d.py

+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+


I have deliberately kept additional methods like feedforward chunking, QKV fusion, etc. out of this class because it helps with the initial review.

src/diffusers/pipelines/lavender_flow/pipeline_lavender_flow.py

yiyixuxu

very nice!
left some comments mainly on attention processor

src/diffusers/models/attention_processor.py

sayakpaul · 2024-07-06T02:45:35Z

Looking into the test failures 👀

bghira · 2024-07-08T18:43:55Z

i've been testing a fork of this with the LoRA support and it works without any changes to just add the peft adapter to the Transformer2D model and the SD3 LoRA loader mixin to the pipeline.

sayakpaul · 2024-07-09T04:12:07Z

@yiyixuxu @DN6 I have addressed the comments. Here are some extended comments from my end:

Run conversion again because of changes needed to address [Core] Add AuraFlow #8796 (comment). Here's the Hub PR: https://huggingface.co/AuraDiffusion/auradiffusion-v0.1a0/discussions/3.
I think the transformer optimizations can be tackled later in a separate PR. We saw for SD3 that chunked feedforward doesn't lead to a lot of memory saving. So, I would like to first see how much that gain is and then add accordingly.
https://github.com/huggingface/diffusers/actions/runs/9850568702/job/27195993322?pr=8796#step:5:173 (circular import for FP32 layer norm).

@bghira, I will add LoRA support in an immediate future PR once this PR is merged to keep the reviewing scope concrete and manageable. It's not just about adding those classes like you mentioned. We need to scale and unscale the layers appropriately for dealing with scale, add features like fuse_lora(), etc. So, keeping that out of this PR.

sayakpaul · 2024-07-09T05:04:13Z

I have also added fast tests and decided to make the default value of negative prompt to be None instead of "This is watermark, jpeg image white background, web image". I think this is better aligned with our other pipelines. Will include this negative prompt in the docs once I start adding them.

yiyixuxu

very nice!!

src/diffusers/models/normalization.py

src/diffusers/models/transformers/auraflow_transformer_2d.py

yiyixuxu · 2024-07-09T16:23:51Z

src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

@@ -158,7 +158,12 @@ def scale_noise(
    def _sigma_to_t(self, sigma):
        return sigma * self.config.num_train_timesteps

-    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+    def set_timesteps(


umm I don't think these changed are introduced in this PR

@yiyixuxu we merged #8799 into the current PR branch. So, the commits will come here. But they still belong to you, so I guess that is okay?

you're right, I was confused 😅

yiyixuxu · 2024-07-09T16:25:51Z

src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

+                padding="max_length",
+                return_tensors="pt",
+            )
+            text_inputs = {k: v.to(device) for k, v in text_inputs.items()}


just curious what else is in text_inputs other than the text_input_ids?

attention_mask

src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

yiyixuxu

thanks!

* add lavender flow transformer --------- Co-authored-by: YiYi Xu <yixu310@gmail.com>

sayakpaul added 11 commits July 4, 2024 14:34

add lavender flow transformer

0ebbd30

progress.

939d990

progress

f08baf3

progress

005a8f6

move out the attention processor.

e238d56

finish implementation of pipeline

570c258

default neg promot

b881190

up

b8237b2

fixes

89eea61

up

2c97d04

up for pr

a50e1ff

sayakpaul requested review from DN6 and yiyixuxu July 5, 2024 08:33

sayakpaul added 2 commits July 5, 2024 14:05

fix copies

b0d29b2

Merge branch 'main' into lavender-flow

ae037cf

move fp32 layer norm to normalization

ad6cb66

sayakpaul commented Jul 5, 2024

View reviewed changes

minor fixes

8ae6be7

sayakpaul commented Jul 5, 2024

View reviewed changes

src/diffusers/pipelines/lavender_flow/pipeline_lavender_flow.py Outdated Show resolved Hide resolved

yiyixuxu reviewed Jul 5, 2024

View reviewed changes

src/diffusers/models/attention_processor.py Outdated Show resolved Hide resolved

src/diffusers/models/attention_processor.py Outdated Show resolved Hide resolved

src/diffusers/models/attention_processor.py Outdated Show resolved Hide resolved

sayakpaul added 3 commits July 5, 2024 15:29

remove boolean flag and resort to norm_type

47ff911

eliminate added_qk_norm

10ed96f

add added_proj_bias

3d9265e

This was referenced Jul 5, 2024

fix loading sharded checkpoints from subfolder #8798

Merged

[lavender-flow] use flow match euler scheduler #8799

Merged

lavender flow -> aura flow

84708c4

sayakpaul mentioned this pull request Jul 6, 2024

Fix the added_proj_bias default value #8800

Merged

sayakpaul added 5 commits July 9, 2024 05:22

resolve conflicts

d9a01f4

more feedback.

8984d23

context_norm_type fix

a281547

fix circular import

8830bf1

fix conversion

4334f72

sayakpaul added 5 commits July 9, 2024 06:55

add fast tests for pipeline

b1dc5ec

fix test file name

942377d

fix test path

f8a08b5

spacxing brtween initialization

e9832f9

style

2c87250

sayakpaul marked this pull request as ready for review July 9, 2024 05:00

add test for the transformer model.

1b3e620

sayakpaul requested review from DN6 and yiyixuxu July 9, 2024 06:59

sayakpaul added 2 commits July 9, 2024 17:54

Merge branch 'main' into lavender-flow

66ff7f5

remove context_norm_type

ed33913

yiyixuxu reviewed Jul 9, 2024

View reviewed changes

sayakpaul added 2 commits July 10, 2024 06:42

remove ada continuous.

6531e54

address yiyi

0f721ac

yiyixuxu approved these changes Jul 10, 2024

View reviewed changes

sayakpaul mentioned this pull request Jul 10, 2024

Stable Audio integration #8716

Merged

5 tasks

yiyixuxu added 2 commits July 11, 2024 08:33

Merge branch 'main' into lavender-flow

95708dc

style

15d3198

yiyixuxu merged commit 2261510 into main Jul 11, 2024
16 of 18 checks passed

yiyixuxu deleted the lavender-flow branch July 11, 2024 18:50

a-r-r-o-w mentioned this pull request Jul 22, 2024

Adding Differential Diffusion to Kolors, Auraflow, HunyuanDiT #8924

Closed

3 tasks

sayakpaul added a commit that referenced this pull request Dec 23, 2024

[Core] Add AuraFlow (#8796)

824ed25

* add lavender flow transformer --------- Co-authored-by: YiYi Xu <yixu310@gmail.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Core] Add AuraFlow #8796

[Core] Add AuraFlow #8796

sayakpaul commented Jul 5, 2024 •

edited

Loading

HuggingFaceDocBuilderDev commented Jul 5, 2024

sayakpaul Jul 5, 2024 •

edited

Loading

sayakpaul Jul 5, 2024

yiyixuxu left a comment

sayakpaul commented Jul 6, 2024

bghira commented Jul 8, 2024

sayakpaul commented Jul 9, 2024 •

edited

Loading

sayakpaul commented Jul 9, 2024

yiyixuxu left a comment

yiyixuxu Jul 9, 2024

sayakpaul Jul 10, 2024

yiyixuxu Jul 10, 2024

yiyixuxu Jul 9, 2024

sayakpaul Jul 10, 2024

yiyixuxu left a comment

		@@ -0,0 +1,401 @@
		# Copyright 2024 Stability AI, Lavender Flow, The HuggingFace Team. All rights reserved.

[Core] Add AuraFlow #8796

[Core] Add AuraFlow #8796

Conversation

sayakpaul commented Jul 5, 2024 • edited Loading

What does this PR do?

TODOS

HuggingFaceDocBuilderDev commented Jul 5, 2024

sayakpaul Jul 5, 2024 • edited Loading

Choose a reason for hiding this comment

sayakpaul Jul 5, 2024

Choose a reason for hiding this comment

yiyixuxu left a comment

Choose a reason for hiding this comment

sayakpaul commented Jul 6, 2024

bghira commented Jul 8, 2024

sayakpaul commented Jul 9, 2024 • edited Loading

sayakpaul commented Jul 9, 2024

yiyixuxu left a comment

Choose a reason for hiding this comment

yiyixuxu Jul 9, 2024

Choose a reason for hiding this comment

sayakpaul Jul 10, 2024

Choose a reason for hiding this comment

yiyixuxu Jul 10, 2024

Choose a reason for hiding this comment

yiyixuxu Jul 9, 2024

Choose a reason for hiding this comment

sayakpaul Jul 10, 2024

Choose a reason for hiding this comment

yiyixuxu left a comment

Choose a reason for hiding this comment

sayakpaul commented Jul 5, 2024 •

edited

Loading

sayakpaul Jul 5, 2024 •

edited

Loading

sayakpaul commented Jul 9, 2024 •

edited

Loading