| 363 | |
| 364 | |
| 365 | class DownBlock3D(nn.Module): |
| 366 | def __init__( |
| 367 | self, |
| 368 | in_channels: int, |
| 369 | out_channels: int, |
| 370 | temb_channels: int, |
| 371 | dropout: float = 0.0, |
| 372 | num_layers: int = 1, |
| 373 | resnet_eps: float = 1e-6, |
| 374 | resnet_time_scale_shift: str = "default", |
| 375 | resnet_act_fn: str = "swish", |
| 376 | resnet_groups: int = 32, |
| 377 | resnet_pre_norm: bool = True, |
| 378 | output_scale_factor=1.0, |
| 379 | add_downsample=True, |
| 380 | downsample_padding=1, |
| 381 | ): |
| 382 | super().__init__() |
| 383 | resnets = [] |
| 384 | |
| 385 | for i in range(num_layers): |
| 386 | in_channels = in_channels if i == 0 else out_channels |
| 387 | resnets.append( |
| 388 | ResnetBlock3D( |
| 389 | in_channels=in_channels, |
| 390 | out_channels=out_channels, |
| 391 | temb_channels=temb_channels, |
| 392 | eps=resnet_eps, |
| 393 | groups=resnet_groups, |
| 394 | dropout=dropout, |
| 395 | time_embedding_norm=resnet_time_scale_shift, |
| 396 | non_linearity=resnet_act_fn, |
| 397 | output_scale_factor=output_scale_factor, |
| 398 | pre_norm=resnet_pre_norm, |
| 399 | ) |
| 400 | ) |
| 401 | |
| 402 | self.resnets = nn.ModuleList(resnets) |
| 403 | |
| 404 | if add_downsample: |
| 405 | self.downsamplers = nn.ModuleList( |
| 406 | [ |
| 407 | Downsample3D( |
| 408 | out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" |
| 409 | ) |
| 410 | ] |
| 411 | ) |
| 412 | else: |
| 413 | self.downsamplers = None |
| 414 | |
| 415 | self.gradient_checkpointing = False |
| 416 | |
| 417 | def forward(self, hidden_states, temb=None): |
| 418 | output_states = () |
| 419 | |
| 420 | for resnet in self.resnets: |
| 421 | if self.training and self.gradient_checkpointing: |
| 422 | |