| 19 | return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
| 20 | |
| 21 | class TemporalBasicTransformerBlock(nn.Module): |
| 22 | def __init__( |
| 23 | self, |
| 24 | dim: int, |
| 25 | num_attention_heads: int, |
| 26 | attention_head_dim: int, |
| 27 | dropout=0.0, |
| 28 | cross_attention_dim: Optional[int] = None, |
| 29 | activation_fn: str = "geglu", |
| 30 | num_embeds_ada_norm: Optional[int] = None, |
| 31 | attention_bias: bool = False, |
| 32 | only_cross_attention: bool = False, |
| 33 | upcast_attention: bool = False, |
| 34 | unet_use_cross_frame_attention=None, |
| 35 | unet_use_temporal_attention=None, |
| 36 | updown=None, |
| 37 | ): |
| 38 | super().__init__() |
| 39 | self.only_cross_attention = only_cross_attention |
| 40 | self.use_ada_layer_norm = num_embeds_ada_norm is not None |
| 41 | self.unet_use_cross_frame_attention = unet_use_cross_frame_attention |
| 42 | self.unet_use_temporal_attention = unet_use_temporal_attention |
| 43 | |
| 44 | # SC-Attn |
| 45 | self.attn1 = Attention( |
| 46 | query_dim=dim, |
| 47 | heads=num_attention_heads, |
| 48 | dim_head=attention_head_dim, |
| 49 | dropout=dropout, |
| 50 | bias=attention_bias, |
| 51 | upcast_attention=upcast_attention, |
| 52 | updown=updown, |
| 53 | ) |
| 54 | self.norm1 = ( |
| 55 | AdaLayerNorm(dim, num_embeds_ada_norm) |
| 56 | if self.use_ada_layer_norm |
| 57 | else nn.LayerNorm(dim) |
| 58 | ) |
| 59 | |
| 60 | # Cross-Attn |
| 61 | if cross_attention_dim is not None: |
| 62 | self.attn2 = Attention( |
| 63 | query_dim=dim, |
| 64 | cross_attention_dim=cross_attention_dim, |
| 65 | heads=num_attention_heads, |
| 66 | dim_head=attention_head_dim, |
| 67 | dropout=dropout, |
| 68 | bias=attention_bias, |
| 69 | upcast_attention=upcast_attention, |
| 70 | ) |
| 71 | else: |
| 72 | self.attn2 = None |
| 73 | |
| 74 | if cross_attention_dim is not None: |
| 75 | self.norm2 = ( |
| 76 | AdaLayerNorm(dim, num_embeds_ada_norm) |
| 77 | if self.use_ada_layer_norm |
| 78 | else nn.LayerNorm(dim) |