Buckets:
AutoencoderDC
SANA 中使用的 2D 自编码器模型,由 MIT HAN Lab 的 Junyu Chen*、Han Cai*、Junsong Chen、Enze Xie、Shang Yang、Haotian Tang、Muyang Li、Yao Lu 和 Song Han 在 DCAE 中提出。
该论文的摘要如下:
我们提出了 Deep Compression Autoencoder(DC-AE),这是一种用于加速高分辨率扩散模型的新型自编码器模型家族。现有的自编码器模型在中等空间压缩比(例如 8x)下已展现出令人印象深刻的结果,但在高空间压缩比(例如 64x)下无法保持令人满意的重建精度。我们通过引入两项关键技术来解决这一挑战:(1)残差自编码(Residual Autoencoding),我们设计模型基于空间到通道(space-to-channel)变换后的特征来学习残差,以缓解高空间压缩自编码器的优化难度;(2)解耦高分辨率适配(Decoupled High-Resolution Adaptation),一种高效的解耦三阶段训练策略,用于减轻高空间压缩自编码器的泛化惩罚。借助这些设计,我们在保持重建质量的同时,将自编码器的空间压缩比提升至 128。将我们的 DC-AE 应用于潜在扩散模型(latent diffusion models)时,我们在不降低精度的情况下实现了显著的速度提升。例如,在 ImageNet 512x512 上,与广泛使用的 SD-VAE-f8 自编码器相比,我们的 DC-AE 在 H100 GPU 上为 UViT-H 提供了 19.1 倍的推理加速和 17.9 倍的训练加速,同时取得了更好的 FID。我们的代码可在 this https URL 获取。
以下 DCAE 模型已在 Diffusers 中发布并得到支持。
该模型由 lawrence-cj 贡献。
使用 from_pretrained() 以 Diffusers 格式加载模型。
from diffusers import AutoencoderDC
ae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers", dtype=torch.float32).to("cuda") # or "mps", "xpu", "cpu"
通过 from_single_file 在 Diffusers 中加载模型
from difusers import AutoencoderDC
ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors"
model = AutoencoderDC.from_single_file(ckpt_path)
AutoencoderDC 模型具有 in 和 mix 两种单文件检查点变体,它们的检查点键相匹配,但使用不同的缩放因子。Diffusers 无法仅根据检查点自动推断出应使用的正确配置文件,因此默认会使用 mix 变体的配置文件来配置模型。若要覆盖自动确定的配置,请在对 in 变体检查点使用单文件加载时,使用 config 参数。
from diffusers import AutoencoderDC
ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f128c512-in-1.0/blob/main/model.safetensors"
model = AutoencoderDC.from_single_file(ckpt_path, config="mit-han-lab/dc-ae-f128c512-in-1.0-diffusers")
AutoencoderDC[[diffusers.AutoencoderDC]]
diffusers.AutoencoderDC[[diffusers.AutoencoderDC]]
diffusers.AutoencoderDC(in_channels: int = 3, latent_channels: int = 32, attention_head_dim: int = 32, encoder_block_types: str | tuple[str] = 'ResBlock', decoder_block_types: str | tuple[str] = 'ResBlock', encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), decoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), encoder_layers_per_block: tuple[int, ...] = (2, 2, 2, 3, 3, 3), decoder_layers_per_block: tuple[int, ...] = (3, 3, 3, 3, 3, 3), encoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), decoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), upsample_block_type: str = 'pixel_shuffle', downsample_block_type: str = 'pixel_unshuffle', decoder_norm_types: str | tuple[str] = 'rms_norm', decoder_act_fns: str | tuple[str] = 'silu', encoder_out_shortcut: bool = True, decoder_in_shortcut: bool = True, decoder_conv_act_fn: str = 'relu', scaling_factor: float = 1.0)
参数:
in_channels (int,默认为 3):样本中的输入通道数。
latent_channels (int,默认为 32):潜在空间表示中的通道数。
encoder_block_types (str | tuple[str],默认为 "ResBlock"):编码器中使用的块类型。
decoder_block_types (str | tuple[str],默认为 "ResBlock"):解码器中使用的块类型。
encoder_block_out_channels (tuple[int, ...],默认为 (128, 256, 512, 512, 1024, 1024)):编码器中每个块的输出通道数。
decoder_block_out_channels (tuple[int, ...],默认为 (128, 256, 512, 512, 1024, 1024)):解码器中每个块的输出通道数。
encoder_layers_per_block (tuple[int],默认为 (2, 2, 2, 3, 3, 3)):编码器中每个块的层数。
decoder_layers_per_block (tuple[int],默认为 (3, 3, 3, 3, 3, 3)):解码器中每个 block 的层数。
encoder_qkv_multiscales (tuple[tuple[int, ...], ...],默认为 ((), (), (), (5,), (5,), (5,))):编码器 QKV(query-key-value)变换的多尺度配置。
decoder_qkv_multiscales (tuple[tuple[int, ...], ...],默认为 ((), (), (), (5,), (5,), (5,))):解码器 QKV(query-key-value)变换的多尺度配置。
upsample_block_type (str,默认为 "pixel_shuffle"):在解码器中用于上采样的块类型。
downsample_block_type (str,默认为 "pixel_unshuffle"):编码器中用于下采样的块类型。
decoder_norm_types (str | tuple[str],默认为 "rms_norm"):解码器中使用的归一化类型。
decoder_act_fns (str | tuple[str],默认为 "silu"):解码器中使用的激活函数。
encoder_out_shortcut (bool, 默认为 True) : 是否在编码器末尾使用 shortcut。
decoder_in_shortcut (bool,默认为 True):是否在解码器开头使用快捷方式。
decoder_conv_act_fn (str,默认为 "relu"):解码器末尾使用的激活函数。
scaling_factor (float,默认为 1.0):潜在特征均方根值的乘法逆元。该值用于在训练扩散模型时将潜在空间缩放至单位方差。在传入扩散模型之前,潜在变量会按公式 z = z * scaling_factor 进行缩放。解码时,潜在变量会按以下公式缩放回原始尺度:z = 1 / scaling_factor * z。
该模型继承自 ModelMixin。请查阅其父类文档,了解为所有模型实现的通用方法(例如下载或保存)。
encode[[diffusers.AutoencoderDC.encode]]
encode(x: Tensor, return_dict: bool = True)
参数:
x (torch.Tensor) : 输入图像批次。
return_dict (bool,默认为 True):是否返回 ~models.vae.EncoderOutput 而不是普通元组。
返回值:
编码视频的潜在表示。如果 return_dict 为 True,则返回 ~models.vae.EncoderOutput,否则返回普通的 tuple。
将一批图像编码为潜在表示。
decode[[diffusers.AutoencoderDC.decode]]
decode(z: Tensor, return_dict: bool = True)
参数:
z (torch.Tensor) : 输入的潜在向量批次。
return_dict (bool,默认为 True):是否返回 ~models.vae.DecoderOutput 而不是普通元组。
返回: ~models.vae.DecoderOutput 或 tuple
如果 return_dict 为 True,则返回 ~models.vae.DecoderOutput,否则返回普通的 tuple。
解码一批图像。
enable_tiling[[diffusers.AutoencoderDC.enable_tiling]]
enable_tiling(tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None, tile_sample_stride_height: float | None = None, tile_sample_stride_width: float | None = None)
参数:
tile_sample_min_height (int, 可选) : 样本在高度维度上被分割为多个图块所需的最小高度。
tile_sample_min_width (int, 可选) : 将样本沿宽度维度分割为多个图块所需的最小宽度。
tile_sample_stride_height (int, 可选) : 两个连续垂直瓦片之间的最小重叠量。这旨在确保在高度维度上不会产生瓦片拼接伪影。
tile_sample_stride_width (int, 可选) : 两个相邻水平瓦片之间的步长。这用于确保在宽度维度上不会产生瓦片伪影。
启用分块 AE 解码。启用此选项后,AE 会将输入张量拆分为多个块,分步执行解码和编码。这有助于节省大量内存,并支持处理更大的图像。
forward[[diffusers.AutoencoderDC.forward]]
forward(sample: Tensor, return_dict: bool = True)
参数:
sample (torch.Tensor) : 输入样本。
return_dict (bool,可选,默认为 True):是否返回 DecoderOutput 而不是普通元组。
返回: ~models.vae.DecoderOutput 或 tuple
如果 return_dict 为 True,则返回 ~models.vae.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)):模型最后一层解码后的输出样本。
解码方法的输出。
Xet Storage Details
- Size:
- 12 kB
- Xet hash:
- 852f549dc866dc26c08f315241f7dee5aff4c4b1cfe95c0c35bfc905553d657d
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.