[SD3] pipe.enable_xformers_memory_efficient_attention

[CanvaChen]

Describe the bug

RuntimeError: The size of tensor a (154) must match the size of tensor b (2304) at non-singleton dimension 1

Reproduction

# StableDiffusion3Pipeline
pipe.enable_xformers_memory_efficient_attention()

Logs

No response

System Info

diffusers==0.29.0
python=3.10
pytorch=2.3

Who can help?

@yiyixuxu @DN6 @sayakpaul

[sayakpaul]

I don’t think we allow xformers attention in the SD3 blocks. Would you be interested in opening a PR? We will be happy to guide you.

[CanvaChen]

It might be a bit challenging for me as my understanding of xformers is currently at the application level.

[sayakpaul]

Oh that is okay. Here are a couple of reference pointers for you:

• This is the standard xformers attention processor class:

  diffusers/src/diffusers/models/attention_processor.py

  Line 1312 in f96e4a1

  class XFormersAttnProcessor:

• So, similarly, you would have to implement one for the joint attention processor block. More specifically, you will have to use ops from xformers in place of the native PyTorch ops here:

  diffusers/src/diffusers/models/attention_processor.py

  Line 1135 in f96e4a1

  hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)

  (making sure the dimensions satisfy the criteria of xformers).

Does this seem like a feature you would be interested in contributing? Not only we would greatly appreciate it but also help you by providing guidance #8276

Let us know.

[CanvaChen]

Thank you for your guidance. I’m interested in attempting this and plan to work on it during my free time after work.

[sayakpaul]

That would be great! As mentioned we will be more than happy to guide you throughout.

[CanvaChen]

1. I have implemented the XFormersJointAttnProcessor, but after calling pipe.enable_xformers_memory_efficient_attention(), self.attn.processor is always set to XFormersAttnProcessor. Where should I configure to select the correct processor?
2. Since the first issue has not been resolved, I did not call pipe.enable_xformers_memory_efficient_attention() and temporarily changed processor = JointAttnProcessor2_0() to processor = XFormersJointAttnProcessor() directly in the JointTransformerBlock for testing. I found that the attention_mask passed to the Processor is None. Why is this happening?

[sayakpaul]

Thanks for your updates!

Do you want to open a PR with your implementation and tag myself and @yiyixuxu there?

It is perfectly okay to have it in an incomplete state.

[CanvaChen]

I’ll submit the PR after I’ve finished testing to make sure there are no issues. Currently, the attention_mask parameter is None, and I’m not sure if this is a problem. Could you help clarify my concerns in the two points above?

[sayakpaul]

I think None mask param value is fine.

Here is an example of how you can set the right processor:

diffusers/src/diffusers/models/attention_processor.py

Line 381 in a899e42

self.set_processor(processor)

[CanvaChen]

I have already opened a PR. Currently, the attention_mask is None, so I am not handling the attention_mask for now.