Diffusers documentation
AutoencoderDC
AutoencoderDC
Le modèle d'auto-encodeur 2D utilisé dans SANA et introduit dans DCAE par les auteurs Junyu Chen*, Han Cai*, Junsong Chen, Enze Xie, Shang Yang, Haotian Tang, Muyang Li, Yao Lu, Song Han du MIT HAN Lab.
L'abstract de l'article est :
Nous présentons Deep Compression Autoencoder (DC-AE), une nouvelle famille de modèles auto-encodeurs destinée à accélérer les modèles de diffusion haute résolution. Les modèles auto-encodeurs existants ont démontré des résultats impressionnants avec un ratio de compression spatiale modéré (par exemple, 8x), mais échouent à maintenir une précision de reconstruction satisfaisante pour des ratios de compression spatiale élevés (par exemple, 64x). Nous relevons ce défi en introduisant deux techniques clés : (1) l'auto-encodage résiduel, où nous concevons nos modèles pour apprendre des résidus basés sur les caractéristiques transformées espace-vers-canal afin d'atténuer la difficulté d'optimisation des auto-encodeurs à forte compression spatiale ; (2) l'adaptation découplée haute résolution, une stratégie d'entraînement découplée en trois phases efficace pour atténuer la pénalité de généralisation des auto-encodeurs à forte compression spatiale. Grâce à ces conceptions, nous améliorons le ratio de compression spatiale de l'auto-encodeur jusqu'à 128 tout en maintenant la qualité de reconstruction. En appliquant notre DC-AE aux modèles de diffusion latente, nous obtenons une accélération significative sans perte de précision. Par exemple, sur ImageNet 512x512, notre DC-AE offre une accélération d'inférence de 19,1x et une accélération d'entraînement de 17,9x sur GPU H100 pour UViT-H, tout en atteignant un FID meilleur, par rapport à l'auto-encodeur SD-VAE-f8 largement utilisé. Notre code est disponible à cette URL https.
Les modèles DCAE suivants sont publiés et pris en charge dans Diffusers.
Ce modèle a été contribué par lawrence-cj.
Charger un modèle au format Diffusers avec from_pretrained().
from diffusers import AutoencoderDC
ae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers", dtype=torch.float32).to("cuda") # or "mps", "xpu", "cpu"Charger un modèle dans Diffusers via from_single_file
from difusers import AutoencoderDC
ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors"
model = AutoencoderDC.from_single_file(ckpt_path)
Le modèle AutoencoderDC dispose de variantes de points de contrôle en fichier unique in et mix qui possèdent des clés de point de contrôle identiques, mais utilisent des facteurs d’échelle différents. Diffusers ne peut pas déduire automatiquement le fichier de configuration correct à utiliser avec le modèle sur la base du seul point de contrôle et se rabattra sur la configuration de la variante mix. Pour remplacer la configuration déterminée automatiquement, veuillez utiliser l’argument config lors de l’utilisation du chargement en fichier unique avec des points de contrôle de la variante in.
from diffusers import AutoencoderDC
ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f128c512-in-1.0/blob/main/model.safetensors"
model = AutoencoderDC.from_single_file(ckpt_path, config="mit-han-lab/dc-ae-f128c512-in-1.0-diffusers")AutoencoderDC
classe diffusers.AutoencoderDC
< source >( in_channels: int = 3latent_channels: int = 32attention_head_dim: int = 32encoder_block_types: str | tuple[str] = 'ResBlock'decoder_block_types: str | tuple[str] = 'ResBlock'encoder_block_out_channels: tuple = (128, 256, 512, 512, 1024, 1024)decoder_block_out_channels: tuple = (128, 256, 512, 512, 1024, 1024)encoder_layers_per_block: tuple = (2, 2, 2, 3, 3, 3)decoder_layers_per_block: tuple = (3, 3, 3, 3, 3, 3)encoder_qkv_multiscales: tuple = ((), (), (), (5,), (5,), (5,))decoder_qkv_multiscales: tuple = ((), (), (), (5,), (5,), (5,))upsample_block_type: str = 'pixel_shuffle'downsample_block_type: str = 'pixel_unshuffle'decoder_norm_types: str | tuple[str] = 'rms_norm'decoder_act_fns: str | tuple[str] = 'silu'encoder_out_shortcut: bool = Truedecoder_in_shortcut: bool = Truedecoder_conv_act_fn: str = 'relu'scaling_factor: float = 1.0 )
Paramètres
- in_channels (
int, par défaut3) — Le nombre de canaux d'entrée dans les échantillons. - latent_channels (
int, par défaut32) — Le nombre de canaux dans la représentation de l'espace latent. - encoder_block_types (
str | tuple[str], par défaut"ResBlock") — Le ou les types de bloc à utiliser dans l’encodeur. - decoder_block_types (
str | tuple[str], par défaut"ResBlock") — Le ou les types de bloc à utiliser dans le décodeur. - encoder_block_out_channels (
tuple[int, ...], par défaut(128, 256, 512, 512, 1024, 1024)) — Le nombre de canaux de sortie pour chaque bloc de l'encodeur. - decoder_block_out_channels (
tuple[int, ...], par défaut(128, 256, 512, 512, 1024, 1024)) — Le nombre de canaux de sortie pour chaque bloc du décodeur. - encoder_layers_per_block (
tuple[int], par défaut(2, 2, 2, 3, 3, 3)) — Le nombre de couches par bloc dans l'encodeur. - decoder_layers_per_block (
tuple[int], par défaut(3, 3, 3, 3, 3, 3)) — Le nombre de couches par bloc dans le décodeur. - encoder_qkv_multiscales (
tuple[tuple[int, ...], ...], par défaut((), (), (), (5,), (5,), (5,))) — Configurations multi-échelles pour les transformations QKV (query-key-value) de l’encodeur. - decoder_qkv_multiscales (
tuple[tuple[int, ...], ...], par défaut((), (), (), (5,), (5,), (5,))) — Configurations multi-échelles pour les transformations QKV (query-key-value) du décodeur. - upsample_block_type (
str, par défaut"pixel_shuffle") — Le type de bloc à utiliser pour le rééchantillonnage par interpolation (upsampling) dans le décodeur. - downsample_block_type (
str, par défaut"pixel_unshuffle") — Le type de bloc à utiliser pour le rééchantillonnage descendant dans l'encodeur. - decoder_norm_types (
str | tuple[str], par défaut"rms_norm") — Le ou les types de normalisation à utiliser dans le décodeur. - decoder_act_fns (
str | tuple[str], par défaut"silu") — La ou les fonctions d'activation à utiliser dans le décodeur. - encoder_out_shortcut (
bool, par défautTrue) — Indique s'il faut utiliser un raccourci à la fin de l'encodeur. - decoder_in_shortcut (
bool, par défautTrue) — Indique s'il faut utiliser un raccourci au début du décodeur. - decoder_conv_act_fn (
str, par défaut"relu") — La fonction d'activation à utiliser à la fin du décodeur. - scaling_factor (
float, par défaut1.0) — L'inverse multiplicatif de la racine carrée de la moyenne des carrés des caractéristiques latentes. Cela sert à mettre à l'échelle l'espace latent pour qu'il ait une variance unitaire lors de l'entraînement du modèle de diffusion. Les latents sont mis à l'échelle avec la formulez = z * scaling_factoravant d'être transmis au modèle de diffusion. Lors du décodage, les latents sont remis à l'échelle d'origine avec la formule :z = 1 / scaling_factor * z.
Un modèle Autoencoder introduit dans DCAE et utilisé dans SANA.
Ce modèle hérite de ModelMixin. Consultez la documentation de la classe parente pour ses méthodes génériques implémentées pour tous les modèles (téléchargement ou sauvegarde, par exemple).
encode
< source >( x: Tensorreturn_dict: bool = True )
Encode un lot d'images en latents.
décoder
< source >( z: Tensorreturn_dict: bool = True ) → ~models.vae.DecoderOutput or tuple
Paramètres
- z (
torch.Tensor) — Lot d’entrée de vecteurs latents. - return_dict (
bool, par défautTrue) — Indique s'il faut retourner un~models.vae.DecoderOutputplutôt qu'un simple tuple.
Retourne
~models.vae.DecoderOutput ou tuple
If return_dict is True, a ~models.vae.DecoderOutput is returned, otherwise a plain tuple is
returned.
Décoder un lot d'images.
enable_tiling
< source >( tile_sample_min_height: int | None = Nonetile_sample_min_width: int | None = Nonetile_sample_stride_height: float | None = Nonetile_sample_stride_width: float | None = None )
Paramètres
- tile_sample_min_height (
int, facultatif) — Hauteur minimale requise pour qu'un échantillon soit séparé en tuiles le long de la dimension de hauteur. - tile_sample_min_width (
int, facultatif) — Largeur minimale requise pour qu'un échantillon soit séparé en tuiles le long de la dimension de la largeur. - tile_sample_stride_height (
int, facultatif) — La quantité minimale de chevauchement entre deux tuiles verticales consécutives. Cela permet de s'assurer qu'aucun artefact de tuilage n'est produit le long de la dimension de hauteur. - tile_sample_stride_width (
int, facultatif) — La distance entre deux tuiles horizontales consécutives. Cela permet de s'assurer qu'aucun artefact de tuilage n'est produit le long de la dimension de la largeur.
Activer le décodage AE par tuiles. Lorsque cette option est activée, l’AE divise le tenseur d’entrée en tuiles pour effectuer le décodage et l’encodage en plusieurs étapes. Cela est utile pour économiser une grande quantité de mémoire et pour permettre le traitement d’images plus grandes.
forward
< source >( sample: Tensorreturn_dict: bool = True ) → ~models.vae.DecoderOutput or tuple
Paramètres
- sample (
torch.Tensor) — Échantillon d'entrée. - return_dict (
bool, facultatif, par défautTrue) — Indique s'il faut retourner unDecoderOutputplutôt qu'un simple tuple.
Retourne
~models.vae.DecoderOutput ou tuple
Si return_dict est True, un ~models.vae.DecoderOutput est renvoyé, sinon un simple tuple est renvoyé.
DecoderOutput
classe diffusers.models.autoencoders.vae.DecoderOutput
< source >( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )
Sortie de la méthode de décodage.