| 1013 | |
| 1014 | |
| 1015 | class UNetMidBlockSpatioTemporal(nn.Module): |
| 1016 | def __init__( |
| 1017 | self, |
| 1018 | in_channels: int, |
| 1019 | temb_channels: int, |
| 1020 | num_layers: int = 1, |
| 1021 | transformer_layers_per_block: Union[int, Tuple[int]] = 1, |
| 1022 | num_attention_heads: int = 1, |
| 1023 | cross_attention_dim: int = 1280, |
| 1024 | ): |
| 1025 | super().__init__() |
| 1026 | |
| 1027 | self.has_cross_attention = True |
| 1028 | self.num_attention_heads = num_attention_heads |
| 1029 | |
| 1030 | # support for variable transformer layers per block |
| 1031 | if isinstance(transformer_layers_per_block, int): |
| 1032 | transformer_layers_per_block = [transformer_layers_per_block] * num_layers |
| 1033 | |
| 1034 | # there is always at least one resnet |
| 1035 | resnets = [ |
| 1036 | SpatioTemporalResBlock( |
| 1037 | in_channels=in_channels, |
| 1038 | out_channels=in_channels, |
| 1039 | temb_channels=temb_channels, |
| 1040 | eps=1e-5, |
| 1041 | ) |
| 1042 | ] |
| 1043 | attentions = [] |
| 1044 | |
| 1045 | for i in range(num_layers): |
| 1046 | attentions.append( |
| 1047 | TransformerSpatioTemporalModel( |
| 1048 | num_attention_heads, |
| 1049 | in_channels // num_attention_heads, |
| 1050 | in_channels=in_channels, |
| 1051 | num_layers=transformer_layers_per_block[i], |
| 1052 | cross_attention_dim=cross_attention_dim, |
| 1053 | ) |
| 1054 | ) |
| 1055 | |
| 1056 | resnets.append( |
| 1057 | SpatioTemporalResBlock( |
| 1058 | in_channels=in_channels, |
| 1059 | out_channels=in_channels, |
| 1060 | temb_channels=temb_channels, |
| 1061 | eps=1e-5, |
| 1062 | ) |
| 1063 | ) |
| 1064 | |
| 1065 | self.attentions = nn.ModuleList(attentions) |
| 1066 | self.resnets = nn.ModuleList(resnets) |
| 1067 | |
| 1068 | self.gradient_checkpointing = False |
| 1069 | |
| 1070 | def forward( |
| 1071 | self, |
| 1072 | hidden_states: torch.Tensor, |