[Core] Add AuraFlow

[sayakpaul]

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]

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]

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]

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

[yiyixuxu]

very nice!
left some comments mainly on attention processor

[sayakpaul]

Looking into the test failures 👀

[bghira]

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]

@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]

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!!

[yiyixuxu]

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

[sayakpaul]

@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?

[yiyixuxu]

you're right, I was confused 😅

[yiyixuxu]

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

[sayakpaul]

attention_mask

[yiyixuxu]

thanks!