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