| 31 | class BasicTransformerBlock(torch.nn.Module): |
| 32 | |
| 33 | def __init__(self, dim, num_attention_heads, attention_head_dim, cross_attention_dim): |
| 34 | super().__init__() |
| 35 | |
| 36 | # 1. Self-Attn |
| 37 | self.norm1 = torch.nn.LayerNorm(dim, elementwise_affine=True) |
| 38 | self.attn1 = Attention(q_dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, bias_out=True) |
| 39 | |
| 40 | # 2. Cross-Attn |
| 41 | self.norm2 = torch.nn.LayerNorm(dim, elementwise_affine=True) |
| 42 | self.attn2 = Attention(q_dim=dim, kv_dim=cross_attention_dim, num_heads=num_attention_heads, head_dim=attention_head_dim, bias_out=True) |
| 43 | |
| 44 | # 3. Feed-forward |
| 45 | self.norm3 = torch.nn.LayerNorm(dim, elementwise_affine=True) |
| 46 | self.act_fn = GEGLU(dim, dim * 4) |
| 47 | self.ff = torch.nn.Linear(dim * 4, dim) |
| 48 | |
| 49 | |
| 50 | def forward(self, hidden_states, encoder_hidden_states, ipadapter_kwargs=None): |