r""" A VQ-VAE model for decoding latent representations. This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). Parameters: in_channels (int, *optional*, defaults to 3)
| 38 | |
| 39 | |
| 40 | class VQModel(ModelMixin, ConfigMixin): |
| 41 | r""" |
| 42 | A VQ-VAE model for decoding latent representations. |
| 43 | |
| 44 | This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented |
| 45 | for all models (such as downloading or saving). |
| 46 | |
| 47 | Parameters: |
| 48 | in_channels (int, *optional*, defaults to 3): Number of channels in the input image. |
| 49 | out_channels (int, *optional*, defaults to 3): Number of channels in the output. |
| 50 | down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): |
| 51 | Tuple of downsample block types. |
| 52 | up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): |
| 53 | Tuple of upsample block types. |
| 54 | block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): |
| 55 | Tuple of block output channels. |
| 56 | layers_per_block (`int`, *optional*, defaults to `1`): Number of layers per block. |
| 57 | act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. |
| 58 | latent_channels (`int`, *optional*, defaults to `3`): Number of channels in the latent space. |
| 59 | sample_size (`int`, *optional*, defaults to `32`): Sample input size. |
| 60 | num_vq_embeddings (`int`, *optional*, defaults to `256`): Number of codebook vectors in the VQ-VAE. |
| 61 | norm_num_groups (`int`, *optional*, defaults to `32`): Number of groups for normalization layers. |
| 62 | vq_embed_dim (`int`, *optional*): Hidden dim of codebook vectors in the VQ-VAE. |
| 63 | scaling_factor (`float`, *optional*, defaults to `0.18215`): |
| 64 | The component-wise standard deviation of the trained latent space computed using the first batch of the |
| 65 | training set. This is used to scale the latent space to have unit variance when training the diffusion |
| 66 | model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the |
| 67 | diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 |
| 68 | / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image |
| 69 | Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper. |
| 70 | norm_type (`str`, *optional*, defaults to `"group"`): |
| 71 | Type of normalization layer to use. Can be one of `"group"` or `"spatial"`. |
| 72 | """ |
| 73 | |
| 74 | @register_to_config |
| 75 | def __init__( |
| 76 | self, |
| 77 | in_channels: int = 3, |
| 78 | out_channels: int = 3, |
| 79 | down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), |
| 80 | up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), |
| 81 | block_out_channels: Tuple[int, ...] = (64,), |
| 82 | layers_per_block: int = 1, |
| 83 | act_fn: str = "silu", |
| 84 | latent_channels: int = 3, |
| 85 | sample_size: int = 32, |
| 86 | num_vq_embeddings: int = 256, |
| 87 | norm_num_groups: int = 32, |
| 88 | vq_embed_dim: Optional[int] = None, |
| 89 | scaling_factor: float = 0.18215, |
| 90 | norm_type: str = "group", # group, spatial |
| 91 | mid_block_add_attention=True, |
| 92 | lookup_from_codebook=False, |
| 93 | force_upcast=False, |
| 94 | ): |
| 95 | super().__init__() |
| 96 | |
| 97 | # pass init params to Encoder |
no outgoing calls