| 299 | |
| 300 | class TemporalBasicTransformerBlock(nn.Module): |
| 301 | def __init__( |
| 302 | self, |
| 303 | dim: int, |
| 304 | num_attention_heads: int, |
| 305 | attention_head_dim: int, |
| 306 | dropout=0.0, |
| 307 | cross_attention_dim: Optional[int] = None, |
| 308 | activation_fn: str = "geglu", |
| 309 | num_embeds_ada_norm: Optional[int] = None, |
| 310 | attention_bias: bool = False, |
| 311 | only_cross_attention: bool = False, |
| 312 | upcast_attention: bool = False, |
| 313 | unet_use_cross_frame_attention=None, |
| 314 | unet_use_temporal_attention=None, |
| 315 | name=None, |
| 316 | ): |
| 317 | super().__init__() |
| 318 | self.only_cross_attention = only_cross_attention |
| 319 | self.use_ada_layer_norm = num_embeds_ada_norm is not None |
| 320 | self.unet_use_cross_frame_attention = unet_use_cross_frame_attention |
| 321 | self.unet_use_temporal_attention = unet_use_temporal_attention |
| 322 | self.name=name |
| 323 | |
| 324 | # SC-Attn |
| 325 | self.attn1 = Attention( |
| 326 | query_dim=dim, |
| 327 | heads=num_attention_heads, |
| 328 | dim_head=attention_head_dim, |
| 329 | dropout=dropout, |
| 330 | bias=attention_bias, |
| 331 | upcast_attention=upcast_attention, |
| 332 | ) |
| 333 | self.norm1 = ( |
| 334 | AdaLayerNorm(dim, num_embeds_ada_norm) |
| 335 | if self.use_ada_layer_norm |
| 336 | else nn.LayerNorm(dim) |
| 337 | ) |
| 338 | |
| 339 | # Cross-Attn |
| 340 | if cross_attention_dim is not None: |
| 341 | self.attn2 = Attention( |
| 342 | query_dim=dim, |
| 343 | cross_attention_dim=cross_attention_dim, |
| 344 | heads=num_attention_heads, |
| 345 | dim_head=attention_head_dim, |
| 346 | dropout=dropout, |
| 347 | bias=attention_bias, |
| 348 | upcast_attention=upcast_attention, |
| 349 | ) |
| 350 | else: |
| 351 | self.attn2 = None |
| 352 | |
| 353 | if cross_attention_dim is not None: |
| 354 | self.norm2 = ( |
| 355 | AdaLayerNorm(dim, num_embeds_ada_norm) |
| 356 | if self.use_ada_layer_norm |
| 357 | else nn.LayerNorm(dim) |
| 358 | ) |