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

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

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

lớp 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 )

Tham số

  • encoder_type (str, tùy chọn, mặc định là "dinov2") — Loại encoder đóng băng cần sử dụng. Một trong số "dinov2", "siglip2" hoặc "mae".
  • encoder_hidden_size (int, tùy chọn, mặc định là 768) — Kích thước ẩn của mô hình encoder.
  • encoder_patch_size (int, tùy chọn, mặc định là 14) — Kích thước patch của mô hình encoder.
  • encoder_num_hidden_layers (int, tùy chọn, mặc định là 12) — Số lượng lớp ẩn trong mô hình encoder.
  • patch_size (int, tùy 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, tùy chọn, mặc định là 224) — Kích thước đầu vào mà bộ mã hóa mong đợi.
  • image_size (int, tùy chọn) — 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, tùy 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 encoder (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 encoder (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 dạng Tensor được chấp nhận và chuyển đổi thành các danh sách có thể serialize trong config.
  • latents_std (list hoặc tuple, tùy chọn) — Độ lệch chuẩn 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.
  • noise_tau (float, tùy chọn, mặc định là 0.0) — Mức độ nhiễu cho huấn luyện (thêm nhiễu vào latents trong quá trình 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 encoder trong loss hay không (cho huấn luyện nâng cao).

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

< >

( 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

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

lớp diffusers.models.autoencoders.vae.DecoderOutput

< >

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

Tham số

  • sample (torch.Tensor có hình dạng (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ã.

Update on GitHub