Buckets:

hf-doc-build/doc / diffusers /main /es /api /models /autoencoder_rae.md
mishig's picture
|
download
raw
9.4 kB

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

AutoencoderRAE[[diffusers.AutoencoderRAE]]

diffusers.AutoencoderRAE[[diffusers.AutoencoderRAE]]

diffusers.AutoencoderRAE(encoder_type: str = 'dinov2', encoder_hidden_size: int = 768, encoder_patch_size: int = 14, encoder_num_hidden_layers: int = 12, decoder_hidden_size: int = 512, decoder_num_hidden_layers: int = 8, decoder_num_attention_heads: int = 16, decoder_intermediate_size: int = 2048, patch_size: int = 16, encoder_input_size: int = 224, image_size: int | None = None, num_channels: int = 3, encoder_norm_mean: list | None = None, encoder_norm_std: list | None = None, latents_mean: typing.Union[list, tuple, torch.Tensor, NoneType] = None, latents_std: typing.Union[list, tuple, torch.Tensor, NoneType] = None, noise_tau: float = 0.0, reshape_to_2d: bool = True, use_encoder_loss: bool = False, scaling_factor: float = 1.0)

Fuente

Parámetros:

encoder_type (str, opcional, por defecto "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 del 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[[diffusers.AutoencoderRAE.encode]]

encode(x: Tensor, return_dict: bool = True, generator: typing.Optional[torch.Generator] = None)

Fuente

decode[[diffusers.AutoencoderRAE.decode]]

decode(z: Tensor, return_dict: bool = True)

Fuente

forward[[diffusers.AutoencoderRAE.forward]]

forward(sample: Tensor, return_dict: bool = True, generator: typing.Optional[torch.Generator] = None)

Fuente

Parámetros:

sample (torch.Tensor) : Muestra de entrada.

return_dict (bool, opcional, por defecto True) : Si se devuelve 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[[diffusers.models.autoencoders.vae.DecoderOutput]]

diffusers.models.autoencoders.vae.DecoderOutput[[diffusers.models.autoencoders.vae.DecoderOutput]]

diffusers.models.autoencoders.vae.DecoderOutput(sample: Tensor, commit_loss: typing.Optional[torch.FloatTensor] = None)

Fuente

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.

Xet Storage Details

Size:
9.4 kB
·
Xet hash:
9a3adb39fb2ac06d0be59106c4e26f36d6d590fdda1b146731e4e4da3be9b4a9

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.