Diffusers documentation
AutoencoderRAE
AutoencoderRAE
El modelo Representation Autoencoder (RAE) introducido en Diffusion Transformers with Representation Autoencoders por Boyang Zheng, Nanye Ma, Shengbang Tong y Saining Xie de NYU VISIONx.
RAE combina un codificador de visión preentrenado congelado (DINOv2, SigLIP2 o MAE) con un decodificador entrenable de estilo ViT-MAE. En la receta de entrenamiento de RAE en dos etapas, el autoencoder se entrena en la etapa 1 (reconstrucción) y, a continuación, un modelo de difusión se entrena en el espacio latente resultante en la etapa 2 (generación).
Los siguientes modelos de RAE se han publicado y están admitidos en Diffusers:
| Modelo | Codificador | Forma latente (entrada 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 (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 |
Carga de un modelo preentrenado
from diffusers import AutoencoderRAE
model = AutoencoderRAE.from_pretrained(
"nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
).to("cuda").eval() # or "mps", "xpu", "cpu"Codificación y decodificación de una imagen 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")Normalización latente
Algunos puntos de control preentrenados incluyen estadísticas por canal latents_mean y latents_std para normalizar el espacio latente. Cuando están presentes, encode y decode aplican automáticamente la normalización y la denormalización, 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).sampleAutoencoderRAE
clase 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, valor predeterminado"dinov2") — Tipo de codificador congelado a utilizar. Uno de"dinov2","siglip2"o"mae". - encoder_hidden_size (
int, opcional, valor predeterminado768) — Tamaño oculto del modelo codificador. - encoder_patch_size (
int, opcional, valor predeterminado14) — Tamaño de parche del modelo codificador. - encoder_num_hidden_layers (
int, opcional, valor predeterminado12) — Número de capas ocultas en el modelo codificador. - patch_size (
int, opcional, valor predeterminado16) — Tamaño del parche del decodificador (se usa para unpatchify y la cabeza del decodificador). - encoder_input_size (
int, opcional, valor predeterminado224) — Tamaño de entrada esperado por el codificador. - image_size (
int, opcional) — Tamaño de la imagen de salida del decodificador. Si esNone, se deriva del número de tokens del codificador y depatch_sizecomo en RAE-main:image_size = patch_size * sqrt(num_patches), dondenum_patches = (encoder_input_size // encoder_patch_size) ** 2. - num_channels (
int, opcional, valor predeterminado3) — Número de canales de entrada/salida. - encoder_norm_mean (
list, opcional, valor predeterminado[0.485, 0.456, 0.406]) — Media por canal para la normalización de la entrada del codificador (valores predeterminados de ImageNet). - encoder_norm_std (
list, opcional, valor predeterminado[0.229, 0.224, 0.225]) — Desviación estándar por canal para la normalización de la entrada del codificador (valores predeterminados de ImageNet). - latents_mean (
listotuple, opcional) — Media opcional para la normalización de latentes. Se aceptan entradas de tensor y se convierten en listas serializables en la configuración. - latents_std (
listotuple, opcional) — Desviación estándar opcional para la normalización de latentes. Se aceptan entradas de Tensor y se convierten en listas serializables en la configuración. - noise_tau (
float, opcional, valor predeterminado0.0) — Nivel de ruido para el entrenamiento (añade ruido a los latentes durante el entrenamiento). - reshape_to_2d (
bool, opcional, valor predeterminadoTrue) — Si se deben redimensionar los latentes al formato 2D (B, C, H, W). - use_encoder_loss (
bool, opcional, valor predeterminadoFalse) — Si se deben usar los estados ocultos del codificador en la pérdida (para entrenamiento avanzado).
Modelo Representation Autoencoder (RAE) para codificar imágenes en latentes y decodificar latentes en imágenes.
Este modelo utiliza un codificador preentrenado congelado (DINOv2, SigLIP2 o MAE) junto con un decodificador ViT entrenable para reconstruir imágenes a partir de representaciones aprendidas.
Este modelo hereda de ModelMixin. Consulta la documentación de la superclase para sus métodos genéricos implementados para todos los modelos (como la descarga o el guardado).
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
Parámetros
- sample (
torch.Tensor) — Muestra de entrada. - return_dict (
bool, opcional, valor predeterminadoTrue) — Indica si se debe devolver unDecoderOutputen lugar de una tupla simple. - generator (
torch.Generator, opcional) — Untorch.Generatorpara hacer que la muestreo sea determinista.
Devuelve
DecoderOutput o tuple
Si return_dict es True, se devuelve un DecoderOutput; en caso contrario, se devuelve una tuple simple.
DecoderOutput
clase diffusers.models.autoencoders.vae.DecoderOutput
< source >( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )
Salida del método de decodificación.