| 1123 | |
| 1124 | |
| 1125 | class DownBlockSpatioTemporal(nn.Module): |
| 1126 | def __init__( |
| 1127 | self, |
| 1128 | in_channels: int, |
| 1129 | out_channels: int, |
| 1130 | temb_channels: int, |
| 1131 | num_layers: int = 1, |
| 1132 | add_downsample: bool = True, |
| 1133 | ): |
| 1134 | super().__init__() |
| 1135 | resnets = [] |
| 1136 | |
| 1137 | for i in range(num_layers): |
| 1138 | in_channels = in_channels if i == 0 else out_channels |
| 1139 | resnets.append( |
| 1140 | SpatioTemporalResBlock( |
| 1141 | in_channels=in_channels, |
| 1142 | out_channels=out_channels, |
| 1143 | temb_channels=temb_channels, |
| 1144 | eps=1e-5, |
| 1145 | ) |
| 1146 | ) |
| 1147 | |
| 1148 | self.resnets = nn.ModuleList(resnets) |
| 1149 | |
| 1150 | if add_downsample: |
| 1151 | self.downsamplers = nn.ModuleList( |
| 1152 | [ |
| 1153 | Downsample2D( |
| 1154 | out_channels, |
| 1155 | use_conv=True, |
| 1156 | out_channels=out_channels, |
| 1157 | name="op", |
| 1158 | ) |
| 1159 | ] |
| 1160 | ) |
| 1161 | else: |
| 1162 | self.downsamplers = None |
| 1163 | |
| 1164 | self.gradient_checkpointing = False |
| 1165 | |
| 1166 | def forward( |
| 1167 | self, |
| 1168 | hidden_states: torch.Tensor, |
| 1169 | temb: Optional[torch.Tensor] = None, |
| 1170 | image_only_indicator: Optional[torch.Tensor] = None, |
| 1171 | ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: |
| 1172 | output_states = () |
| 1173 | for resnet in self.resnets: |
| 1174 | if self.training and self.gradient_checkpointing: |
| 1175 | |
| 1176 | def create_custom_forward(module): |
| 1177 | def custom_forward(*inputs): |
| 1178 | return module(*inputs) |
| 1179 | |
| 1180 | return custom_forward |
| 1181 | |
| 1182 | if is_torch_version(">=", "1.11.0"): |