Describe the bug

Like cross_attention_kwargs in UNet, I want to modify the attention processor of the FLUX model, and pass the extra parameter by the joint_attention_kwargs which written in the FluxPipeline:

noise_pred = self.transformer(
                   hidden_states=latents,
                   timestep=timestep / 1000,
                   guidance=guidance,
                   pooled_projections=pooled_prompt_embeds,
                   encoder_hidden_states=prompt_embeds,
                   txt_ids=text_ids,
                   img_ids=latent_image_ids,
                   joint_attention_kwargs=self.joint_attention_kwargs, # here
                   return_dict=False,
               )[0]

But it doesn't work. I read the resourse code and find that the joint_attention_kwargs is not passed to the inner blocks of the transformer.

https://github.com/huggingface/diffusers/blob/28f9d84549c0b1d24ef00d69a4c723f3a11cffb6/src/diffusers/models/transformers/transformer_flux.py#L495C1-L500C18

we can find that joint_attention_kwargs is missing!

Reproduction

not needed

Logs

No response

System Info

diffusers 0.31.0.dev0

Who can help?

@yiyixuxu