| 900 | |
| 901 | |
| 902 | class MidBlockTemporalDecoder(nn.Module): |
| 903 | def __init__( |
| 904 | self, |
| 905 | in_channels: int, |
| 906 | out_channels: int, |
| 907 | attention_head_dim: int = 512, |
| 908 | num_layers: int = 1, |
| 909 | upcast_attention: bool = False, |
| 910 | ): |
| 911 | super().__init__() |
| 912 | |
| 913 | resnets = [] |
| 914 | attentions = [] |
| 915 | for i in range(num_layers): |
| 916 | input_channels = in_channels if i == 0 else out_channels |
| 917 | resnets.append( |
| 918 | SpatioTemporalResBlock( |
| 919 | in_channels=input_channels, |
| 920 | out_channels=out_channels, |
| 921 | temb_channels=None, |
| 922 | eps=1e-6, |
| 923 | temporal_eps=1e-5, |
| 924 | merge_factor=0.0, |
| 925 | merge_strategy="learned", |
| 926 | switch_spatial_to_temporal_mix=True, |
| 927 | ) |
| 928 | ) |
| 929 | |
| 930 | attentions.append( |
| 931 | Attention( |
| 932 | query_dim=in_channels, |
| 933 | heads=in_channels // attention_head_dim, |
| 934 | dim_head=attention_head_dim, |
| 935 | eps=1e-6, |
| 936 | upcast_attention=upcast_attention, |
| 937 | norm_num_groups=32, |
| 938 | bias=True, |
| 939 | residual_connection=True, |
| 940 | ) |
| 941 | ) |
| 942 | |
| 943 | self.attentions = nn.ModuleList(attentions) |
| 944 | self.resnets = nn.ModuleList(resnets) |
| 945 | |
| 946 | def forward( |
| 947 | self, |
| 948 | hidden_states: torch.Tensor, |
| 949 | image_only_indicator: torch.Tensor, |
| 950 | ): |
| 951 | hidden_states = self.resnets[0]( |
| 952 | hidden_states, |
| 953 | image_only_indicator=image_only_indicator, |
| 954 | ) |
| 955 | for resnet, attn in zip(self.resnets[1:], self.attentions): |
| 956 | hidden_states = attn(hidden_states) |
| 957 | hidden_states = resnet( |
| 958 | hidden_states, |
| 959 | image_only_indicator=image_only_indicator, |