Flux followup

[yiyixuxu]

to-do

• refactor rotary embedding
  • removed Flux-specific methods/classes to use rotary embedding: i.e. EmbedND, apply_rope, rope;
  • created a FluxPosEmbed that uses diffusers' existing get_1d_rotary_pos_embed method, which is already used by hunyuan dit, lumina and stable audio
  • changed the flux transformer inputs img_ids and txt_ids: currently, these are 3d tensors with a batch dimension; I changed them to 2d since these are just positional ids that are used to encode rotary embedding, so we do not need to add a batch dimension here, and get_1d_rotary_pos_embed method does not accept batched positional ids so we remove it here to be consistent across the library - note that this is a breaking change, so I made sure to deprecate it, also add a test to make sure the previous inputs will still work
• refactor attention processor (combine into one, deprecated FluxSingleAttnProcessor2_0)

# flux unit test for rotary embedding refactor import torch from diffusers import FluxPipeline model_path = "black-forest-labs/FLUX.1-dev" pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16) pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power prompt = "A cat holding a sign that says hello world" image = pipe( prompt, height=1024, width=1024, guidance_scale=3.5, num_inference_steps=50, max_sequence_length=512, generator=torch.Generator("cpu").manual_seed(0), ).images[0] image.save(f"yiyi_test_4_out{branch}.png")

main	this PR

[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.

[wangqixun]

Could you please inform when the related PR for the 'flux transformer block' and 'pipeline' can be completed and merged? We are currently working on the adaptation and training of controlnet and other related plugins on the old version. After this PR is completed, we will reorganize the code to adapt to the new coding style, and then we will submit a new PR for controlnet and ipadapter.

prompt = "A girl with green long hair. she is wearing a yellow suit. half body, background is sky and cloud. anime style"

[yiyixuxu]

@wangqixun let's not wait for this refactor to be done for the PR! we can refactor the ip-adapter and controlnet together once the PR is in

[sayakpaul]

Superb!

So, IIUC, the batched txt_ids and img_ids are not at a necessity now because of how we're doing the RoPE in the refactored class (FluxPosEmbed)? Or are there any additional differences to be aware of?

[sayakpaul]

Let's maybe deprecate it earlier? Not a strong opinion, though.

[sayakpaul]

Maybe a reference to the original BFL inference code?

[yiyixuxu]

@sayakpaul the results for flux are identical with this refactor
the only other difference is here, where we downcast the dtype for mps see #9133 for more details

[sayakpaul]

Aye, thanks!

[DN6]

Nice 👍🏽

[DN6]

LGTM 👍🏽