r""" A basic Transformer block for video like data. Parameters: dim (`int`): The number of channels in the input and output. time_mix_inner_dim (`int`): The number of channels for temporal attention. num_attention_heads (`int`): The number of heads to use for multi-h
| 362 | |
| 363 | @maybe_allow_in_graph |
| 364 | class TemporalBasicTransformerBlock(nn.Module): |
| 365 | r""" |
| 366 | A basic Transformer block for video like data. |
| 367 | |
| 368 | Parameters: |
| 369 | dim (`int`): The number of channels in the input and output. |
| 370 | time_mix_inner_dim (`int`): The number of channels for temporal attention. |
| 371 | num_attention_heads (`int`): The number of heads to use for multi-head attention. |
| 372 | attention_head_dim (`int`): The number of channels in each head. |
| 373 | cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. |
| 374 | """ |
| 375 | |
| 376 | def __init__( |
| 377 | self, |
| 378 | dim: int, |
| 379 | time_mix_inner_dim: int, |
| 380 | num_attention_heads: int, |
| 381 | attention_head_dim: int, |
| 382 | cross_attention_dim: Optional[int] = None, |
| 383 | ): |
| 384 | super().__init__() |
| 385 | self.is_res = dim == time_mix_inner_dim |
| 386 | |
| 387 | self.norm_in = nn.LayerNorm(dim) |
| 388 | |
| 389 | # Define 3 blocks. Each block has its own normalization layer. |
| 390 | # 1. Self-Attn |
| 391 | self.norm_in = nn.LayerNorm(dim) |
| 392 | self.ff_in = FeedForward( |
| 393 | dim, |
| 394 | dim_out=time_mix_inner_dim, |
| 395 | activation_fn="geglu", |
| 396 | ) |
| 397 | |
| 398 | self.norm1 = nn.LayerNorm(time_mix_inner_dim) |
| 399 | self.attn1 = Attention( |
| 400 | query_dim=time_mix_inner_dim, |
| 401 | heads=num_attention_heads, |
| 402 | dim_head=attention_head_dim, |
| 403 | cross_attention_dim=None, |
| 404 | ) |
| 405 | |
| 406 | # 2. Cross-Attn |
| 407 | if cross_attention_dim is not None: |
| 408 | # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. |
| 409 | # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during |
| 410 | # the second cross attention block. |
| 411 | self.norm2 = nn.LayerNorm(time_mix_inner_dim) |
| 412 | self.attn2 = Attention( |
| 413 | query_dim=time_mix_inner_dim, |
| 414 | cross_attention_dim=cross_attention_dim, |
| 415 | heads=num_attention_heads, |
| 416 | dim_head=attention_head_dim, |
| 417 | ) # is self-attn if encoder_hidden_states is none |
| 418 | else: |
| 419 | self.norm2 = None |
| 420 | self.attn2 = None |
| 421 |