Hugging Face's logo

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

O modelo Representation Autoencoder (RAE) apresentado em Diffusion Transformers with Representation Autoencoders, de Boyang Zheng, Nanye Ma, Shengbang Tong e Saining Xie, do NYU VISIONx.

O RAE combina um codificador de visão pré-treinado congelado (DINOv2, SigLIP2 ou MAE) com um decodificador treinável no estilo ViT-MAE. Na receita de treinamento de duas etapas do RAE, o autoencoder é treinado na etapa 1 (reconstrução) e, em seguida, um modelo de difusão é treinado no espaço latente resultante na etapa 2 (geração).

Os seguintes modelos RAE foram lançados e são suportados no Diffusers:

ModelCodificadorFormato latente (entrada de 224px)
nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08DINOv2-base768 x 16 x 16
nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512DINOv2-base (512px)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

Carregando um modelo pré-treinado

from diffusers import AutoencoderRAE

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

Codificação e decodificação de uma imagem real

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")

Normalização latente

Alguns checkpoints pré-treinados incluem estatísticas por canal latents_mean e latents_std para normalizar o espaço latente. Quando presentes, encode e decode aplicam automaticamente a normalização e a denormalização, respectivamente.

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 )

Parâmetros

  • encoder_type (str, opcional, padrão "dinov2") — Tipo de codificador congelado a ser usado. Um de "dinov2", "siglip2" ou "mae".
  • encoder_hidden_size (int, opcional, padrão 768) — Tamanho oculto do modelo codificador.
  • encoder_patch_size (int, opcional, padrão 14) — Tamanho do patch do modelo codificador.
  • encoder_num_hidden_layers (int, opcional, padrão 12) — Número de camadas ocultas no modelo codificador.
  • patch_size (int, opcional, padrão 16) — Tamanho do patch do decodificador (usado para unpatchify e para a cabeça do decodificador).
  • encoder_input_size (int, opcional, padrão 224) — Tamanho de entrada esperado pelo codificador.
  • image_size (int, opcional) — Tamanho da imagem de saída do decodificador. Se for None, é derivado da contagem de tokens do codificador e de patch_size como no RAE-main: image_size = patch_size * sqrt(num_patches), onde num_patches = (encoder_input_size // encoder_patch_size) ** 2.
  • num_channels (int, opcional, padrão 3) — Número de canais de entrada/saída.
  • encoder_norm_mean (list, opcional, padrão [0.485, 0.456, 0.406]) — Média por canal para a normalização da entrada do codificador (padrões do ImageNet).
  • encoder_norm_std (list, opcional, padrão [0.229, 0.224, 0.225]) — Desvio padrão por canal para a normalização da entrada do codificador (padrões do ImageNet).
  • latents_mean (list ou tuple, opcional) — Média opcional para normalização de latentes. Entradas de Tensor são aceitas e convertidas em listas serializáveis de configuração.
  • latents_std (list ou tuple, opcional) — Desvio padrão opcional para normalização de latentes. Entradas de Tensor são aceitas e convertidas em listas serializáveis de configuração.
  • noise_tau (float, opcional, padrão 0.0) — Nível de ruído para o treinamento (adiciona ruído aos latentes durante o treinamento).
  • reshape_to_2d (bool, opcional, padrão True) — Se deve redimensionar os latentes para o formato 2D (B, C, H, W).
  • use_encoder_loss (bool, opcional, padrão False) — Se deve usar os estados ocultos do codificador na perda (para treinamento avançado).

Modelo Representation Autoencoder (RAE) para codificar imagens em latentes e decodificar latentes em imagens.

Este modelo utiliza um codificador pré-treinado congelado (DINOv2, SigLIP2 ou MAE) com um decodificador ViT treinável para reconstruir imagens a partir de representações aprendidas.

Este modelo herda de ModelMixin. Consulte a documentação da superclasse para seus métodos genéricos implementados para todos os modelos (como download ou salvamento).

encode

< >

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

decode

< >

( z: Tensorreturn_dict: bool = True )

forward

< >

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

Parâmetros

  • sample (torch.Tensor) — Amostra de entrada.
  • return_dict (bool, opcional, padrão True) — Se deve ou não retornar um DecoderOutput em vez de uma tupla simples.
  • generator (torch.Generator, opcional) — Um torch.Generator para tornar a amostragem determinística.

Retorna

DecoderOutput ou tuple

Se return_dict for True, um DecoderOutput é retornado; caso contrário, uma tuple simples é retornada.

DecoderOutput

classe diffusers.models.autoencoders.vae.DecoderOutput

< >

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

Parâmetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width)) — A amostra de saída decodificada da última camada do modelo.

Saída do método de decodificação.

Update on GitHub

←AutoencoderKLWan AutoencoderSAME→