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

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:

ModeloCodificadorForma latente (entrada de 224 px)
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

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

AutoencoderRAE

clase 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, valor predeterminado "dinov2") — Tipo de codificador congelado a utilizar. Uno de "dinov2", "siglip2" o "mae".
  • encoder_hidden_size (int, opcional, valor predeterminado 768) — Tamaño oculto del modelo codificador.
  • encoder_patch_size (int, opcional, valor predeterminado 14) — Tamaño de parche del modelo codificador.
  • encoder_num_hidden_layers (int, opcional, valor predeterminado 12) — Número de capas ocultas en el modelo codificador.
  • patch_size (int, opcional, valor predeterminado 16) — Tamaño del parche del decodificador (se usa para unpatchify y la cabeza del decodificador).
  • encoder_input_size (int, opcional, valor predeterminado 224) — Tamaño de entrada esperado por el codificador.
  • image_size (int, opcional) — Tamaño de la imagen de salida del decodificador. Si es None, se deriva del número de tokens del codificador y de patch_size como en RAE-main: image_size = patch_size * sqrt(num_patches), donde num_patches = (encoder_input_size // encoder_patch_size) ** 2.
  • num_channels (int, opcional, valor predeterminado 3) — 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 (list o tuple, 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 (list o tuple, 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 predeterminado 0.0) — Nivel de ruido para el entrenamiento (añade ruido a los latentes durante el entrenamiento).
  • reshape_to_2d (bool, opcional, valor predeterminado True) — Si se deben redimensionar los latentes al formato 2D (B, C, H, W).
  • use_encoder_loss (bool, opcional, valor predeterminado False) — 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

< >

( 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) — Muestra de entrada.
  • return_dict (bool, opcional, valor predeterminado True) — Indica si se debe devolver un DecoderOutput en lugar de una tupla simple.
  • generator (torch.Generator, opcional) — Un torch.Generator para 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

< >

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

Parámetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width)) — La muestra de salida decodificada de la última capa del modelo.

Salida del método de decodificación.

Update on GitHub