(
self,
in_channels: int,
out_channels: int,
resolution_idx: Optional[int] = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default", # default, spatial
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
resnet_pre_norm: bool = True,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
temb_channels: Optional[int] = None,
)
| 2769 | |
| 2770 | class UpDecoderBlock2D(nn.Module): |
| 2771 | def __init__( |
| 2772 | self, |
| 2773 | in_channels: int, |
| 2774 | out_channels: int, |
| 2775 | resolution_idx: Optional[int] = None, |
| 2776 | dropout: float = 0.0, |
| 2777 | num_layers: int = 1, |
| 2778 | resnet_eps: float = 1e-6, |
| 2779 | resnet_time_scale_shift: str = "default", # default, spatial |
| 2780 | resnet_act_fn: str = "swish", |
| 2781 | resnet_groups: int = 32, |
| 2782 | resnet_pre_norm: bool = True, |
| 2783 | output_scale_factor: float = 1.0, |
| 2784 | add_upsample: bool = True, |
| 2785 | temb_channels: Optional[int] = None, |
| 2786 | ): |
| 2787 | super().__init__() |
| 2788 | resnets = [] |
| 2789 | |
| 2790 | for i in range(num_layers): |
| 2791 | input_channels = in_channels if i == 0 else out_channels |
| 2792 | |
| 2793 | if resnet_time_scale_shift == "spatial": |
| 2794 | resnets.append( |
| 2795 | ResnetBlockCondNorm2D( |
| 2796 | in_channels=input_channels, |
| 2797 | out_channels=out_channels, |
| 2798 | temb_channels=temb_channels, |
| 2799 | eps=resnet_eps, |
| 2800 | groups=resnet_groups, |
| 2801 | dropout=dropout, |
| 2802 | time_embedding_norm="spatial", |
| 2803 | non_linearity=resnet_act_fn, |
| 2804 | output_scale_factor=output_scale_factor, |
| 2805 | ) |
| 2806 | ) |
| 2807 | else: |
| 2808 | resnets.append( |
| 2809 | ResnetBlock2D( |
| 2810 | in_channels=input_channels, |
| 2811 | out_channels=out_channels, |
| 2812 | temb_channels=temb_channels, |
| 2813 | eps=resnet_eps, |
| 2814 | groups=resnet_groups, |
| 2815 | dropout=dropout, |
| 2816 | time_embedding_norm=resnet_time_scale_shift, |
| 2817 | non_linearity=resnet_act_fn, |
| 2818 | output_scale_factor=output_scale_factor, |
| 2819 | pre_norm=resnet_pre_norm, |
| 2820 | ) |
| 2821 | ) |
| 2822 | |
| 2823 | self.resnets = nn.ModuleList(resnets) |
| 2824 | |
| 2825 | if add_upsample: |
| 2826 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 2827 | else: |
| 2828 | self.upsamplers = None |
nothing calls this directly
no test coverage detected