| 54 | |
| 55 | # 1024 * 1024 * 3 -> 16 * 16 * 512 -> 1024 * 1024 * 3 |
| 56 | class UNet1024(ModelMixin, ConfigMixin): |
| 57 | @register_to_config |
| 58 | def __init__( |
| 59 | self, |
| 60 | in_channels: int = 3, |
| 61 | out_channels: int = 3, |
| 62 | down_block_types: Tuple[str] = ("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"), |
| 63 | up_block_types: Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"), |
| 64 | block_out_channels: Tuple[int] = (32, 32, 64, 128, 256, 512, 512), |
| 65 | layers_per_block: int = 2, |
| 66 | mid_block_scale_factor: float = 1, |
| 67 | downsample_padding: int = 1, |
| 68 | downsample_type: str = "conv", |
| 69 | upsample_type: str = "conv", |
| 70 | dropout: float = 0.0, |
| 71 | act_fn: str = "silu", |
| 72 | attention_head_dim: Optional[int] = 8, |
| 73 | norm_num_groups: int = 4, |
| 74 | norm_eps: float = 1e-5, |
| 75 | latent_c: int = 4, |
| 76 | ): |
| 77 | super().__init__() |
| 78 | |
| 79 | # input |
| 80 | self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)) |
| 81 | self.latent_conv_in = zero_module(nn.Conv2d(latent_c, block_out_channels[2], kernel_size=1)) |
| 82 | |
| 83 | self.down_blocks = nn.ModuleList([]) |
| 84 | self.mid_block = None |
| 85 | self.up_blocks = nn.ModuleList([]) |
| 86 | |
| 87 | # down |
| 88 | output_channel = block_out_channels[0] |
| 89 | for i, down_block_type in enumerate(down_block_types): |
| 90 | input_channel = output_channel |
| 91 | output_channel = block_out_channels[i] |
| 92 | is_final_block = i == len(block_out_channels) - 1 |
| 93 | |
| 94 | down_block = get_down_block( |
| 95 | down_block_type, |
| 96 | num_layers=layers_per_block, |
| 97 | in_channels=input_channel, |
| 98 | out_channels=output_channel, |
| 99 | temb_channels=None, |
| 100 | add_downsample=not is_final_block, |
| 101 | resnet_eps=norm_eps, |
| 102 | resnet_act_fn=act_fn, |
| 103 | resnet_groups=norm_num_groups, |
| 104 | attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel, |
| 105 | downsample_padding=downsample_padding, |
| 106 | resnet_time_scale_shift="default", |
| 107 | downsample_type=downsample_type, |
| 108 | dropout=dropout, |
| 109 | ) |
| 110 | self.down_blocks.append(down_block) |
| 111 | |
| 112 | # mid |
| 113 | self.mid_block = UNetMidBlock2D( |