Pass additional information into Key and Value weights of Stable Diffusion

Hello everyone, I’m working on integrating some code to work with Stable Diffusion. The code is found here (GitHub - lucidrains/perfusion-pytorch: Implementation of Key-Locked Rank One Editing, from Nvidia AI). What I need to do is to replace the Key and Value weights inside each cross attention block within Stable Diffusion. Essentially what it does is wrap each Key Value matrix and do some additional calculations and return the output. This is extremely close to the Custom Diffusion process that you can find here. However, the code I’m using requires additional parameters to be passed in during training/inference. What I would like to do is to pass in the additional information in the call() method of the UNet2DConditionModel class, and then have that be passed into the wrapper. What I’ve thought of so far is using a custom AttnProcessor class, similar to the CustomDiffusionAttnProcessor class. However I can’t figure out how to use these classes to accept the additional parameters that would be passed into the UNet2DConditionModel. Any help would be greatly appreciated.

1 Like

Hey Irow, i am working on the same project let me know if we can contact through mail or or messages on huggingface.