Diffusers documentation
AutoencoderRAE
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).sampleAutoencoderRAE
lớp diffusers.AutoencoderRAE
< source >( 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_sizegiố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 (
listhoặctuple, 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 (
listhoặctuple, 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
< source >( x: Tensorreturn_dict: bool = Truegenerator: typing.Optional[torch.Generator] = None )
forward
< source >( 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ộtDecoderOutputthay vì một tuple thông thường hay không. - generator (
torch.Generator, tùy chọn) — Mộttorch.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
< source >( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )
Đầu ra của phương pháp giải mã.