|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from dataclasses import dataclass |
|
|
from typing import Tuple, Union |
|
|
|
|
|
import torch |
|
|
|
|
|
from ..configuration_utils import ConfigMixin, register_to_config |
|
|
from ..utils import BaseOutput, apply_forward_hook |
|
|
from .modeling_utils import ModelMixin |
|
|
from .vae import DecoderOutput, DecoderTiny, EncoderTiny |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class AutoencoderTinyOutput(BaseOutput): |
|
|
""" |
|
|
Output of AutoencoderTiny encoding method. |
|
|
|
|
|
Args: |
|
|
latents (`torch.Tensor`): Encoded outputs of the `Encoder`. |
|
|
|
|
|
""" |
|
|
|
|
|
latents: torch.Tensor |
|
|
|
|
|
|
|
|
class AutoencoderTiny(ModelMixin, ConfigMixin): |
|
|
r""" |
|
|
A tiny distilled VAE model for encoding images into latents and decoding latent representations into images. |
|
|
|
|
|
[`AutoencoderTiny`] is a wrapper around the original implementation of `TAESD`. |
|
|
|
|
|
This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods implemented for |
|
|
all models (such as downloading or saving). |
|
|
|
|
|
Parameters: |
|
|
in_channels (`int`, *optional*, defaults to 3): Number of channels in the input image. |
|
|
out_channels (`int`, *optional*, defaults to 3): Number of channels in the output. |
|
|
encoder_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`): |
|
|
Tuple of integers representing the number of output channels for each encoder block. The length of the |
|
|
tuple should be equal to the number of encoder blocks. |
|
|
decoder_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`): |
|
|
Tuple of integers representing the number of output channels for each decoder block. The length of the |
|
|
tuple should be equal to the number of decoder blocks. |
|
|
act_fn (`str`, *optional*, defaults to `"relu"`): |
|
|
Activation function to be used throughout the model. |
|
|
latent_channels (`int`, *optional*, defaults to 4): |
|
|
Number of channels in the latent representation. The latent space acts as a compressed representation of |
|
|
the input image. |
|
|
upsampling_scaling_factor (`int`, *optional*, defaults to 2): |
|
|
Scaling factor for upsampling in the decoder. It determines the size of the output image during the |
|
|
upsampling process. |
|
|
num_encoder_blocks (`Tuple[int]`, *optional*, defaults to `(1, 3, 3, 3)`): |
|
|
Tuple of integers representing the number of encoder blocks at each stage of the encoding process. The |
|
|
length of the tuple should be equal to the number of stages in the encoder. Each stage has a different |
|
|
number of encoder blocks. |
|
|
num_decoder_blocks (`Tuple[int]`, *optional*, defaults to `(3, 3, 3, 1)`): |
|
|
Tuple of integers representing the number of decoder blocks at each stage of the decoding process. The |
|
|
length of the tuple should be equal to the number of stages in the decoder. Each stage has a different |
|
|
number of decoder blocks. |
|
|
latent_magnitude (`float`, *optional*, defaults to 3.0): |
|
|
Magnitude of the latent representation. This parameter scales the latent representation values to control |
|
|
the extent of information preservation. |
|
|
latent_shift (float, *optional*, defaults to 0.5): |
|
|
Shift applied to the latent representation. This parameter controls the center of the latent space. |
|
|
scaling_factor (`float`, *optional*, defaults to 1.0): |
|
|
The component-wise standard deviation of the trained latent space computed using the first batch of the |
|
|
training set. This is used to scale the latent space to have unit variance when training the diffusion |
|
|
model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the |
|
|
diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 |
|
|
/ scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image |
|
|
Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper. For this Autoencoder, |
|
|
however, no such scaling factor was used, hence the value of 1.0 as the default. |
|
|
force_upcast (`bool`, *optional*, default to `False`): |
|
|
If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE |
|
|
can be fine-tuned / trained to a lower range without losing too much precision, in which case |
|
|
`force_upcast` can be set to `False` (see this fp16-friendly |
|
|
[AutoEncoder](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix)). |
|
|
""" |
|
|
_supports_gradient_checkpointing = True |
|
|
|
|
|
@register_to_config |
|
|
def __init__( |
|
|
self, |
|
|
in_channels=3, |
|
|
out_channels=3, |
|
|
encoder_block_out_channels: Tuple[int] = (64, 64, 64, 64), |
|
|
decoder_block_out_channels: Tuple[int] = (64, 64, 64, 64), |
|
|
act_fn: str = "relu", |
|
|
latent_channels: int = 4, |
|
|
upsampling_scaling_factor: int = 2, |
|
|
num_encoder_blocks: Tuple[int] = (1, 3, 3, 3), |
|
|
num_decoder_blocks: Tuple[int] = (3, 3, 3, 1), |
|
|
latent_magnitude: int = 3, |
|
|
latent_shift: float = 0.5, |
|
|
force_upcast: float = False, |
|
|
scaling_factor: float = 1.0, |
|
|
): |
|
|
super().__init__() |
|
|
|
|
|
if len(encoder_block_out_channels) != len(num_encoder_blocks): |
|
|
raise ValueError("`encoder_block_out_channels` should have the same length as `num_encoder_blocks`.") |
|
|
if len(decoder_block_out_channels) != len(num_decoder_blocks): |
|
|
raise ValueError("`decoder_block_out_channels` should have the same length as `num_decoder_blocks`.") |
|
|
|
|
|
self.encoder = EncoderTiny( |
|
|
in_channels=in_channels, |
|
|
out_channels=latent_channels, |
|
|
num_blocks=num_encoder_blocks, |
|
|
block_out_channels=encoder_block_out_channels, |
|
|
act_fn=act_fn, |
|
|
) |
|
|
|
|
|
self.decoder = DecoderTiny( |
|
|
in_channels=latent_channels, |
|
|
out_channels=out_channels, |
|
|
num_blocks=num_decoder_blocks, |
|
|
block_out_channels=decoder_block_out_channels, |
|
|
upsampling_scaling_factor=upsampling_scaling_factor, |
|
|
act_fn=act_fn, |
|
|
) |
|
|
|
|
|
self.latent_magnitude = latent_magnitude |
|
|
self.latent_shift = latent_shift |
|
|
self.scaling_factor = scaling_factor |
|
|
|
|
|
def _set_gradient_checkpointing(self, module, value=False): |
|
|
if isinstance(module, (EncoderTiny, DecoderTiny)): |
|
|
module.gradient_checkpointing = value |
|
|
|
|
|
def scale_latents(self, x): |
|
|
"""raw latents -> [0, 1]""" |
|
|
return x.div(2 * self.latent_magnitude).add(self.latent_shift).clamp(0, 1) |
|
|
|
|
|
def unscale_latents(self, x): |
|
|
"""[0, 1] -> raw latents""" |
|
|
return x.sub(self.latent_shift).mul(2 * self.latent_magnitude) |
|
|
|
|
|
@apply_forward_hook |
|
|
def encode( |
|
|
self, x: torch.FloatTensor, return_dict: bool = True |
|
|
) -> Union[AutoencoderTinyOutput, Tuple[torch.FloatTensor]]: |
|
|
output = self.encoder(x) |
|
|
|
|
|
if not return_dict: |
|
|
return (output,) |
|
|
|
|
|
return AutoencoderTinyOutput(latents=output) |
|
|
|
|
|
@apply_forward_hook |
|
|
def decode(self, x: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.FloatTensor]]: |
|
|
output = self.decoder(x) |
|
|
|
|
|
|
|
|
output = output.mul_(2).sub_(1) |
|
|
|
|
|
if not return_dict: |
|
|
return (output,) |
|
|
|
|
|
return DecoderOutput(sample=output) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
sample: torch.FloatTensor, |
|
|
return_dict: bool = True, |
|
|
) -> Union[DecoderOutput, Tuple[torch.FloatTensor]]: |
|
|
r""" |
|
|
Args: |
|
|
sample (`torch.FloatTensor`): Input sample. |
|
|
return_dict (`bool`, *optional*, defaults to `True`): |
|
|
Whether or not to return a [`DecoderOutput`] instead of a plain tuple. |
|
|
""" |
|
|
enc = self.encode(sample).latents |
|
|
scaled_enc = self.scale_latents(enc).mul_(255).round_().byte() |
|
|
unscaled_enc = self.unscale_latents(scaled_enc) |
|
|
dec = self.decode(unscaled_enc) |
|
|
|
|
|
if not return_dict: |
|
|
return (dec,) |
|
|
return DecoderOutput(sample=dec) |
|
|
|
