diffusers/main/vi/api/models/autoencoder_rae.md

AutoencoderRAE

Mô hình Representation Autoencoder (RAE) được giới thiệu trong Diffusion Transformers with Representation Autoencoders của Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie từ NYU VISIONx.

RAE kết hợp một bộ mã hóa thị giác được huấn luyện trước và đóng băng (DINOv2, SigLIP2, hoặc MAE) với một bộ giải mã có thể huấn luyện theo phong cách ViT-MAE. Trong quy trình huấn luyện RAE hai giai đoạn, mô hình tự mã hóa được huấn luyện ở giai đoạn 1 (tái tạo), sau đó một mô hình khuếch tán được huấn luyện trên không gian ẩn thu được ở giai đoạn 2 (sinh tạo).

Các mô hình RAE sau đây đã được phát hành và hỗ trợ trong Diffusers:

Mô hình	Bộ mã hóa	Hình dạng ẩn (đầu vào 224px)
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

Tải một mô hình được huấn luyện trước

from diffusers import AutoencoderRAE

model = AutoencoderRAE.from_pretrained(
    "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
).to("cuda").eval()  # or "mps", "xpu", "cpu"

Mã hóa và giải mã một hình ảnh thực

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

Chuẩn hóa ẩn

Một số checkpoint được huấn luyện trước bao gồm các thống kê latents_mean và latents_std theo từng kênh để chuẩn hóa không gian ẩn. Khi có mặt, encode và decode sẽ tự động áp dụng lần lượt phép chuẩn hóa và khử chuẩn hóa.

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)

Nguồn

Tham số:

encoder_type (str, tùy chọn, mặc định là "dinov2") : Loại bộ mã hóa đóng băng cần sử dụng. Một trong số "dinov2", "siglip2" hoặc "mae".

encoder_hidden_size (int, tuỳ chọn, mặc định là 768) : Kích thước ẩn của mô hình encoder.

encoder_patch_size (int, tuỳ chọn, mặc định là 14) : Kích thước patch của mô hình encoder.

encoder_num_hidden_layers (int, tuỳ chọn, mặc định là 12) : Số lượng lớp ẩn trong mô hình encoder.

patch_size (int, tuỳ chọn, mặc định là 16) : Kích thước patch của bộ giải mã (được sử dụng cho unpatchify và đầu ra của bộ giải mã).

encoder_input_size (int, tuỳ chọn, mặc định là 224) : Kích thước đầu vào mà bộ mã hóa mong đợi.

image_size (int, optional) : Kích thước ảnh đầu ra của bộ giải mã. Nếu là None, giá trị này được suy ra từ số lượng token của bộ mã hóa và patch_size giống như RAE-main: image_size = patch_size * sqrt(num_patches), trong đó num_patches = (encoder_input_size // encoder_patch_size) ** 2.

num_channels (int, tuỳ chọn, mặc định là 3) : Số lượng kênh đầu vào/đầu ra.

encoder_norm_mean (list, tùy chọn, mặc định là [0.485, 0.456, 0.406]) : Giá trị trung bình theo kênh để chuẩn hóa đầu vào của bộ mã hóa (giá trị mặc định của ImageNet).

encoder_norm_std (list, tùy chọn, mặc định là [0.229, 0.224, 0.225]) : Độ lệch chuẩn theo kênh cho việc chuẩn hóa đầu vào của bộ mã hóa (giá trị mặc định của ImageNet).

latents_mean (list hoặc tuple, tùy chọn) : Giá trị trung bình tùy chọn để chuẩn hóa latent. Các đầu vào Tensor được chấp nhận và chuyển đổi thành các danh sách có thể chuỗi hóa cấu hình.

latents_std (list hoặc tuple, tùy chọn) : Độ lệch chuẩn tùy chọn cho việc chuẩn hóa latent. Các đầu vào Tensor được chấp nhận và chuyển đổi thành các danh sách có thể tuần tự hóa cấu hình.

noise_tau (float, tuỳ chọn, mặc định là 0.0) : Mức độ nhiễu cho quá trình huấn luyện (thêm nhiễu vào latents trong khi huấn luyện).

reshape_to_2d (bool, tùy chọn, mặc định là True) : Có nên định dạng lại latents thành định dạng 2D (B, C, H, W) hay không.

use_encoder_loss (bool, tùy chọn, mặc định là False) : Có sử dụng trạng thái ẩn của bộ mã hóa trong hàm mất mát (cho huấn luyện nâng cao) hay không.

Mô hình Representation Autoencoder (RAE) để mã hóa ảnh thành các ẩn (latents) và giải mã các ẩn thành ảnh.

Mô hình này sử dụng một encoder được huấn luyện trước và đóng băng (DINOv2, SigLIP2, hoặc MAE) cùng với một decoder ViT có thể huấn luyện để tái tạo hình ảnh từ các biểu diễn đã học.

Mô hình này kế thừa từ ModelMixin. Xem tài liệu của lớp cha để biết các phương thức chung được triển khai cho tất cả các mô hình (chẳng hạn như tải xuống hoặc lưu).

encode[[diffusers.AutoencoderRAE.encode]]

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

Nguồn

decode[[diffusers.AutoencoderRAE.decode]]

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

Nguồn

forward[[diffusers.AutoencoderRAE.forward]]

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

Nguồn

Tham số:

sample (torch.Tensor) : Mẫu đầu vào.

return_dict (bool, tùy chọn, mặc định là True) : Có nên trả về một DecoderOutput thay vì một tuple thông thường hay không.

generator (torch.Generator, tùy chọn) : Một torch.Generator để làm cho việc lấy mẫu trở nên tất định.

Trả về: DecoderOutput hoặc tuple

Nếu return_dict là True, một DecoderOutput sẽ được trả về, ngược lại một tuple thông thường sẽ được trả về.

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)

Nguồn

Tham số:

sample (torch.Tensor có shape (batch_size, num_channels, height, width)) : Mẫu đầu ra đã giải mã từ lớp cuối cùng của mô hình.

Đầu ra của phương pháp giải mã.