Lightweight VAE Decoder for SD2 with halved channels. Accepts latent z of shape [B, 4, H/8, W/8] and outputs RGB image [B, 3, H, W].
| 25 | |
| 26 | |
| 27 | class LightDecoder(nn.Module): |
| 28 | """ |
| 29 | Lightweight VAE Decoder for SD2 with halved channels. |
| 30 | |
| 31 | Accepts latent z of shape [B, 4, H/8, W/8] and outputs RGB image [B, 3, H, W]. |
| 32 | """ |
| 33 | |
| 34 | def __init__( |
| 35 | self, |
| 36 | in_channels: int = 4, |
| 37 | out_channels: int = 3, |
| 38 | block_out_channels=(64, 128, 256, 256), |
| 39 | layers_per_block: int = 2, |
| 40 | norm_num_groups: int = 32, |
| 41 | act_fn: str = "silu", |
| 42 | mid_block_add_attention: bool = True, |
| 43 | ): |
| 44 | super().__init__() |
| 45 | |
| 46 | self.in_channels = in_channels |
| 47 | self.out_channels = out_channels |
| 48 | self.block_out_channels = block_out_channels |
| 49 | self.layers_per_block = layers_per_block |
| 50 | |
| 51 | # Reversed block_out_channels for decoder (goes from deepest to shallowest) |
| 52 | reversed_block_out_channels = list(reversed(block_out_channels)) |
| 53 | |
| 54 | # conv_in: latent_channels -> deepest channel dim |
| 55 | self.conv_in = nn.Conv2d( |
| 56 | in_channels, |
| 57 | reversed_block_out_channels[0], |
| 58 | kernel_size=3, |
| 59 | stride=1, |
| 60 | padding=1, |
| 61 | ) |
| 62 | |
| 63 | # Mid block with attention |
| 64 | # attention_head_dim should match in_channels (single-head attention), |
| 65 | # consistent with diffusers Decoder which uses attention_head_dim=block_out_channels[-1] |
| 66 | self.mid_block = UNetMidBlock2D( |
| 67 | in_channels=reversed_block_out_channels[0], |
| 68 | temb_channels=None, |
| 69 | dropout=0.0, |
| 70 | num_layers=1, |
| 71 | resnet_eps=1e-6, |
| 72 | resnet_act_fn=act_fn, |
| 73 | resnet_groups=norm_num_groups, |
| 74 | add_attention=mid_block_add_attention, |
| 75 | attention_head_dim=reversed_block_out_channels[0], |
| 76 | ) |
| 77 | |
| 78 | # Up blocks |
| 79 | self.up_blocks = nn.ModuleList([]) |
| 80 | output_channel = reversed_block_out_channels[0] |
| 81 | for i, up_block_out_channel in enumerate(reversed_block_out_channels): |
| 82 | prev_output_channel = output_channel |
| 83 | output_channel = up_block_out_channel |
| 84 | is_final_block = i == len(reversed_block_out_channels) - 1 |