Diffusers documentation

微型自编码器

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

微型自编码器

Ollin Boer Bohan 在 madebyollin/taesd 中提出了 Stable Diffusion 的微型自编码器(TAESD)。它是 Stable Diffusion VAE 的微型蒸馏版本,能够几乎即时地快速解码 StableDiffusionPipeline 或 StableDiffusionXLPipeline 中的潜在表示。

要在 Stable Diffusion v-2.1 中使用:

import torch
from diffusers import DiffusionPipeline, AutoencoderTiny

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1-base", dtype=torch.float16
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd", dtype=torch.float16)
pipe = pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "slice of delicious New York-style berry cheesecake"
image = pipe(prompt, num_inference_steps=25).images[0]
image

与 Stable Diffusion XL 1.0 配合使用

import torch
from diffusers import DiffusionPipeline, AutoencoderTiny

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", dtype=torch.float16)
pipe = pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "slice of delicious New York-style berry cheesecake"
image = pipe(prompt, num_inference_steps=25).images[0]
image

AutoencoderTiny

class diffusers.AutoencoderTiny

< >

( in_channels: int = 3out_channels: int = 3encoder_block_out_channels: tuple[int, ...] = (64, 64, 64, 64)decoder_block_out_channels: tuple[int, ...] = (64, 64, 64, 64)act_fn: str = 'relu'upsample_fn: str = 'nearest'latent_channels: int = 4upsampling_scaling_factor: int = 2num_encoder_blocks: tuple[int, ...] = (1, 3, 3, 3)num_decoder_blocks: tuple[int, ...] = (3, 3, 3, 1)latent_magnitude: int = 3latent_shift: float = 0.5force_upcast: bool = Falsescaling_factor: float = 1.0shift_factor: float = 0.0 )

参数

  • in_channels (int,可选,默认为 3) — 输入图像的通道数。
  • out_channels (int,可选,默认为 3) — 输出中的通道数。
  • encoder_block_out_channels (tuple[int],可选,默认为 (64, 64, 64, 64)) — 一个整数元组,表示每个编码器块的输出通道数。该元组的长度应等于编码器块的数量。
  • decoder_block_out_channels (tuple[int],可选,默认为 (64, 64, 64, 64)) — 一个整数元组,表示每个解码器块的输出通道数。该元组的长度应等于解码器块的数量。
  • act_fn (str,可选,默认为 "relu") — 在整个模型中使用的激活函数。
  • latent_channels (int, 可选, 默认为 4) — 潜在表示中的通道数。潜在空间充当输入图像的压缩表示。
  • upsampling_scaling_factor (int,可选,默认为 2) — 解码器中上采样的缩放因子。它决定了上采样过程中输出图像的尺寸。
  • num_encoder_blocks (tuple[int],可选,默认为 (1, 3, 3, 3)) — 一个整数元组,表示编码过程中每个阶段的编码器块数量。该元组的长度应等于编码器中的阶段数。每个阶段具有不同数量的编码器块。
  • num_decoder_blocks (tuple[int],可选,默认为 (3, 3, 3, 1)) — 一个整数元组,表示解码过程中每个阶段的解码器块数量。该元组的长度应等于解码器中的阶段数。每个阶段具有不同数量的解码器块。
  • latent_magnitude (float, 可选,默认为 3.0) — 潜在表示的幅度。该参数对潜在表示的值进行缩放,以控制信息保留的程度。
  • latent_shift(float,可选,默认为 0.5)— 应用于潜在表示的偏移量。此参数控制潜在空间的中心。
  • scaling_factor (float, 可选,默认为 1.0) — 使用训练集的第一个批次计算得出的训练后潜在空间的逐分量标准差。在训练扩散模型时,用于将潜在空间缩放至单位方差。在将潜在变量传递给扩散模型之前,会按公式 z = z * scaling_factor 进行缩放。在解码时,会按公式 z = 1 / scaling_factor * z 将潜在变量缩放回原始尺度。更多细节请参阅 High-Resolution Image Synthesis with Latent Diffusion Models 论文的 4.3.2 节和 D.1 节。不过,对于此 Autoencoder,并未使用此类缩放因子,因此默认值为 1.0。
  • force_upcast (bool,可选,默认为 False) — 如果启用,将强制 VAE 以 float32 运行,适用于高分辨率图像流水线,例如 SD-XL。VAE 可以微调/训练到更低的数值范围而不会损失太多精度,在这种情况下可以将 force_upcast 设置为 False(参见此 fp16 友好的 AutoEncoder)。

一个小型蒸馏 VAE 模型,用于将图像编码为潜在表示,并将潜在表示解码为图像。

AutoencoderTiny 是 TAESD 原始实现的封装。

该模型继承自 ModelMixin。有关为所有模型实现的通用方法(例如下载或保存),请参阅其父类文档。

forward

< >

( sample: Tensorreturn_dict: bool = True ) → DecoderOutput or tuple

参数

  • sample (torch.Tensor) — 输入样本。
  • return_dict (bool,可选,默认为 True) — 是否返回 DecoderOutput 而不是普通元组。

返回值

DecoderOutput 或 tuple

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

scale_latents

< >

( x: Tensor )

原始潜在变量 -> [0, 1]

unscale_latents

< >

( x: Tensor )

[0, 1] -> 原始潜在变量

AutoencoderTinyOutput

class diffusers.models.autoencoders.autoencoder_tiny.AutoencoderTinyOutput

< >

( latents: Tensor )

参数

  • latents (torch.Tensor) — Encoder 的编码输出。

AutoencoderTiny 编码方法的输出。

Update on GitHub