Diffusers documentation

AutoencoderRAE

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

AutoencoderRAE

Le modèle Representation Autoencoder (RAE) introduit dans Diffusion Transformers with Representation Autoencoders par Boyang Zheng, Nanye Ma, Shengbang Tong et Saining Xie de NYU VISIONx.

RAE combine un encodeur visuel pré-entraîné figé (DINOv2, SigLIP2 ou MAE) avec un décodeur de type ViT-MAE entraînable. Dans la recette d'entraînement en deux étapes de RAE, l'auto-encodeur est entraîné à l'étape 1 (reconstruction), puis un modèle de diffusion est entraîné sur l'espace latent résultant à l'étape 2 (génération).

Les modèles RAE suivants sont publiés et pris en charge dans Diffusers :

ModèleEncodeurForme latente (entrée de 224 px)
nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08DINOv2-base768 x 16 x 16
nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512DINOv2-base (512 px)768 x 32 x 32
nyu-visionx/RAE-dinov2-wReg-small-ViTXL-n08DINOv2-small384 x 16 x 16
nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08DINOv2-large1024 x 16 x 16
nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08SigLIP2-base768 x 16 x 16
nyu-visionx/RAE-mae-base-p16-ViTXL-n08MAE-base768 x 16 x 16

Chargement d'un modèle pré-entraîné

from diffusers import AutoencoderRAE

model = AutoencoderRAE.from_pretrained(
    "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
).to("cuda").eval()  # or "mps", "xpu", "cpu"

Encodage et décodage d'une image réelle

import torch
from diffusers import AutoencoderRAE
from diffusers.utils import load_image
from torchvision.transforms.functional import to_tensor, to_pil_image

model = AutoencoderRAE.from_pretrained(
    "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
).to("cuda").eval()  # or "mps", "xpu", "cpu"

image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")
image = image.convert("RGB").resize((224, 224))
x = to_tensor(image).unsqueeze(0).to("cuda")  # (1, 3, 224, 224), values in [0, 1]

with torch.no_grad():
    latents = model.encode(x).latent        # (1, 768, 16, 16)
    recon = model.decode(latents).sample     # (1, 3, 256, 256)

recon_image = to_pil_image(recon[0].clamp(0, 1).cpu())
recon_image.save("recon.png")

Normalisation latente

Certains points de contrôle préentraînés incluent des statistiques par canal latents_mean et latents_std pour normaliser l’espace latent. Lorsqu’elles sont présentes, encode et decode appliquent automatiquement la normalisation et la dénormalisation, respectivement.

model = AutoencoderRAE.from_pretrained(
    "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
).to("cuda").eval()  # or "mps", "xpu", "cpu"

# Latent normalization is handled automatically inside encode/decode
# when the checkpoint config includes latents_mean/latents_std.
with torch.no_grad():
    latents = model.encode(x).latent   # normalized latents
    recon = model.decode(latents).sample

AutoencoderRAE

classe diffusers.AutoencoderRAE

< >

( encoder_type: str = 'dinov2'encoder_hidden_size: int = 768encoder_patch_size: int = 14encoder_num_hidden_layers: int = 12decoder_hidden_size: int = 512decoder_num_hidden_layers: int = 8decoder_num_attention_heads: int = 16decoder_intermediate_size: int = 2048patch_size: int = 16encoder_input_size: int = 224image_size: int | None = Nonenum_channels: int = 3encoder_norm_mean: list | None = Noneencoder_norm_std: list | None = Nonelatents_mean: typing.Union[list, tuple, torch.Tensor, NoneType] = Nonelatents_std: typing.Union[list, tuple, torch.Tensor, NoneType] = Nonenoise_tau: float = 0.0reshape_to_2d: bool = Trueuse_encoder_loss: bool = Falsescaling_factor: float = 1.0 )

Paramètres

  • encoder_type (str, facultatif, par défaut "dinov2") — Type d’encodeur figé à utiliser. L’un des "dinov2", "siglip2" ou "mae".
  • encoder_hidden_size (int, facultatif, par défaut 768) — Taille cachée du modèle encodeur.
  • encoder_patch_size (int, facultatif, valeur par défaut 14) — Taille des patches du modèle encodeur.
  • encoder_num_hidden_layers (int, facultatif, valeur par défaut 12) — Nombre de couches cachées dans le modèle encodeur.
  • patch_size (int, facultatif, par défaut 16) — Taille des patches du décodeur (utilisée pour l'unpatchify et la tête du décodeur).
  • encoder_input_size (int, facultatif, valeur par défaut 224) — Taille d'entrée attendue par l'encodeur.
  • image_size (int, facultatif) — Taille de l'image en sortie du décodeur. Si None, elle est dérivée du nombre de jetons du encodeur et de patch_size, comme dans RAE-main : image_size = patch_size * sqrt(num_patches), où num_patches = (encoder_input_size // encoder_patch_size) ** 2.
  • num_channels (int, facultatif, par défaut 3) — Nombre de canaux d'entrée/sortie.
  • encoder_norm_mean (list, facultatif, valeur par défaut [0.485, 0.456, 0.406]) — Moyenne par canal pour la normalisation de l'entrée de l'encodeur (valeurs par défaut d'ImageNet).
  • encoder_norm_std (list, facultatif, par défaut [0.229, 0.224, 0.225]) — Écart-type par canal pour la normalisation de l'entrée de l'encodeur (valeurs par défaut d'ImageNet).
  • latents_mean (list ou tuple, facultatif) — Moyenne facultative pour la normalisation des latents. Les entrées de type Tensor sont acceptées et converties en listes sérialisables dans la configuration.
  • latents_std (list ou tuple, facultatif) — Écart-type facultatif pour la normalisation des latents. Les entrées de type Tensor sont acceptées et converties en listes sérialisables dans la configuration.
  • noise_tau (float, facultatif, par défaut 0.0) — Niveau de bruit pour l'entraînement (ajoute du bruit aux latents pendant l'entraînement).
  • reshape_to_2d (bool, facultatif, par défaut True) — Indique s'il faut redimensionner les latents au format 2D (B, C, H, W).
  • use_encoder_loss (bool, facultatif, par défaut False) — Indique s'il faut utiliser les états cachés de l'encodeur dans la perte (pour l'entraînement avancé).

Modèle Representation Autoencoder (RAE) pour encoder des images en latents et décoder des latents en images.

Ce modèle utilise un encodeur pré-entraîné figé (DINOv2, SigLIP2 ou MAE) avec un décodeur ViT entraînable pour reconstruire des images à partir de représentations apprises.

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

< >

( x: Tensorreturn_dict: bool = Truegenerator: typing.Optional[torch.Generator] = None )

décoder

< >

( z: Tensorreturn_dict: bool = True )

forward

< >

( sample: Tensorreturn_dict: bool = Truegenerator: typing.Optional[torch.Generator] = None ) → DecoderOutput or tuple

Paramètres

  • sample (torch.Tensor) — Échantillon d'entrée.
  • return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un DecoderOutput plutôt qu'un simple tuple.
  • generator (torch.Generator, facultatif) — Un torch.Generator pour rendre l'échantillonnage déterministe.

Retourne

DecoderOutput ou tuple

Si return_dict est True, un DecoderOutput est renvoyé, sinon un simple tuple est renvoyé.

DecoderOutput

classe diffusers.models.autoencoders.vae.DecoderOutput

< >

( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )

Paramètres

  • sample (torch.Tensor de forme (batch_size, num_channels, height, width)) — L'échantillon de sortie décodé de la dernière couche du modèle.

Sortie de la méthode de décodage.

Update on GitHub