| 575 | |
| 576 | |
| 577 | class UpBlock3D(nn.Module): |
| 578 | def __init__( |
| 579 | self, |
| 580 | in_channels: int, |
| 581 | prev_output_channel: int, |
| 582 | out_channels: int, |
| 583 | temb_channels: int, |
| 584 | dropout: float = 0.0, |
| 585 | num_layers: int = 1, |
| 586 | resnet_eps: float = 1e-6, |
| 587 | resnet_time_scale_shift: str = "default", |
| 588 | resnet_act_fn: str = "swish", |
| 589 | resnet_groups: int = 32, |
| 590 | resnet_pre_norm: bool = True, |
| 591 | output_scale_factor=1.0, |
| 592 | add_upsample=True, |
| 593 | ): |
| 594 | super().__init__() |
| 595 | resnets = [] |
| 596 | |
| 597 | for i in range(num_layers): |
| 598 | res_skip_channels = in_channels if (i == num_layers - 1) else out_channels |
| 599 | resnet_in_channels = prev_output_channel if i == 0 else out_channels |
| 600 | |
| 601 | resnets.append( |
| 602 | ResnetBlock3D( |
| 603 | in_channels=resnet_in_channels + res_skip_channels, |
| 604 | out_channels=out_channels, |
| 605 | temb_channels=temb_channels, |
| 606 | eps=resnet_eps, |
| 607 | groups=resnet_groups, |
| 608 | dropout=dropout, |
| 609 | time_embedding_norm=resnet_time_scale_shift, |
| 610 | non_linearity=resnet_act_fn, |
| 611 | output_scale_factor=output_scale_factor, |
| 612 | pre_norm=resnet_pre_norm, |
| 613 | ) |
| 614 | ) |
| 615 | |
| 616 | self.resnets = nn.ModuleList(resnets) |
| 617 | |
| 618 | if add_upsample: |
| 619 | self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 620 | else: |
| 621 | self.upsamplers = None |
| 622 | |
| 623 | self.gradient_checkpointing = False |
| 624 | |
| 625 | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None): |
| 626 | for resnet in self.resnets: |
| 627 | # pop res hidden states |
| 628 | res_hidden_states = res_hidden_states_tuple[-1] |
| 629 | res_hidden_states_tuple = res_hidden_states_tuple[:-1] |
| 630 | hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) |
| 631 | |
| 632 | if self.training and self.gradient_checkpointing: |
| 633 | |
| 634 | def create_custom_forward(module): |