(
self,
in_channels=3,
out_channels=3,
up_block_types=("UpDecoderBlock2D",),
block_out_channels=(64,),
layers_per_block=2,
norm_num_groups=32,
act_fn="silu",
norm_type="group", # group, spatial
)
| 151 | |
| 152 | class Decoder(nn.Module): |
| 153 | def __init__( |
| 154 | self, |
| 155 | in_channels=3, |
| 156 | out_channels=3, |
| 157 | up_block_types=("UpDecoderBlock2D",), |
| 158 | block_out_channels=(64,), |
| 159 | layers_per_block=2, |
| 160 | norm_num_groups=32, |
| 161 | act_fn="silu", |
| 162 | norm_type="group", # group, spatial |
| 163 | ): |
| 164 | super().__init__() |
| 165 | self.layers_per_block = layers_per_block |
| 166 | |
| 167 | self.conv_in = nn.Conv2d( |
| 168 | in_channels, |
| 169 | block_out_channels[-1], |
| 170 | kernel_size=3, |
| 171 | stride=1, |
| 172 | padding=1, |
| 173 | ) |
| 174 | |
| 175 | self.mid_block = None |
| 176 | self.up_blocks = nn.ModuleList([]) |
| 177 | |
| 178 | temb_channels = in_channels if norm_type == "spatial" else None |
| 179 | |
| 180 | # mid |
| 181 | self.mid_block = UNetMidBlock2D( |
| 182 | in_channels=block_out_channels[-1], |
| 183 | resnet_eps=1e-6, |
| 184 | resnet_act_fn=act_fn, |
| 185 | output_scale_factor=1, |
| 186 | resnet_time_scale_shift="default" if norm_type == "group" else norm_type, |
| 187 | attn_num_head_channels=None, |
| 188 | resnet_groups=norm_num_groups, |
| 189 | temb_channels=temb_channels, |
| 190 | ) |
| 191 | |
| 192 | # up |
| 193 | reversed_block_out_channels = list(reversed(block_out_channels)) |
| 194 | output_channel = reversed_block_out_channels[0] |
| 195 | for i, up_block_type in enumerate(up_block_types): |
| 196 | prev_output_channel = output_channel |
| 197 | output_channel = reversed_block_out_channels[i] |
| 198 | |
| 199 | is_final_block = i == len(block_out_channels) - 1 |
| 200 | |
| 201 | up_block = get_up_block( |
| 202 | up_block_type, |
| 203 | num_layers=self.layers_per_block + 1, |
| 204 | in_channels=prev_output_channel, |
| 205 | out_channels=output_channel, |
| 206 | prev_output_channel=None, |
| 207 | add_upsample=not is_final_block, |
| 208 | resnet_eps=1e-6, |
| 209 | resnet_act_fn=act_fn, |
| 210 | resnet_groups=norm_num_groups, |
nothing calls this directly
no test coverage detected