mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
c6180a311c
* add only cross attention to simple attention blocks * add test for only_cross_attention re: @patrickvonplaten * mid_block_only_cross_attention better default allow mid_block_only_cross_attention to default to `only_cross_attention` when `only_cross_attention` is given as a single boolean
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor
|
|
|
|
|
|
class AttnAddedKVProcessorTests(unittest.TestCase):
|
|
def get_constructor_arguments(self, only_cross_attention: bool = False):
|
|
query_dim = 10
|
|
|
|
if only_cross_attention:
|
|
cross_attention_dim = 12
|
|
else:
|
|
# when only cross attention is not set, the cross attention dim must be the same as the query dim
|
|
cross_attention_dim = query_dim
|
|
|
|
return {
|
|
"query_dim": query_dim,
|
|
"cross_attention_dim": cross_attention_dim,
|
|
"heads": 2,
|
|
"dim_head": 4,
|
|
"added_kv_proj_dim": 6,
|
|
"norm_num_groups": 1,
|
|
"only_cross_attention": only_cross_attention,
|
|
"processor": AttnAddedKVProcessor(),
|
|
}
|
|
|
|
def get_forward_arguments(self, query_dim, added_kv_proj_dim):
|
|
batch_size = 2
|
|
|
|
hidden_states = torch.rand(batch_size, query_dim, 3, 2)
|
|
encoder_hidden_states = torch.rand(batch_size, 4, added_kv_proj_dim)
|
|
attention_mask = None
|
|
|
|
return {
|
|
"hidden_states": hidden_states,
|
|
"encoder_hidden_states": encoder_hidden_states,
|
|
"attention_mask": attention_mask,
|
|
}
|
|
|
|
def test_only_cross_attention(self):
|
|
# self and cross attention
|
|
|
|
torch.manual_seed(0)
|
|
|
|
constructor_args = self.get_constructor_arguments(only_cross_attention=False)
|
|
attn = Attention(**constructor_args)
|
|
|
|
self.assertTrue(attn.to_k is not None)
|
|
self.assertTrue(attn.to_v is not None)
|
|
|
|
forward_args = self.get_forward_arguments(
|
|
query_dim=constructor_args["query_dim"], added_kv_proj_dim=constructor_args["added_kv_proj_dim"]
|
|
)
|
|
|
|
self_and_cross_attn_out = attn(**forward_args)
|
|
|
|
# only self attention
|
|
|
|
torch.manual_seed(0)
|
|
|
|
constructor_args = self.get_constructor_arguments(only_cross_attention=True)
|
|
attn = Attention(**constructor_args)
|
|
|
|
self.assertTrue(attn.to_k is None)
|
|
self.assertTrue(attn.to_v is None)
|
|
|
|
forward_args = self.get_forward_arguments(
|
|
query_dim=constructor_args["query_dim"], added_kv_proj_dim=constructor_args["added_kv_proj_dim"]
|
|
)
|
|
|
|
only_cross_attn_out = attn(**forward_args)
|
|
|
|
self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all())
|