AsymmetricAutoencoderKL
改进的更大的变分自动编码器 (VAE) 模型,使用 KL 损失进行修复任务:为 StableDiffusion 设计更好的非对称 VQGAN,作者为 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 节中找到。
可用检查点
- https://huggingface.co/cross-attention/asymmetric-autoencoder-kl-x-1-5
- https://huggingface.co/cross-attention/asymmetric-autoencoder-kl-x-2
示例用法
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
类 diffusers.AsymmetricAutoencoderKL
< 源代码 >( in_channels: int = 3 out_channels: int = 3 down_block_types: Tuple = ('DownEncoderBlock2D',) down_block_out_channels: Tuple = (64,) layers_per_down_block: int = 1 up_block_types: Tuple = ('UpDecoderBlock2D',) up_block_out_channels: Tuple = (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",)
) — 上采样块类型的元组。 - 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
缩放回原始尺度。有关更多详细信息,请参阅论文 使用潜在扩散模型的高分辨率图像合成 的第 4.3.2 节和 D.1 节。
为 StableDiffusion 设计更好的非对称 VQGAN https://arxiv.org/abs/2306.04632 。一个使用 KL 损失的 VAE 模型,用于将图像编码为潜变量,并将潜变量表示解码为图像。
此模型继承自 ModelMixin。查看超类文档以了解为所有模型实现的通用方法(例如下载或保存)。
前向传播
< 源代码 > ( sample: 张量 mask: 可选 = None sample_posterior: bool = False return_dict: bool = True generator: 可选 = None )
AutoencoderKLOutput
类 diffusers.models.modeling_outputs.AutoencoderKLOutput
< 源代码 >( latent_dist: DiagonalGaussianDistribution )
AutoencoderKL 编码方法的输出。
DecoderOutput
类 diffusers.models.autoencoders.vae.DecoderOutput
< 源代码 >( sample: Tensor commit_loss: Optional = None )
解码方法的输出。