Diffusers documentation
AutoencoderRAE
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èle | Encodeur | Forme latente (entrée de 224 px) |
|---|---|---|
nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08 | DINOv2-base | 768 x 16 x 16 |
nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512 | DINOv2-base (512 px) | 768 x 32 x 32 |
nyu-visionx/RAE-dinov2-wReg-small-ViTXL-n08 | DINOv2-small | 384 x 16 x 16 |
nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08 | DINOv2-large | 1024 x 16 x 16 |
nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08 | SigLIP2-base | 768 x 16 x 16 |
nyu-visionx/RAE-mae-base-p16-ViTXL-n08 | MAE-base | 768 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).sampleAutoencoderRAE
classe diffusers.AutoencoderRAE
< source >( 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éfaut768) — Taille cachée du modèle encodeur. - encoder_patch_size (
int, facultatif, valeur par défaut14) — Taille des patches du modèle encodeur. - encoder_num_hidden_layers (
int, facultatif, valeur par défaut12) — Nombre de couches cachées dans le modèle encodeur. - patch_size (
int, facultatif, par défaut16) — 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éfaut224) — Taille d'entrée attendue par l'encodeur. - image_size (
int, facultatif) — Taille de l'image en sortie du décodeur. SiNone, elle est dérivée du nombre de jetons du encodeur et depatch_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éfaut3) — 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 (
listoutuple, 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 (
listoutuple, 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éfaut0.0) — Niveau de bruit pour l'entraînement (ajoute du bruit aux latents pendant l'entraînement). - reshape_to_2d (
bool, facultatif, par défautTrue) — Indique s'il faut redimensionner les latents au format 2D (B, C, H, W). - use_encoder_loss (
bool, facultatif, par défautFalse) — 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
< source >( x: Tensorreturn_dict: bool = Truegenerator: typing.Optional[torch.Generator] = None )
forward
< source >( 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éfautTrue) — Indique s'il faut retourner unDecoderOutputplutôt qu'un simple tuple. - generator (
torch.Generator, facultatif) — Untorch.Generatorpour 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
< source >( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )
Sortie de la méthode de décodage.