| 557 | |
| 558 | |
| 559 | class BasicTransformerBlock(nn.Module): |
| 560 | |
| 561 | def __init__(self, dim, n_heads, d_head, context_dim=None, gated_ff=True): |
| 562 | super().__init__() |
| 563 | self.attn1 = CrossAttention( |
| 564 | query_dim=dim, |
| 565 | heads=n_heads, |
| 566 | dim_head=d_head, |
| 567 | ) |
| 568 | self.ff = FeedForward(dim, glu=gated_ff) |
| 569 | self.attn2 = CrossAttention( |
| 570 | query_dim=dim, |
| 571 | context_dim=context_dim, |
| 572 | heads=n_heads, |
| 573 | dim_head=d_head, |
| 574 | ) |
| 575 | self.norm1 = LayerNormANE(dim) |
| 576 | self.norm2 = LayerNormANE(dim) |
| 577 | self.norm3 = LayerNormANE(dim) |
| 578 | |
| 579 | def forward(self, hidden_states, context=None): |
| 580 | hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states |
| 581 | hidden_states = self.attn2(self.norm2(hidden_states), |
| 582 | context=context) + hidden_states |
| 583 | hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states |
| 584 | return hidden_states |
| 585 | |
| 586 | |
| 587 | class FeedForward(nn.Module): |