| 1335 | |
| 1336 | |
| 1337 | class UpBlockSpatioTemporal(nn.Module): |
| 1338 | def __init__( |
| 1339 | self, |
| 1340 | in_channels: int, |
| 1341 | prev_output_channel: int, |
| 1342 | out_channels: int, |
| 1343 | temb_channels: int, |
| 1344 | resolution_idx: Optional[int] = None, |
| 1345 | num_layers: int = 1, |
| 1346 | resnet_eps: float = 1e-6, |
| 1347 | add_upsample: bool = True, |
| 1348 | ): |
| 1349 | super().__init__() |
| 1350 | resnets = [] |
| 1351 | |
| 1352 | for i in range(num_layers): |
| 1353 | res_skip_channels = in_channels if (i == num_layers - 1) else out_channels |
| 1354 | resnet_in_channels = prev_output_channel if i == 0 else out_channels |
| 1355 | |
| 1356 | resnets.append( |
| 1357 | SpatioTemporalResBlock( |
| 1358 | in_channels=resnet_in_channels + res_skip_channels, |
| 1359 | out_channels=out_channels, |
| 1360 | temb_channels=temb_channels, |
| 1361 | eps=resnet_eps, |
| 1362 | ) |
| 1363 | ) |
| 1364 | |
| 1365 | self.resnets = nn.ModuleList(resnets) |
| 1366 | |
| 1367 | if add_upsample: |
| 1368 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 1369 | else: |
| 1370 | self.upsamplers = None |
| 1371 | |
| 1372 | self.gradient_checkpointing = False |
| 1373 | self.resolution_idx = resolution_idx |
| 1374 | |
| 1375 | def forward( |
| 1376 | self, |
| 1377 | hidden_states: torch.Tensor, |
| 1378 | res_hidden_states_tuple: Tuple[torch.Tensor, ...], |
| 1379 | temb: Optional[torch.Tensor] = None, |
| 1380 | image_only_indicator: Optional[torch.Tensor] = None, |
| 1381 | ) -> torch.Tensor: |
| 1382 | for resnet in self.resnets: |
| 1383 | # pop res hidden states |
| 1384 | res_hidden_states = res_hidden_states_tuple[-1] |
| 1385 | res_hidden_states_tuple = res_hidden_states_tuple[:-1] |
| 1386 | |
| 1387 | hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) |
| 1388 | |
| 1389 | if self.training and self.gradient_checkpointing: |
| 1390 | |
| 1391 | def create_custom_forward(module): |
| 1392 | def custom_forward(*inputs): |
| 1393 | return module(*inputs) |
| 1394 | |