Diffusers documentation
AutoencoderRAE
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).sampleAutoencoderRAE
class 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 )
パラメータ
- 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
< 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
パラメータ
- 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
< source >( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )
デコードメソッドの出力。