| 797 | |
| 798 | |
| 799 | class UpBlock3D(nn.Module): |
| 800 | def __init__( |
| 801 | self, |
| 802 | in_channels: int, |
| 803 | prev_output_channel: int, |
| 804 | out_channels: int, |
| 805 | temb_channels: int, |
| 806 | dropout: float = 0.0, |
| 807 | num_layers: int = 1, |
| 808 | resnet_eps: float = 1e-6, |
| 809 | resnet_time_scale_shift: str = "default", |
| 810 | resnet_act_fn: str = "swish", |
| 811 | resnet_groups: int = 32, |
| 812 | resnet_pre_norm: bool = True, |
| 813 | output_scale_factor: float = 1.0, |
| 814 | add_upsample: bool = True, |
| 815 | resolution_idx: Optional[int] = None, |
| 816 | ): |
| 817 | super().__init__() |
| 818 | resnets = [] |
| 819 | temp_convs = [] |
| 820 | |
| 821 | for i in range(num_layers): |
| 822 | res_skip_channels = in_channels if (i == num_layers - 1) else out_channels |
| 823 | resnet_in_channels = prev_output_channel if i == 0 else out_channels |
| 824 | |
| 825 | resnets.append( |
| 826 | ResnetBlock2D( |
| 827 | in_channels=resnet_in_channels + res_skip_channels, |
| 828 | out_channels=out_channels, |
| 829 | temb_channels=temb_channels, |
| 830 | eps=resnet_eps, |
| 831 | groups=resnet_groups, |
| 832 | dropout=dropout, |
| 833 | time_embedding_norm=resnet_time_scale_shift, |
| 834 | non_linearity=resnet_act_fn, |
| 835 | output_scale_factor=output_scale_factor, |
| 836 | pre_norm=resnet_pre_norm, |
| 837 | ) |
| 838 | ) |
| 839 | temp_convs.append( |
| 840 | TemporalConvLayer( |
| 841 | out_channels, |
| 842 | out_channels, |
| 843 | dropout=0.1, |
| 844 | norm_num_groups=resnet_groups, |
| 845 | ) |
| 846 | ) |
| 847 | |
| 848 | self.resnets = nn.ModuleList(resnets) |
| 849 | self.temp_convs = nn.ModuleList(temp_convs) |
| 850 | |
| 851 | if add_upsample: |
| 852 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 853 | else: |
| 854 | self.upsamplers = None |
| 855 | |
| 856 | self.gradient_checkpointing = False |