r""" The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample. Args: in_channels (`int`, *optional*, defaults to 3): The number of input channels. out_channels (`int`, *optional*, defaults to 3): Th
| 183 | |
| 184 | |
| 185 | class Decoder(nn.Module): |
| 186 | r""" |
| 187 | The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample. |
| 188 | |
| 189 | Args: |
| 190 | in_channels (`int`, *optional*, defaults to 3): |
| 191 | The number of input channels. |
| 192 | out_channels (`int`, *optional*, defaults to 3): |
| 193 | The number of output channels. |
| 194 | up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): |
| 195 | The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. |
| 196 | block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): |
| 197 | The number of output channels for each block. |
| 198 | layers_per_block (`int`, *optional*, defaults to 2): |
| 199 | The number of layers per block. |
| 200 | norm_num_groups (`int`, *optional*, defaults to 32): |
| 201 | The number of groups for normalization. |
| 202 | act_fn (`str`, *optional*, defaults to `"silu"`): |
| 203 | The activation function to use. See `~diffusers.models.activations.get_activation` for available options. |
| 204 | norm_type (`str`, *optional*, defaults to `"group"`): |
| 205 | The normalization type to use. Can be either `"group"` or `"spatial"`. |
| 206 | """ |
| 207 | |
| 208 | def __init__( |
| 209 | self, |
| 210 | in_channels: int = 3, |
| 211 | out_channels: int = 3, |
| 212 | up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), |
| 213 | block_out_channels: Tuple[int, ...] = (64,), |
| 214 | layers_per_block: int = 2, |
| 215 | norm_num_groups: int = 32, |
| 216 | act_fn: str = "silu", |
| 217 | norm_type: str = "group", # group, spatial |
| 218 | mid_block_add_attention=True, |
| 219 | ): |
| 220 | super().__init__() |
| 221 | self.layers_per_block = layers_per_block |
| 222 | |
| 223 | self.conv_in = nn.Conv2d( |
| 224 | in_channels, |
| 225 | block_out_channels[-1], |
| 226 | kernel_size=3, |
| 227 | stride=1, |
| 228 | padding=1, |
| 229 | ) |
| 230 | |
| 231 | self.mid_block = None |
| 232 | self.up_blocks = nn.ModuleList([]) |
| 233 | |
| 234 | temb_channels = in_channels if norm_type == "spatial" else None |
| 235 | |
| 236 | # mid |
| 237 | self.mid_block = UNetMidBlock2D( |
| 238 | in_channels=block_out_channels[-1], |
| 239 | resnet_eps=1e-6, |
| 240 | resnet_act_fn=act_fn, |
| 241 | output_scale_factor=1, |
| 242 | resnet_time_scale_shift="default" if norm_type == "group" else norm_type, |