(
self,
in_channels: int = 3,
out_channels: int = 3,
up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
block_out_channels: Tuple[int, ...] = (64,),
layers_per_block: int = 2,
norm_num_groups: int = 32,
act_fn: str = "silu",
norm_type: str = "group", # group, spatial
mid_block_add_attention=True,
)
| 206 | """ |
| 207 | |
| 208 | def __init__( |
| 209 | self, |
| 210 | in_channels: int = 3, |
| 211 | out_channels: int = 3, |
| 212 | up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), |
| 213 | block_out_channels: Tuple[int, ...] = (64,), |
| 214 | layers_per_block: int = 2, |
| 215 | norm_num_groups: int = 32, |
| 216 | act_fn: str = "silu", |
| 217 | norm_type: str = "group", # group, spatial |
| 218 | mid_block_add_attention=True, |
| 219 | ): |
| 220 | super().__init__() |
| 221 | self.layers_per_block = layers_per_block |
| 222 | |
| 223 | self.conv_in = nn.Conv2d( |
| 224 | in_channels, |
| 225 | block_out_channels[-1], |
| 226 | kernel_size=3, |
| 227 | stride=1, |
| 228 | padding=1, |
| 229 | ) |
| 230 | |
| 231 | self.mid_block = None |
| 232 | self.up_blocks = nn.ModuleList([]) |
| 233 | |
| 234 | temb_channels = in_channels if norm_type == "spatial" else None |
| 235 | |
| 236 | # mid |
| 237 | self.mid_block = UNetMidBlock2D( |
| 238 | in_channels=block_out_channels[-1], |
| 239 | resnet_eps=1e-6, |
| 240 | resnet_act_fn=act_fn, |
| 241 | output_scale_factor=1, |
| 242 | resnet_time_scale_shift="default" if norm_type == "group" else norm_type, |
| 243 | attention_head_dim=block_out_channels[-1], |
| 244 | resnet_groups=norm_num_groups, |
| 245 | temb_channels=temb_channels, |
| 246 | add_attention=mid_block_add_attention, |
| 247 | ) |
| 248 | |
| 249 | # up |
| 250 | reversed_block_out_channels = list(reversed(block_out_channels)) |
| 251 | output_channel = reversed_block_out_channels[0] |
| 252 | for i, up_block_type in enumerate(up_block_types): |
| 253 | prev_output_channel = output_channel |
| 254 | output_channel = reversed_block_out_channels[i] |
| 255 | |
| 256 | is_final_block = i == len(block_out_channels) - 1 |
| 257 | |
| 258 | up_block = get_up_block( |
| 259 | up_block_type, |
| 260 | num_layers=self.layers_per_block + 1, |
| 261 | in_channels=prev_output_channel, |
| 262 | out_channels=output_channel, |
| 263 | prev_output_channel=None, |
| 264 | add_upsample=not is_final_block, |
| 265 | resnet_eps=1e-6, |
nothing calls this directly
no test coverage detected