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

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

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

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

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

パラメータ

  • 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、任意、デフォルトは 224) — エンコーダーが期待する入力サイズです。
  • image_size (int, 任意) — デコーダーの出力画像サイズ。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 入力は受け入れられ、config シリアライズ可能なリストに変換されます。
  • 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

< >

( 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

パラメータ

  • sample (torch.Tensor) — 入力サンプル。
  • return_dict (bool、省略可、デフォルトは True) — DecoderOutput を返すか、プレーンなタプルを返すかを指定します。
  • generator (torch.Generator、省略可) — サンプリングを決定論的に行うためのtorch.Generator。

戻り値

DecoderOutput または tuple

return_dict が True の場合、DecoderOutput が返され、それ以外の場合は通常の tuple が返されます。

DecoderOutput

class diffusers.models.autoencoders.vae.DecoderOutput

< >

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

パラメータ

  • sample (torch.Tensor、形状 (batch_size, num_channels, height, width)) — モデルの最終層からのデコード済み出力サンプルです。

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

Update on GitHub