| 515 | |
| 516 | |
| 517 | class VanillaTemporalModule(nn.Module): |
| 518 | def __init__( |
| 519 | self, |
| 520 | in_channels, |
| 521 | num_attention_heads = 8, |
| 522 | num_transformer_block = 2, |
| 523 | attention_block_types =( "Temporal_Self", "Temporal_Self" ), |
| 524 | cross_frame_attention_mode = None, |
| 525 | temporal_position_encoding = False, |
| 526 | temporal_position_encoding_max_len = 24, |
| 527 | temporal_attention_dim_div = 1, |
| 528 | zero_initialize = True, |
| 529 | ): |
| 530 | super().__init__() |
| 531 | |
| 532 | self.temporal_transformer = TemporalTransformer3DModel( |
| 533 | in_channels=in_channels, |
| 534 | num_attention_heads=num_attention_heads, |
| 535 | attention_head_dim=in_channels // num_attention_heads // temporal_attention_dim_div, |
| 536 | num_layers=num_transformer_block, |
| 537 | attention_block_types=attention_block_types, |
| 538 | cross_frame_attention_mode=cross_frame_attention_mode, |
| 539 | temporal_position_encoding=temporal_position_encoding, |
| 540 | temporal_position_encoding_max_len=temporal_position_encoding_max_len, |
| 541 | ) |
| 542 | |
| 543 | if zero_initialize: |
| 544 | self.temporal_transformer.proj_out = zero_module(self.temporal_transformer.proj_out) |
| 545 | |
| 546 | def forward(self, input_tensor, temb, encoder_hidden_states, attention_mask=None, anchor_frame_idx=None): |
| 547 | hidden_states = input_tensor |
| 548 | hidden_states = self.temporal_transformer(hidden_states, encoder_hidden_states, attention_mask) |
| 549 | |
| 550 | output = hidden_states |
| 551 | return output |
| 552 | |
| 553 | class TemporalTransformer3DModel(nn.Module): |
| 554 | def __init__( |