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:

Model	Codificador	Formato latente (entrada de 224px)
`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 (512px)	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

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

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

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

< source >

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

decode

< source >

( z: Tensorreturn_dict: bool = True )

forward

< source >

( 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

< source >

( 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