(
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
)
| 453 | """ |
| 454 | |
| 455 | def __init__( |
| 456 | self, |
| 457 | in_channels: int = 3, |
| 458 | out_channels: int = 3, |
| 459 | up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), |
| 460 | block_out_channels: Tuple[int, ...] = (64,), |
| 461 | layers_per_block: int = 2, |
| 462 | norm_num_groups: int = 32, |
| 463 | act_fn: str = "silu", |
| 464 | norm_type: str = "group", # group, spatial |
| 465 | ): |
| 466 | super().__init__() |
| 467 | self.layers_per_block = layers_per_block |
| 468 | |
| 469 | self.conv_in = nn.Conv2d( |
| 470 | in_channels, |
| 471 | block_out_channels[-1], |
| 472 | kernel_size=3, |
| 473 | stride=1, |
| 474 | padding=1, |
| 475 | ) |
| 476 | |
| 477 | self.mid_block = None |
| 478 | self.up_blocks = nn.ModuleList([]) |
| 479 | |
| 480 | temb_channels = in_channels if norm_type == "spatial" else None |
| 481 | |
| 482 | # mid |
| 483 | self.mid_block = UNetMidBlock2D( |
| 484 | in_channels=block_out_channels[-1], |
| 485 | resnet_eps=1e-6, |
| 486 | resnet_act_fn=act_fn, |
| 487 | output_scale_factor=1, |
| 488 | resnet_time_scale_shift="default" if norm_type == "group" else norm_type, |
| 489 | attention_head_dim=block_out_channels[-1], |
| 490 | resnet_groups=norm_num_groups, |
| 491 | temb_channels=temb_channels, |
| 492 | ) |
| 493 | |
| 494 | # up |
| 495 | reversed_block_out_channels = list(reversed(block_out_channels)) |
| 496 | output_channel = reversed_block_out_channels[0] |
| 497 | for i, up_block_type in enumerate(up_block_types): |
| 498 | prev_output_channel = output_channel |
| 499 | output_channel = reversed_block_out_channels[i] |
| 500 | |
| 501 | is_final_block = i == len(block_out_channels) - 1 |
| 502 | |
| 503 | up_block = get_up_block( |
| 504 | up_block_type, |
| 505 | num_layers=self.layers_per_block + 1, |
| 506 | in_channels=prev_output_channel, |
| 507 | out_channels=output_channel, |
| 508 | prev_output_channel=None, |
| 509 | add_upsample=not is_final_block, |
| 510 | resnet_eps=1e-6, |
| 511 | resnet_act_fn=act_fn, |
| 512 | resnet_groups=norm_num_groups, |
nothing calls this directly
no test coverage detected