Diffusers 文档

AsymmetricAutoencoderKL

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

AsymmetricAutoencoderKL

改进的更大的变分自编码器 (VAE) 模型,具有 KL 损失,用于图像修复任务:Designing a Better Asymmetric VQGAN for StableDiffusion,作者:Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang, Yinpeng Chen, Lu Yuan, Gang Hua。

论文摘要如下:

StableDiffusion 是一种革命性的文本到图像生成器,在图像生成和编辑领域引起了轰动。与在像素空间中学习扩散模型的传统方法不同,StableDiffusion 通过 VQGAN 在潜在空间中学习扩散模型,确保了效率和质量。它不仅支持图像生成任务,还支持真实图像的图像编辑,例如图像修复和局部编辑。然而,我们观察到 StableDiffusion 中使用的原始 VQGAN 会导致严重的信息丢失,即使在未编辑的图像区域也会引起失真伪影。为此,我们提出了一种新的非对称 VQGAN,它具有两个简单的设计。首先,除了来自编码器的输入外,解码器还包含一个条件分支,该分支结合了来自特定任务先验的信息,例如图像修复中未被遮罩的图像区域。其次,解码器比编码器重得多,从而可以进行更精细的恢复,而总推理成本仅略有增加。我们的非对称 VQGAN 的训练成本很低,我们只需要重新训练一个新的非对称解码器,同时保持原始 VQGAN 编码器和 StableDiffusion 不变。我们的非对称 VQGAN 可以广泛用于基于 StableDiffusion 的图像修复和局部编辑方法。广泛的实验表明,它可以显着提高图像修复和编辑性能,同时保持原始的文本到图像能力。代码可在 https://github.com/buxiangzhiren/Asymmetric_VQGAN 获取

评估结果可以在原始论文的第 4.1 节中找到。

可用检查点

使用示例

from diffusers import AsymmetricAutoencoderKL, StableDiffusionInpaintPipeline
from diffusers.utils import load_image, make_image_grid


prompt = "a photo of a person with beard"
img_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/celeba_hq_256.png"
mask_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"

original_image = load_image(img_url).resize((512, 512))
mask_image = load_image(mask_url).resize((512, 512))

pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
pipe.vae = AsymmetricAutoencoderKL.from_pretrained("cross-attention/asymmetric-autoencoder-kl-x-1-5")
pipe.to("cuda")

image = pipe(prompt=prompt, image=original_image, mask_image=mask_image).images[0]
make_image_grid([original_image, mask_image, image], rows=1, cols=3)

AsymmetricAutoencoderKL

class diffusers.AsymmetricAutoencoderKL

< >

( in_channels: int = 3 out_channels: int = 3 down_block_types: typing.Tuple[str, ...] = ('DownEncoderBlock2D',) down_block_out_channels: typing.Tuple[int, ...] = (64,) layers_per_down_block: int = 1 up_block_types: typing.Tuple[str, ...] = ('UpDecoderBlock2D',) up_block_out_channels: typing.Tuple[int, ...] = (64,) layers_per_up_block: int = 1 act_fn: str = 'silu' latent_channels: int = 4 norm_num_groups: int = 32 sample_size: int = 32 scaling_factor: float = 0.18215 )

参数

  • in_channels (int, 可选, 默认为 3) — 输入图像中的通道数。
  • out_channels (int, 可选, 默认为 3) — 输出中的通道数。
  • down_block_types (Tuple[str], 可选, 默认为 ("DownEncoderBlock2D",)) — 下采样块类型的元组。
  • down_block_out_channels (Tuple[int], 可选, 默认为 (64,)) — 下块输出通道的元组。
  • layers_per_down_block (int, 可选, 默认为 1) — 每个下块的层数。
  • up_block_types (Tuple[str], 可选, 默认为 ("UpDecoderBlock2D",)) — 上采样模块类型的元组。
  • up_block_out_channels (Tuple[int], 可选, 默认为 (64,)) — 上采样模块输出通道的元组。
  • layers_per_up_block (int, 可选, 默认为 1) — 每个上采样模块的层数。
  • act_fn (str, 可选, 默认为 "silu") — 要使用的激活函数。
  • latent_channels (int, 可选, 默认为 4) — 潜在空间中的通道数。
  • sample_size (int, 可选, 默认为 32) — 采样输入大小。
  • norm_num_groups (int, 可选, 默认为 32) — 在 ResNet 块的第一个归一化层中使用的组数。
  • scaling_factor (float, 可选, 默认为 0.18215) — 使用训练集的第一批数据计算出的训练后潜在空间的组件级标准差。这用于在训练扩散模型时缩放潜在空间以使其具有单位方差。潜在变量在传递给扩散模型之前使用公式 z = z * scaling_factor 进行缩放。解码时,潜在变量使用公式 z = 1 / scaling_factor * z 缩放回原始尺度。有关更多详细信息,请参阅 High-Resolution Image Synthesis with Latent Diffusion Models 论文的 4.3.2 和 D.1 节。

Designing a Better Asymmetric VQGAN for StableDiffusion https://arxiv.org/abs/2306.04632 。具有 KL 损失的 VAE 模型,用于将图像编码为潜在变量并将潜在表示解码为图像。

此模型继承自 ModelMixin。查看超类文档以了解为其所有模型实现的通用方法(例如下载或保存)。

forward

< >

( sample: Tensor mask: typing.Optional[torch.Tensor] = None sample_posterior: bool = False return_dict: bool = True generator: typing.Optional[torch._C.Generator] = None )

参数

  • sample (torch.Tensor) — 输入样本。
  • mask (torch.Tensor, 可选, 默认为 None) — 可选的图像修复掩码。
  • sample_posterior (bool, 可选, 默认为 False) — 是否从后验分布中采样。
  • return_dict (bool, 可选, 默认为 True) — 是否返回 DecoderOutput 而不是普通元组。

AutoencoderKLOutput

class diffusers.models.modeling_outputs.AutoencoderKLOutput

< >

( latent_dist: DiagonalGaussianDistribution )

参数

  • latent_dist (DiagonalGaussianDistribution) — Encoder 的编码输出,表示为 DiagonalGaussianDistribution 的均值和 logvar。 DiagonalGaussianDistribution 允许从分布中采样潜在变量。

AutoencoderKL 编码方法的输出。

DecoderOutput

class diffusers.models.autoencoders.vae.DecoderOutput

< >

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

参数

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

解码方法的输出。

< > 在 GitHub 上更新