r""" The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation. Args: in_channels (`int`, *optional*, defaults to 3): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number
| 44 | |
| 45 | |
| 46 | class Encoder(nn.Module): |
| 47 | r""" |
| 48 | The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation. |
| 49 | |
| 50 | Args: |
| 51 | in_channels (`int`, *optional*, defaults to 3): |
| 52 | The number of input channels. |
| 53 | out_channels (`int`, *optional*, defaults to 3): |
| 54 | The number of output channels. |
| 55 | down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): |
| 56 | The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available |
| 57 | options. |
| 58 | block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): |
| 59 | The number of output channels for each block. |
| 60 | layers_per_block (`int`, *optional*, defaults to 2): |
| 61 | The number of layers per block. |
| 62 | norm_num_groups (`int`, *optional*, defaults to 32): |
| 63 | The number of groups for normalization. |
| 64 | act_fn (`str`, *optional*, defaults to `"silu"`): |
| 65 | The activation function to use. See `~diffusers.models.activations.get_activation` for available options. |
| 66 | double_z (`bool`, *optional*, defaults to `True`): |
| 67 | Whether to double the number of output channels for the last block. |
| 68 | """ |
| 69 | |
| 70 | def __init__( |
| 71 | self, |
| 72 | in_channels: int = 3, |
| 73 | out_channels: int = 3, |
| 74 | down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), |
| 75 | block_out_channels: Tuple[int, ...] = (64,), |
| 76 | layers_per_block: int = 2, |
| 77 | norm_num_groups: int = 32, |
| 78 | act_fn: str = "silu", |
| 79 | double_z: bool = True, |
| 80 | mid_block_add_attention=True, |
| 81 | ): |
| 82 | super().__init__() |
| 83 | self.layers_per_block = layers_per_block |
| 84 | |
| 85 | self.conv_in = nn.Conv2d( |
| 86 | in_channels, |
| 87 | block_out_channels[0], |
| 88 | kernel_size=3, |
| 89 | stride=1, |
| 90 | padding=1, |
| 91 | ) |
| 92 | |
| 93 | self.mid_block = None |
| 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( |