| 41 | |
| 42 | |
| 43 | class VanillaTemporalModule(nn.Module): |
| 44 | def __init__( |
| 45 | self, |
| 46 | in_channels, |
| 47 | num_attention_heads=8, |
| 48 | num_transformer_block=2, |
| 49 | attention_block_types=("Temporal_Self", "Temporal_Self"), |
| 50 | cross_frame_attention_mode=None, |
| 51 | temporal_position_encoding=False, |
| 52 | temporal_position_encoding_max_len=24, |
| 53 | temporal_attention_dim_div=1, |
| 54 | zero_initialize=True, |
| 55 | ): |
| 56 | super().__init__() |
| 57 | |
| 58 | self.temporal_transformer = TemporalTransformer3DModel( |
| 59 | in_channels=in_channels, |
| 60 | num_attention_heads=num_attention_heads, |
| 61 | attention_head_dim=in_channels |
| 62 | // num_attention_heads |
| 63 | // temporal_attention_dim_div, |
| 64 | num_layers=num_transformer_block, |
| 65 | attention_block_types=attention_block_types, |
| 66 | cross_frame_attention_mode=cross_frame_attention_mode, |
| 67 | temporal_position_encoding=temporal_position_encoding, |
| 68 | temporal_position_encoding_max_len=temporal_position_encoding_max_len, |
| 69 | ) |
| 70 | |
| 71 | if zero_initialize: |
| 72 | self.temporal_transformer.proj_out = zero_module( |
| 73 | self.temporal_transformer.proj_out |
| 74 | ) |
| 75 | |
| 76 | def forward( |
| 77 | self, |
| 78 | input_tensor, |
| 79 | temb, |
| 80 | encoder_hidden_states, |
| 81 | attention_mask=None, |
| 82 | anchor_frame_idx=None, |
| 83 | ): |
| 84 | hidden_states = input_tensor |
| 85 | hidden_states = self.temporal_transformer( |
| 86 | hidden_states, encoder_hidden_states, attention_mask |
| 87 | ) |
| 88 | |
| 89 | output = hidden_states |
| 90 | return output |
| 91 | |
| 92 | |
| 93 | class TemporalTransformer3DModel(nn.Module): |