huggingface
Support out_dim argument for Attention block
Is your feature request related to a problem? Please describe.
When i feed the out_dim argument in __init__ in Attention block it will raise the shape error, because the query_dim != out_dim. In this case, the following code try to keep the given channel of hidden_states.
https://github.com/huggingface/diffusers/blob/b69fd990ad8026f21893499ab396d969b62bb8cc/src/diffusers/models/attention_processor.py#L1393 But it should change the channel as the output of
hidden_states = attn.to_out[0](hidden_states).
Describe the solution you'd like.
I suggest the change of code base : https://github.com/huggingface/diffusers/blob/b69fd990ad8026f21893499ab396d969b62bb8cc/src/diffusers/models/attention_processor.py#L1393
to hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, -1, height, width), then it will respect the channel of hidden_states.
Maybe I will make a PR later.
Describe alternatives you've considered. None.
Additional context. None.
