| 81 | """ |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | in_channels: int = 3, |
| 86 | out_channels: int = 3, |
| 87 | down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",), |
| 88 | block_out_channels: tuple[int, ...] = (64,), |
| 89 | layers_per_block: int = 2, |
| 90 | norm_num_groups: int = 32, |
| 91 | act_fn: str = "silu", |
| 92 | double_z: bool = True, |
| 93 | mid_block_add_attention=True, |
| 94 | ): |
| 95 | super().__init__() |
| 96 | self.layers_per_block = layers_per_block |
| 97 | |
| 98 | self.conv_in = nn.Conv2d( |
| 99 | in_channels, |
| 100 | block_out_channels[0], |
| 101 | kernel_size=3, |
| 102 | stride=1, |
| 103 | padding=1, |
| 104 | ) |
| 105 | |
| 106 | self.down_blocks = nn.ModuleList([]) |
| 107 | |
| 108 | # down |
| 109 | output_channel = block_out_channels[0] |
| 110 | for i, down_block_type in enumerate(down_block_types): |
| 111 | input_channel = output_channel |
| 112 | output_channel = block_out_channels[i] |
| 113 | is_final_block = i == len(block_out_channels) - 1 |
| 114 | |
| 115 | down_block = get_down_block( |
| 116 | down_block_type, |
| 117 | num_layers=self.layers_per_block, |
| 118 | in_channels=input_channel, |
| 119 | out_channels=output_channel, |
| 120 | add_downsample=not is_final_block, |
| 121 | resnet_eps=1e-6, |
| 122 | downsample_padding=0, |
| 123 | resnet_act_fn=act_fn, |
| 124 | resnet_groups=norm_num_groups, |
| 125 | attention_head_dim=output_channel, |
| 126 | temb_channels=None, |
| 127 | ) |
| 128 | self.down_blocks.append(down_block) |
| 129 | |
| 130 | # mid |
| 131 | self.mid_block = UNetMidBlock2D( |
| 132 | in_channels=block_out_channels[-1], |
| 133 | resnet_eps=1e-6, |
| 134 | resnet_act_fn=act_fn, |
| 135 | output_scale_factor=1, |
| 136 | resnet_time_scale_shift="default", |
| 137 | attention_head_dim=block_out_channels[-1], |
| 138 | resnet_groups=norm_num_groups, |
| 139 | temb_channels=None, |
| 140 | add_attention=mid_block_add_attention, |