Better mechanism for getting values to attention processors

[spezialspezial]

I think it would be nice to generally have some better infrastructure/mechanism for getting values to attention processors.

At the moment the main means of transportation seems to be cross_attention_kwargs but they are finally expanded as function arguments and have to match ALL function signatures that receive them. That's why there is a lot of dict gymnastics before the call. Just further passing the dict via reference where every processor may pick what it needs not bothering with other items would be more flexible imo.

In my mind pipeline and processors should be the malleable 'loose ends' to modify and swap with infrastructure in between rarely in need of adjustment.

[sayakpaul]

Great discussion start! Do you have any API design blocks in mind?

Cc: @patrickvonplaten @DN6

[spezialspezial]

Other than the mentioned I'm not versed enough to talk about unet architecture but a pre-unet-evaluation callback inside the inference loop would probably help not having to modify all pipelines in some cases. BTW I really liked the callback_kwargs[k] = locals()[k] construct for the step_end_callback. Very elegant solution covering most needs. Maybe if the output would be generalized too it could be easily reused in other places.

[sayakpaul]

@yiyixuxu WDYT?

[asomoza]

I would also like to support this issue, for example for IP Adapters we can just pass the image embeds in the cross_attention_kwargs and manage them inside the IPAttentionProcessor but you seem to be against this since for this to happen you need to enable the pass of the **kwargs in the attention processors, but in this comment it clearly states that you're not in favor of this:

#6573 (comment)

So the implementation right now is that you concatenate the the image embeds with the encoder_hidden_states like this:

image_embeds = added_cond_kwargs.get("image_embeds")
image_embeds = self.encoder_hid_proj(image_embeds).to(encoder_hidden_states.dtype)
encoder_hidden_states = torch.cat([encoder_hidden_states, image_embeds], dim=1)

and then in the attention processor, you revert it to get the image_embeds:

# split hidden states
end_pos = encoder_hidden_states.shape[1] - self.num_tokens
encoder_hidden_states, ip_hidden_states = (
    encoder_hidden_states[:, :end_pos, :],
    encoder_hidden_states[:, end_pos:, :],
)

this forces me to monkey patch the unet forward method and implement my own version of the AttentionProcesors just to be able to manage the image_embeds.

[asomoza]

I prefer to write here instead of the IP Adapters PR since this is more of a general design issue and not specific to IP Adapters, but now I understand the reasoning behind not using **kwargs with the comment from @patrickvonplaten in #6573 (comment). For me it works whatever method you decide is best for diffusers that allows to pass values directly to the attention processors.

Most systems that reuse the inputs for multiple generations, like UIs or node systems (a good node system doesn't run nodes that haven't changed), can have a noticeable speed up if we just pass the values, in this case ip_hidden_states, directly to the attention processors without the dict (or tensor) gymnastics as @spezialspezial mentioned.

P.S.: The recent refactor of the unet made it so I don't have to monkeypatch the forward method anymore so thanks for that.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

I think this can be safely closed?

[yiyixuxu]

fixed in #6915