| 183 | |
| 184 | |
| 185 | class TemporalTransformerBlock(nn.Module): |
| 186 | def __init__( |
| 187 | self, |
| 188 | dim, |
| 189 | num_attention_heads, |
| 190 | attention_head_dim, |
| 191 | attention_block_types=( |
| 192 | "Temporal_Self", |
| 193 | "Temporal_Self", |
| 194 | ), |
| 195 | dropout=0.0, |
| 196 | norm_num_groups=32, |
| 197 | cross_attention_dim=768, |
| 198 | activation_fn="geglu", |
| 199 | attention_bias=False, |
| 200 | upcast_attention=False, |
| 201 | cross_frame_attention_mode=None, |
| 202 | temporal_position_encoding=False, |
| 203 | temporal_position_encoding_max_len=24, |
| 204 | ): |
| 205 | super().__init__() |
| 206 | |
| 207 | attention_blocks = [] |
| 208 | norms = [] |
| 209 | |
| 210 | for block_name in attention_block_types: |
| 211 | attention_blocks.append( |
| 212 | VersatileAttention( |
| 213 | attention_mode=block_name.split("_")[0], |
| 214 | cross_attention_dim=cross_attention_dim |
| 215 | if block_name.endswith("_Cross") |
| 216 | else None, |
| 217 | query_dim=dim, |
| 218 | heads=num_attention_heads, |
| 219 | dim_head=attention_head_dim, |
| 220 | dropout=dropout, |
| 221 | bias=attention_bias, |
| 222 | upcast_attention=upcast_attention, |
| 223 | cross_frame_attention_mode=cross_frame_attention_mode, |
| 224 | temporal_position_encoding=temporal_position_encoding, |
| 225 | temporal_position_encoding_max_len=temporal_position_encoding_max_len, |
| 226 | ) |
| 227 | ) |
| 228 | norms.append(nn.LayerNorm(dim)) |
| 229 | |
| 230 | self.attention_blocks = nn.ModuleList(attention_blocks) |
| 231 | self.norms = nn.ModuleList(norms) |
| 232 | |
| 233 | self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn) |
| 234 | self.ff_norm = nn.LayerNorm(dim) |
| 235 | |
| 236 | def forward( |
| 237 | self, |
| 238 | hidden_states, |
| 239 | encoder_hidden_states=None, |
| 240 | attention_mask=None, |
| 241 | video_length=None, |
| 242 | ): |