| 13 | |
| 14 | |
| 15 | class Attention(torch.nn.Module): |
| 16 | |
| 17 | def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False): |
| 18 | super().__init__() |
| 19 | dim_inner = head_dim * num_heads |
| 20 | kv_dim = kv_dim if kv_dim is not None else q_dim |
| 21 | self.num_heads = num_heads |
| 22 | self.head_dim = head_dim |
| 23 | |
| 24 | self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q) |
| 25 | self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) |
| 26 | self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv) |
| 27 | self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out) |
| 28 | |
| 29 | def interact_with_ipadapter(self, hidden_states, q, ip_k, ip_v, scale=1.0): |
| 30 | batch_size = q.shape[0] |
| 31 | ip_k = ip_k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) |
| 32 | ip_v = ip_v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) |
| 33 | ip_hidden_states = torch.nn.functional.scaled_dot_product_attention(q, ip_k, ip_v) |
| 34 | hidden_states = hidden_states + scale * ip_hidden_states |
| 35 | return hidden_states |
| 36 | |
| 37 | def torch_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None): |
| 38 | if encoder_hidden_states is None: |
| 39 | encoder_hidden_states = hidden_states |
| 40 | |
| 41 | batch_size = encoder_hidden_states.shape[0] |
| 42 | |
| 43 | q = self.to_q(hidden_states) |
| 44 | k = self.to_k(encoder_hidden_states) |
| 45 | v = self.to_v(encoder_hidden_states) |
| 46 | |
| 47 | q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) |
| 48 | k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) |
| 49 | v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) |
| 50 | |
| 51 | if qkv_preprocessor is not None: |
| 52 | q, k, v = qkv_preprocessor(q, k, v) |
| 53 | |
| 54 | hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) |
| 55 | if ipadapter_kwargs is not None: |
| 56 | hidden_states = self.interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs) |
| 57 | hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) |
| 58 | hidden_states = hidden_states.to(q.dtype) |
| 59 | |
| 60 | hidden_states = self.to_out(hidden_states) |
| 61 | |
| 62 | return hidden_states |
| 63 | |
| 64 | def xformers_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None): |
| 65 | if encoder_hidden_states is None: |
| 66 | encoder_hidden_states = hidden_states |
| 67 | |
| 68 | q = self.to_q(hidden_states) |
| 69 | k = self.to_k(encoder_hidden_states) |
| 70 | v = self.to_v(encoder_hidden_states) |
| 71 | |
| 72 | q = rearrange(q, "b f (n d) -> (b n) f d", n=self.num_heads) |