| 722 | |
| 723 | |
| 724 | class SkipFFTransformerBlock(nn.Module): |
| 725 | def __init__( |
| 726 | self, |
| 727 | dim: int, |
| 728 | num_attention_heads: int, |
| 729 | attention_head_dim: int, |
| 730 | kv_input_dim: int, |
| 731 | kv_input_dim_proj_use_bias: bool, |
| 732 | dropout=0.0, |
| 733 | cross_attention_dim: Optional[int] = None, |
| 734 | attention_bias: bool = False, |
| 735 | attention_out_bias: bool = True, |
| 736 | ): |
| 737 | super().__init__() |
| 738 | if kv_input_dim != dim: |
| 739 | self.kv_mapper = nn.Linear(kv_input_dim, dim, kv_input_dim_proj_use_bias) |
| 740 | else: |
| 741 | self.kv_mapper = None |
| 742 | |
| 743 | self.norm1 = RMSNorm(dim, 1e-06) |
| 744 | |
| 745 | self.attn1 = Attention( |
| 746 | query_dim=dim, |
| 747 | heads=num_attention_heads, |
| 748 | dim_head=attention_head_dim, |
| 749 | dropout=dropout, |
| 750 | bias=attention_bias, |
| 751 | cross_attention_dim=cross_attention_dim, |
| 752 | out_bias=attention_out_bias, |
| 753 | ) |
| 754 | |
| 755 | self.norm2 = RMSNorm(dim, 1e-06) |
| 756 | |
| 757 | self.attn2 = Attention( |
| 758 | query_dim=dim, |
| 759 | cross_attention_dim=cross_attention_dim, |
| 760 | heads=num_attention_heads, |
| 761 | dim_head=attention_head_dim, |
| 762 | dropout=dropout, |
| 763 | bias=attention_bias, |
| 764 | out_bias=attention_out_bias, |
| 765 | ) |
| 766 | |
| 767 | def forward(self, hidden_states, encoder_hidden_states, cross_attention_kwargs): |
| 768 | cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} |
| 769 | |
| 770 | if self.kv_mapper is not None: |
| 771 | encoder_hidden_states = self.kv_mapper(F.silu(encoder_hidden_states)) |
| 772 | |
| 773 | norm_hidden_states = self.norm1(hidden_states) |
| 774 | |
| 775 | attn_output = self.attn1( |
| 776 | norm_hidden_states, |
| 777 | encoder_hidden_states=encoder_hidden_states, |
| 778 | **cross_attention_kwargs, |
| 779 | ) |
| 780 | |
| 781 | hidden_states = attn_output + hidden_states |