AutoencoderRAE

由纽约大学 VISIONx 的 Boyang Zheng、Nanye Ma、Shengbang Tong 和 Saining Xie 在 Diffusion Transformers with Representation Autoencoders 中提出的表示自编码器（RAE）模型。

RAE 将冻结的预训练视觉编码器（DINOv2、SigLIP2 或 MAE）与可训练的 ViT-MAE 风格解码器相结合。在 RAE 的两阶段训练流程中，第一阶段（重建）训练自编码器，第二阶段（生成）则在所得的潜在空间上训练扩散模型。

以下 RAE 模型已在 Diffusers 中发布并得到支持：

模型	编码器	潜在形状（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

加载预训练模型

from diffusers import AutoencoderRAE

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

对真实图像进行编码和解码

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

潜在归一化

某些预训练检查点包含用于归一化潜在空间的逐通道 latents_mean 和 latents_std 统计量。如果存在这些统计量，encode 和 decode 将分别自动应用归一化和反归一化。

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

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

参数

encoder_type (str，可选，默认为 "dinov2") — 要使用的冻结编码器类型。取值为 "dinov2"、"siglip2" 或 "mae" 之一。
encoder_hidden_size (int，可选，默认为 768) — 编码器模型的隐藏层大小。
encoder_patch_size (int，可选，默认为 14) — 编码器模型的 patch 大小。
encoder_num_hidden_layers (int，可选，默认为 12) — 编码器模型中的隐藏层数量。
patch_size (int，可选，默认为 16) — 解码器 patch 大小（用于 unpatchify 和解码器头）。
encoder_input_size (int，可选，默认为 224) — 编码器期望的输入尺寸。
image_size (int, 可选) — 解码器输出图像的尺寸。如果为 None，则根据编码器 token 数量和 patch_size 推导，方式与 RAE-main 相同：image_size = patch_size * sqrt(num_patches)，其中 num_patches = (encoder_input_size // encoder_patch_size) ** 2。
num_channels (int，可选，默认为 3) — 输入/输出通道数。
encoder_norm_mean (list，可选，默认为 [0.485, 0.456, 0.406]) — 用于编码器输入归一化的逐通道均值（ImageNet 默认值）。
encoder_norm_std (list，可选，默认为 [0.229, 0.224, 0.225]) — 用于编码器输入归一化的逐通道标准差（ImageNet 默认值）。
latents_mean (list 或 tuple，可选) — 用于潜在变量归一化的可选均值。接受张量输入，并将其转换为可序列化为配置的列表。
latents_std (list 或 tuple，可选) — 用于潜在变量归一化的可选标准差。接受张量输入，并将其转换为可序列化为配置的列表。
noise_tau (float，可选，默认为 0.0) — 训练时的噪声水平（在训练过程中向潜变量添加噪声）。
reshape_to_2d (bool，可选，默认为 True) — 是否将潜变量重塑为 2D (B, C, H, W) 格式。
use_encoder_loss (bool，可选，默认为 False) — 是否在损失函数中使用编码器隐藏状态（用于高级训练）。