Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /zh /api /models /autoencoder_rae.md

mishig

2 days ago

preview code

download

raw

8.71 kB

# AutoencoderRAE

由纽约大学 VISIONx 的 Boyang Zheng、Nanye Ma、Shengbang Tong 和 Saining Xie 在 [Diffusion Transformers with Representation Autoencoders](https://huggingface.co/papers/2510.11690) 中提出的表示自编码器（RAE）模型。

RAE 将冻结的预训练视觉编码器（DINOv2、SigLIP2 或 MAE）与可训练的 ViT-MAE 风格解码器相结合。在 RAE 的两阶段训练流程中，第一阶段（重建）训练自编码器，第二阶段（生成）则在所得的潜在空间上训练扩散模型。

以下 RAE 模型已在 Diffusers 中发布并得到支持：

| 模型 | 编码器 | 潜在形状（224px 输入） |
|:------|:--------|:---------------------------|
| [`nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08) | DINOv2-base | 768 x 16 x 16 |
| [`nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512`](https://huggingface.co/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`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-small-ViTXL-n08) | DINOv2-small | 384 x 16 x 16 |
| [`nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08) | DINOv2-large | 1024 x 16 x 16 |
| [`nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08) | SigLIP2-base | 768 x 16 x 16 |
| [`nyu-visionx/RAE-mae-base-p16-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-mae-base-p16-ViTXL-n08) | MAE-base | 768 x 16 x 16 |

## 加载预训练模型

```python
from diffusers import AutoencoderRAE
model = AutoencoderRAE.from_pretrained(
    "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
).to("cuda").eval()  # or "mps", "xpu", "cpu"
```

## 对真实图像进行编码和解码

```python
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` 将分别自动应用归一化和反归一化。

```python
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]]

```python
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)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L393)

**参数：**

encoder_type (`str`，*可选*，默认为 `"dinov2"`)：要使用的冻结编码器类型。可选值为 `"dinov2"`、`"siglip2"` 或 `"mae"`。
encoder_hidden_size (`int`，*可选*，默认为 `768`)：编码器模型的隐藏层大小。
encoder_patch_size (`int`，*可选*，默认值为 `14`)：编码器模型的 patch size。
encoder_num_hidden_layers (`int`，*可选*，默认值为 `12`)：编码器模型中的隐藏层数。

patch_size (`int`，*可选*，默认为 `16`)：解码器 patch 大小（用于 unpatchify 和解码器头）。
encoder_input_size (`int`，*可选*，默认为 `224`)：编码器所需的输入尺寸。
image_size (`int`, *optional*) : 解码器输出图像尺寸。如果为 `None`，则根据编码器 token 数量和 `patch_size` 推导，方式与 RAE-main 相同：`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`，*可选*)：用于潜在变量归一化的可选均值。接受张量输入，并将其转换为可序列化为配置的列表。

latents_std (`list` 或 `tuple`，*可选*)：用于潜在变量归一化的可选标准差。接受张量输入，并将其转换为可序列化为配置的列表。
noise_tau (`float`, *可选*，默认为 `0.0`)：训练时的噪声水平（在训练过程中向潜变量添加噪声）。

reshape_to_2d (`bool`，*可选*，默认为 `True`)：是否将潜变量重塑为 2D (B, C, H, W) 格式。

use_encoder_loss (`bool`，*可选*，默认为 `False`)：是否在损失中使用编码器隐藏状态（用于高级训练）。

用于将图像编码为潜在表示（latents）并将潜在表示解码为图像的表示自编码器（RAE）模型。

该模型使用冻结的预训练编码器（DINOv2、SigLIP2 或 MAE）配合可训练的 ViT 解码器，从学习到的表示中重建图像。

该模型继承自 [ModelMixin](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin)。有关为所有模型实现的通用方法（例如下载或保存），请查阅其父类文档。

#### encode[[diffusers.AutoencoderRAE.encode]]

```python
encode(x: Tensor, return_dict: bool = True, generator: typing.Optional[torch.Generator] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L640)

#### decode[[diffusers.AutoencoderRAE.decode]]

```python
decode(z: Tensor, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L671)

#### forward[[diffusers.AutoencoderRAE.forward]]

```python
forward(sample: Tensor, return_dict: bool = True, generator: typing.Optional[torch.Generator] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L682)

**参数：**

sample (`torch.Tensor`) : 输入样本。

return_dict (`bool`，*可选*，默认为 `True`)：是否返回 `DecoderOutput` 而不是普通元组。
generator (`torch.Generator`，*可选*)：一个 [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html)，用于使采样过程具有确定性。
**返回：** `DecoderOutput` 或 `tuple`
如果 `return_dict` 为 True，则返回 `DecoderOutput`，否则返回普通的 `tuple`。

## DecoderOutput[[diffusers.models.autoencoders.vae.DecoderOutput]]

#### diffusers.models.autoencoders.vae.DecoderOutput[[diffusers.models.autoencoders.vae.DecoderOutput]]

```python
diffusers.models.autoencoders.vae.DecoderOutput(sample: Tensor, commit_loss: typing.Optional[torch.FloatTensor] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/vae.py#L49)

**参数：**

sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`)：模型最后一层解码后的输出样本。

解码方法的输出。

Xet Storage Details

Size:: 8.71 kB
Xet hash:: ec71a181fb17c0607100e7f4012341b0b30d26e934fdec12c07f96047dca9ecd

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# AutoencoderRAE

	由纽约大学 VISIONx 的 Boyang Zheng、Nanye Ma、Shengbang Tong 和 Saining Xie 在 [Diffusion Transformers with Representation Autoencoders](https://huggingface.co/papers/2510.11690) 中提出的表示自编码器（RAE）模型。

	RAE 将冻结的预训练视觉编码器（DINOv2、SigLIP2 或 MAE）与可训练的 ViT-MAE 风格解码器相结合。在 RAE 的两阶段训练流程中，第一阶段（重建）训练自编码器，第二阶段（生成）则在所得的潜在空间上训练扩散模型。

	以下 RAE 模型已在 Diffusers 中发布并得到支持：

	\| 模型 \| 编码器 \| 潜在形状（224px 输入） \|
	\|:------\|:--------\|:---------------------------\|
	\| [`nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08) \| DINOv2-base \| 768 x 16 x 16 \|
	\| [`nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512`](https://huggingface.co/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`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-small-ViTXL-n08) \| DINOv2-small \| 384 x 16 x 16 \|
	\| [`nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08) \| DINOv2-large \| 1024 x 16 x 16 \|
	\| [`nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08) \| SigLIP2-base \| 768 x 16 x 16 \|
	\| [`nyu-visionx/RAE-mae-base-p16-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-mae-base-p16-ViTXL-n08) \| MAE-base \| 768 x 16 x 16 \|

	## 加载预训练模型

	```python
	from diffusers import AutoencoderRAE

	model = AutoencoderRAE.from_pretrained(
	"nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
	).to("cuda").eval() # or "mps", "xpu", "cpu"
	```

	## 对真实图像进行编码和解码

	```python
	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` 将分别自动应用归一化和反归一化。

	```python
	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]]

	```python
	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)
	```

	[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L393)

	参数：

	encoder_type (`str`，可选，默认为 `"dinov2"`)：要使用的冻结编码器类型。可选值为 `"dinov2"`、`"siglip2"` 或 `"mae"`。

	encoder_hidden_size (`int`，可选，默认为 `768`)：编码器模型的隐藏层大小。

	encoder_patch_size (`int`，可选，默认值为 `14`)：编码器模型的 patch size。

	encoder_num_hidden_layers (`int`，可选，默认值为 `12`)：编码器模型中的隐藏层数。

	patch_size (`int`，可选，默认为 `16`)：解码器 patch 大小（用于 unpatchify 和解码器头）。

	encoder_input_size (`int`，可选，默认为 `224`)：编码器所需的输入尺寸。

	image_size (`int`, optional) : 解码器输出图像尺寸。如果为 `None`，则根据编码器 token 数量和 `patch_size` 推导，方式与 RAE-main 相同：`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`，可选)：用于潜在变量归一化的可选均值。接受张量输入，并将其转换为可序列化为配置的列表。

	latents_std (`list` 或 `tuple`，可选)：用于潜在变量归一化的可选标准差。接受张量输入，并将其转换为可序列化为配置的列表。

	noise_tau (`float`, 可选，默认为 `0.0`)：训练时的噪声水平（在训练过程中向潜变量添加噪声）。

	reshape_to_2d (`bool`，可选，默认为 `True`)：是否将潜变量重塑为 2D (B, C, H, W) 格式。

	use_encoder_loss (`bool`，可选，默认为 `False`)：是否在损失中使用编码器隐藏状态（用于高级训练）。

	用于将图像编码为潜在表示（latents）并将潜在表示解码为图像的表示自编码器（RAE）模型。

	该模型使用冻结的预训练编码器（DINOv2、SigLIP2 或 MAE）配合可训练的 ViT 解码器，从学习到的表示中重建图像。

	该模型继承自 [ModelMixin](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin)。有关为所有模型实现的通用方法（例如下载或保存），请查阅其父类文档。

	#### encode[[diffusers.AutoencoderRAE.encode]]

	```python
	encode(x: Tensor, return_dict: bool = True, generator: typing.Optional[torch.Generator] = None)
	```

	[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L640)

	#### decode[[diffusers.AutoencoderRAE.decode]]

	```python
	decode(z: Tensor, return_dict: bool = True)
	```

	[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L671)

	#### forward[[diffusers.AutoencoderRAE.forward]]

	```python
	forward(sample: Tensor, return_dict: bool = True, generator: typing.Optional[torch.Generator] = None)
	```

	[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/autoencoder_rae.py#L682)

	参数：

	sample (`torch.Tensor`) : 输入样本。

	return_dict (`bool`，可选，默认为 `True`)：是否返回 `DecoderOutput` 而不是普通元组。

	generator (`torch.Generator`，可选)：一个 [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html)，用于使采样过程具有确定性。

	返回： `DecoderOutput` 或 `tuple`

	如果 `return_dict` 为 True，则返回 `DecoderOutput`，否则返回普通的 `tuple`。

	## DecoderOutput[[diffusers.models.autoencoders.vae.DecoderOutput]]

	#### diffusers.models.autoencoders.vae.DecoderOutput[[diffusers.models.autoencoders.vae.DecoderOutput]]

	```python
	diffusers.models.autoencoders.vae.DecoderOutput(sample: Tensor, commit_loss: typing.Optional[torch.FloatTensor] = None)
	```

	[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/autoencoders/vae.py#L49)

	参数：

	sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`)：模型最后一层解码后的输出样本。

	解码方法的输出。