Hugging Face's logo

Diffusers documentation

AutoencoderRAE

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

AutoencoderRAE

Модель Representation Autoencoder (RAE), представленная в статье Diffusion Transformers with Representation Autoencoders авторами Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie из NYU VISIONx.

RAE объединяет замороженный предобученный визуальный энкодер (DINOv2, SigLIP2 или MAE) с обучаемым декодером в стиле ViT-MAE. В двухэтапном рецепте обучения RAE автоэнкодер обучается на этапе 1 (восстановление), а затем модель диффузии обучается на полученном латентном пространстве на этапе 2 (генерация).

Следующие модели RAE выпущены и поддерживаются в Diffusers:

МодельКодерФорма латентовых представлений (вход 224 пикс.)
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

Загрузка предобученной модели

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

класс 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 )

Параметры

  • encoder_type (str, необязательный, по умолчанию "dinov2") — Тип используемого замороженного энкодера. Одно из значений: "dinov2", "siglip2" или "mae".
  • encoder_hidden_size (int, необязательный, по умолчанию 768) — Размер скрытого слоя модели-энкодера.
  • encoder_patch_size (int, необязательный, по умолчанию 14) — Размер патча модели-энкодера.
  • encoder_num_hidden_layers (int, необязательный, по умолчанию 12) — Количество скрытых слоев в модели-кодере.
  • patch_size (int, необязательный, по умолчанию 16) — Размер патча декодера (используется для unpatchify и головы декодера).
  • encoder_input_size (int, необязательный, по умолчанию 224) — Размер входных данных, ожидаемый кодировщиком.
  • image_size (int, необязательный) — Размер выходного изображения декодера. Если None, он вычисляется на основе количества токенов энкодера и 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) — Использовать ли скрытые состояния энкодера в функции потерь (для расширенного обучения).

Модель Representation Autoencoder (RAE) для кодирования изображений в латенты и декодирования латентов в изображения.

Эта модель использует замороженный предобученный энкодер (DINOv2, SigLIP2 или MAE) и обучаемый декодер ViT для восстановления изображений из полученных представлений.

Эта модель наследуется от ModelMixin. См. документацию суперкласса для общих методов, реализованных для всех моделей (например, загрузка или сохранение).

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

Параметры

  • sample (torch.Tensor) — Входной образец.
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать DecoderOutput вместо обычного кортежа.
  • generator (torch.Generator, необязательный) — torch.Generator для обеспечения детерминированности выборки.

Возвращает

DecoderOutput или tuple

Если return_dict равно True, возвращается DecoderOutput, в противном случае возвращается обычный tuple.

DecoderOutput

класс diffusers.models.autoencoders.vae.DecoderOutput

< >

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

Параметры

  • sample (torch.Tensor формы (batch_size, num_channels, height, width)) — Декодированный вывод из последнего слоя модели.

Результат метода декодирования.

Update on GitHub

←AutoencoderKLWan AutoencoderSAME→