tensor shape mismatching for SD3 basic pipeline with xformers

[shenzhiy21]

Hi, I'm new to diffusers and I tried to enable xformers for SD3 inference speedup.

import torch
from diffusers import StableDiffusion3Pipeline

path = "/mnt/data"

pipe = StableDiffusion3Pipeline.from_pretrained(path, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

pipe.enable_xformers_memory_efficient_attention()

image = pipe(
    "A cat holding a sign that says Hello World",
    negative_prompt="",
    num_inference_steps=28,
    guidance_scale=7.0,
).images[0]

image.save("./output/x.png")

The error msg:

Traceback (most recent call last):
  File "/home/superbench/zhishen/playground/sd3.py", line 11, in <module>
    image = pipe(
            ^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py", line 877, in __call__
    noise_pred = self.transformer(
                 ^^^^^^^^^^^^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/diffusers/models/transformers/transformer_sd3.py", line 335, in forward
    encoder_hidden_states, hidden_states = block(
                                           ^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/superbench/anaconda3/envs/comfy/lib/python3.12/site-packages/diffusers/models/attention.py", line 197, in forward
    encoder_hidden_states = encoder_hidden_states + context_attn_output
                            ~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (333) must match the size of tensor b (4096) at non-singleton dimension 1

After removing this line

pipe.enable_xformers_memory_efficient_attention()

everything would work perfectly as expected.

I'm not sure if xformers doesn't support SD3 models. Thanks for helping me with this!

[a-r-r-o-w]

I think we will be removing xformers support going forward in a future release, so support for it is not guaranteed to work out-of-the-box with newer models. This does look xformers-specific so I suspect it is due to us not having a SD3-specific xformers attention processor. It should be quite simple to create your own and use it by modifying the code here. Once you have created the attention processor, you can set it using: pipe.transformer.set_attn_processor(MyCustomXformersSD3AttnProcessor())

[shenzhiy21]

@a-r-r-o-w Thank you so much!

May I further ask that, what's the meaning of "SD3-specific xformers attention processor"? (Or what's the difference for the attention mechanism between SD3 and previous versions like SD1.5?)