declare-lab/tango2
92
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from dataclasses import dataclass15from typing import Optional, Tuple, Union16 17import torch18import torch.nn as nn19 20from ..configuration_utils import ConfigMixin, register_to_config21from ..utils import BaseOutput, apply_forward_hook22from .modeling_utils import ModelMixin23from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder24 25 26@dataclass27class AutoencoderKLOutput(BaseOutput):28 """29 Output of AutoencoderKL encoding method.30 31 Args:32 latent_dist (`DiagonalGaussianDistribution`):33 Encoded outputs of `Encoder` represented as the mean and logvar of `DiagonalGaussianDistribution`.34 `DiagonalGaussianDistribution` allows for sampling latents from the distribution.35 """36 37 latent_dist: "DiagonalGaussianDistribution"38 39 40class AutoencoderKL(ModelMixin, ConfigMixin):41 r"""Variational Autoencoder (VAE) model with KL loss from the paper Auto-Encoding Variational Bayes by Diederik P. Kingma42 and Max Welling.43 44 This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library45 implements for all the model (such as downloading or saving, etc.)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 :51 obj:`("DownEncoderBlock2D",)`): Tuple of downsample block types.52 up_block_types (`Tuple[str]`, *optional*, defaults to :53 obj:`("UpDecoderBlock2D",)`): Tuple of upsample block types.54 block_out_channels (`Tuple[int]`, *optional*, defaults to :55 obj:`(64,)`): Tuple of block output channels.56 act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.57 latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.58 sample_size (`int`, *optional*, defaults to `32`): TODO59 scaling_factor (`float`, *optional*, defaults to 0.18215):60 The component-wise standard deviation of the trained latent space computed using the first batch of the61 training set. This is used to scale the latent space to have unit variance when training the diffusion62 model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the63 diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 164 / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image65 Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.66 """67 68 _supports_gradient_checkpointing = True69 70 @register_to_config71 def __init__(72 self,73 in_channels: int = 3,74 out_channels: int = 3,75 down_block_types: Tuple[str] = ("DownEncoderBlock2D",),76 up_block_types: Tuple[str] = ("UpDecoderBlock2D",),77 block_out_channels: Tuple[int] = (64,),78 layers_per_block: int = 1,79 act_fn: str = "silu",80 latent_channels: int = 4,81 norm_num_groups: int = 32,82 sample_size: int = 32,83 scaling_factor: float = 0.18215,84 ):85 super().__init__()86 87 # pass init params to Encoder88 self.encoder = Encoder(89 in_channels=in_channels,90 out_channels=latent_channels,91 down_block_types=down_block_types,92 block_out_channels=block_out_channels,93 layers_per_block=layers_per_block,94 act_fn=act_fn,95 norm_num_groups=norm_num_groups,96 double_z=True,97 )98 99 # pass init params to Decoder100 self.decoder = Decoder(101 in_channels=latent_channels,102 out_channels=out_channels,103 up_block_types=up_block_types,104 block_out_channels=block_out_channels,105 layers_per_block=layers_per_block,106 norm_num_groups=norm_num_groups,107 act_fn=act_fn,108 )109 110 self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)111 self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1)112 113 self.use_slicing = False114 self.use_tiling = False115 116 # only relevant if vae tiling is enabled117 self.tile_sample_min_size = self.config.sample_size118 sample_size = (119 self.config.sample_size[0]120 if isinstance(self.config.sample_size, (list, tuple))121 else self.config.sample_size122 )123 self.tile_latent_min_size = int(sample_size / (2 ** (len(self.block_out_channels) - 1)))124 self.tile_overlap_factor = 0.25125 126 def _set_gradient_checkpointing(self, module, value=False):127 if isinstance(module, (Encoder, Decoder)):128 module.gradient_checkpointing = value129 130 def enable_tiling(self, use_tiling: bool = True):131 r"""132 Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to133 compute decoding and encoding in several steps. This is useful to save a large amount of memory and to allow134 the processing of larger images.135 """136 self.use_tiling = use_tiling137 138 def disable_tiling(self):139 r"""140 Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to141 computing decoding in one step.142 """143 self.enable_tiling(False)144 145 def enable_slicing(self):146 r"""147 Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to148 compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.149 """150 self.use_slicing = True151 152 def disable_slicing(self):153 r"""154 Disable sliced VAE decoding. If `enable_slicing` was previously invoked, this method will go back to computing155 decoding in one step.156 """157 self.use_slicing = False158 159 @apply_forward_hook160 def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:161 if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):162 return self.tiled_encode(x, return_dict=return_dict)163 164 h = self.encoder(x)165 moments = self.quant_conv(h)166 posterior = DiagonalGaussianDistribution(moments)167 168 if not return_dict:169 return (posterior,)170 171 return AutoencoderKLOutput(latent_dist=posterior)172 173 def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:174 if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):175 return self.tiled_decode(z, return_dict=return_dict)176 177 z = self.post_quant_conv(z)178 dec = self.decoder(z)179 180 if not return_dict:181 return (dec,)182 183 return DecoderOutput(sample=dec)184 185 @apply_forward_hook186 def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:187 if self.use_slicing and z.shape[0] > 1:188 decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]189 decoded = torch.cat(decoded_slices)190 else:191 decoded = self._decode(z).sample192 193 if not return_dict:194 return (decoded,)195 196 return DecoderOutput(sample=decoded)197 198 def blend_v(self, a, b, blend_extent):199 for y in range(min(a.shape[2], b.shape[2], blend_extent)):200 b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)201 return b202 203 def blend_h(self, a, b, blend_extent):204 for x in range(min(a.shape[3], b.shape[3], blend_extent)):205 b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)206 return b207 208 def tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:209 r"""Encode a batch of images using a tiled encoder.210 211 Args:212 When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several213 steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is:214 different from non-tiled encoding due to each tile using a different encoder. To avoid tiling artifacts, the215 tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the216 look of the output, but they should be much less noticeable.217 x (`torch.FloatTensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`):218 Whether or not to return a [`AutoencoderKLOutput`] instead of a plain tuple.219 """220 overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))221 blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)222 row_limit = self.tile_latent_min_size - blend_extent223 224 # Split the image into 512x512 tiles and encode them separately.225 rows = []226 for i in range(0, x.shape[2], overlap_size):227 row = []228 for j in range(0, x.shape[3], overlap_size):229 tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]230 tile = self.encoder(tile)231 tile = self.quant_conv(tile)232 row.append(tile)233 rows.append(row)234 result_rows = []235 for i, row in enumerate(rows):236 result_row = []237 for j, tile in enumerate(row):238 # blend the above tile and the left tile239 # to the current tile and add the current tile to the result row240 if i > 0:241 tile = self.blend_v(rows[i - 1][j], tile, blend_extent)242 if j > 0:243 tile = self.blend_h(row[j - 1], tile, blend_extent)244 result_row.append(tile[:, :, :row_limit, :row_limit])245 result_rows.append(torch.cat(result_row, dim=3))246 247 moments = torch.cat(result_rows, dim=2)248 posterior = DiagonalGaussianDistribution(moments)249 250 if not return_dict:251 return (posterior,)252 253 return AutoencoderKLOutput(latent_dist=posterior)254 255 def tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:256 r"""Decode a batch of images using a tiled decoder.257 258 Args:259 When this option is enabled, the VAE will split the input tensor into tiles to compute decoding in several260 steps. This is useful to keep memory use constant regardless of image size. The end result of tiled decoding is:261 different from non-tiled decoding due to each tile using a different decoder. To avoid tiling artifacts, the262 tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the263 look of the output, but they should be much less noticeable.264 z (`torch.FloatTensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to265 `True`):266 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.267 """268 overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))269 blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)270 row_limit = self.tile_sample_min_size - blend_extent271 272 # Split z into overlapping 64x64 tiles and decode them separately.273 # The tiles have an overlap to avoid seams between tiles.274 rows = []275 for i in range(0, z.shape[2], overlap_size):276 row = []277 for j in range(0, z.shape[3], overlap_size):278 tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]279 tile = self.post_quant_conv(tile)280 decoded = self.decoder(tile)281 row.append(decoded)282 rows.append(row)283 result_rows = []284 for i, row in enumerate(rows):285 result_row = []286 for j, tile in enumerate(row):287 # blend the above tile and the left tile288 # to the current tile and add the current tile to the result row289 if i > 0:290 tile = self.blend_v(rows[i - 1][j], tile, blend_extent)291 if j > 0:292 tile = self.blend_h(row[j - 1], tile, blend_extent)293 result_row.append(tile[:, :, :row_limit, :row_limit])294 result_rows.append(torch.cat(result_row, dim=3))295 296 dec = torch.cat(result_rows, dim=2)297 if not return_dict:298 return (dec,)299 300 return DecoderOutput(sample=dec)301 302 def forward(303 self,304 sample: torch.FloatTensor,305 sample_posterior: bool = False,306 return_dict: bool = True,307 generator: Optional[torch.Generator] = None,308 ) -> Union[DecoderOutput, torch.FloatTensor]:309 r"""310 Args:311 sample (`torch.FloatTensor`): Input sample.312 sample_posterior (`bool`, *optional*, defaults to `False`):313 Whether to sample from the posterior.314 return_dict (`bool`, *optional*, defaults to `True`):315 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.316 """317 x = sample318 posterior = self.encode(x).latent_dist319 if sample_posterior:320 z = posterior.sample(generator=generator)321 else:322 z = posterior.mode()323 dec = self.decode(z).sample324 325 if not return_dict:326 return (dec,)327 328 return DecoderOutput(sample=dec)329