Buckets:

hf-doc-build/doc / diffusers /v0.43.0 /ru /api /models /autoencoder_rae.md
mishig's picture
|
download
raw
11.9 kB

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

Источник

Параметры:

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, optional, по умолчанию 16) : Размер патча декодера (используется для unpatchify и головы декодера).

encoder_input_size (int, необязательный, по умолчанию 224) : Ожидаемый кодировщиком размер входных данных.

image_size (int, optional) : Размер выходного изображения декодера. Если 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, optional, по умолчанию [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[[diffusers.AutoencoderRAE.encode]]

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

Источник

decode[[diffusers.AutoencoderRAE.decode]]

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

Источник

forward[[diffusers.AutoencoderRAE.forward]]

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

Источник

Параметры:

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]]

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

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

Источник

Параметры:

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

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

Xet Storage Details

Size:
11.9 kB
·
Xet hash:
f9d6b602129297508442d2354dca7bdc2f2394474f892f32513ebe7d9bdabaf3

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