| 1261 | |
| 1262 | class SkipFFTransformerBlock(nn.Module): |
| 1263 | def __init__( |
| 1264 | self, |
| 1265 | dim: int, |
| 1266 | num_attention_heads: int, |
| 1267 | attention_head_dim: int, |
| 1268 | kv_input_dim: int, |
| 1269 | kv_input_dim_proj_use_bias: bool, |
| 1270 | dropout=0.0, |
| 1271 | cross_attention_dim: int | None = None, |
| 1272 | attention_bias: bool = False, |
| 1273 | attention_out_bias: bool = True, |
| 1274 | ): |
| 1275 | super().__init__() |
| 1276 | if kv_input_dim != dim: |
| 1277 | self.kv_mapper = nn.Linear(kv_input_dim, dim, kv_input_dim_proj_use_bias) |
| 1278 | else: |
| 1279 | self.kv_mapper = None |
| 1280 | |
| 1281 | self.norm1 = RMSNorm(dim, 1e-06) |
| 1282 | |
| 1283 | self.attn1 = Attention( |
| 1284 | query_dim=dim, |
| 1285 | heads=num_attention_heads, |
| 1286 | dim_head=attention_head_dim, |
| 1287 | dropout=dropout, |
| 1288 | bias=attention_bias, |
| 1289 | cross_attention_dim=cross_attention_dim, |
| 1290 | out_bias=attention_out_bias, |
| 1291 | ) |
| 1292 | |
| 1293 | self.norm2 = RMSNorm(dim, 1e-06) |
| 1294 | |
| 1295 | self.attn2 = Attention( |
| 1296 | query_dim=dim, |
| 1297 | cross_attention_dim=cross_attention_dim, |
| 1298 | heads=num_attention_heads, |
| 1299 | dim_head=attention_head_dim, |
| 1300 | dropout=dropout, |
| 1301 | bias=attention_bias, |
| 1302 | out_bias=attention_out_bias, |
| 1303 | ) |
| 1304 | |
| 1305 | def forward(self, hidden_states, encoder_hidden_states, cross_attention_kwargs): |
| 1306 | cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} |