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

AutoencoderRAE

NYU VISIONxのBoyang Zheng、Nanye Ma、Shengbang Tong、Saining XieによるDiffusion Transformers with Representation Autoencodersで紹介された、Representation Autoencoder (RAE) モデル。

RAEは、凍結された事前学習済みビジョンエンコーダー（DINOv2、SigLIP2、またはMAE）と、学習可能なViT-MAEスタイルのデコーダーを組み合わせます。2段階のRAE学習レシピでは、1段階目でオートエンコーダーが学習（再構成）され、その後、2段階目で得られた潜在空間上で拡散モデルが学習（生成）されます。

Diffusers でリリースおよびサポートされている RAE モデルは以下の通りです:

モデル	エンコーダー	潜在形状（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

事前学習済みモデルの読み込み

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、任意、デフォルトは16) : デコーダーのパッチサイズ（unpatchifyおよびデコーダーヘッドで使用）。

encoder_input_size (int、optional、デフォルトは224) : エンコーダーが想定する入力サイズ。

image_size (int, optional) : デコーダーの出力画像サイズ。None の場合、RAE-main と同様に、エンコーダーのトークン数と patch_size から導出されます: 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、省略可) : 潜在変数の正規化に使用する任意の平均値。Tensor 入力は受け入れられ、設定にシリアライズ可能なリストに変換されます。

latents_std (list または tuple、任意) : 潜在変数の正規化のための任意の標準偏差。Tensor 入力は受け入れられ、設定にシリアライズ可能なリストに変換されます。

noise_tau (float、任意、デフォルトは 0.0) : 学習時のノイズレベル（学習中に潜変数にノイズを追加します）。

reshape_to_2d (bool、任意、デフォルトはTrue) : latents を 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)) : モデルの最終層からのデコード済み出力サンプル。

デコードメソッドの出力。