Diffusers documentation
一致性模型
一致性模型
一致性模型(Consistency Models)由 Yang Song、Prafulla Dhariwal、Mark Chen 和 Ilya Sutskever 在 Consistency Models 中提出。
该论文的摘要如下:
扩散模型在图像、音频和视频生成领域取得了显著进展,但它们依赖于迭代采样过程,导致生成速度缓慢。为了克服这一局限性,我们提出了一致性模型(consistency models),这是一种新型模型,通过直接将噪声映射到数据来生成高质量样本。它们在设计上支持快速的一步生成,同时仍允许多步采样,以在计算量和样本质量之间进行权衡。它们还支持零样本数据编辑,例如图像修复、上色和超分辨率,而无需针对这些任务进行显式训练。一致性模型可以通过蒸馏预训练的扩散模型来训练,也可以作为独立的生成模型进行训练。通过大量实验,我们证明它们在单步和少步采样中优于现有的扩散模型蒸馏技术,在 CIFAR-10 上实现了 3.55 的新 SOTA FID,在 ImageNet 64x64 上实现了 6.20 的新 SOTA FID(均为单步生成)。当独立训练时,一致性模型成为一种新的生成模型家族,在 CIFAR-10、ImageNet 64x64 和 LSUN 256x256 等标准基准上,能够超越现有的单步、非对抗性生成模型。
原始代码库可在 openai/consistency_models 中找到,更多检查点可在 openai 上获取。
提示
如需进一步提速,请使用 torch.compile 在 1 秒内生成多张图片:
import torch
from diffusers import ConsistencyModelPipeline
device = "cuda" # or "mps", "xpu", "cpu"
# Load the cd_bedroom256_lpips checkpoint.
model_id_or_path = "openai/diffusers-cd_bedroom256_lpips"
pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, dtype=torch.float16)
pipe.to(device)
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
# Multistep sampling
# Timesteps can be explicitly specified; the particular timesteps below are from the original GitHub repo:
# https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L83
for _ in range(10):
image = pipe(timesteps=[17, 0]).images[0]
image.show()ConsistencyModelPipeline
类 diffusers.ConsistencyModelPipeline
< source >( unet: UNet2DModelscheduler: CMStochasticIterativeScheduler )
参数
- unet (UNet2DModel) —
用于对编码后的图像潜变量进行去噪的
UNet2DModel。 - scheduler (SchedulerMixin) —
用于与
unet配合,对编码后的图像潜变量进行去噪的调度器。目前仅与 CMStochasticIterativeScheduler 兼容。
用于无条件或类别条件图像生成的 Pipeline。
该模型继承自 DiffusionPipeline。有关为所有 pipeline 实现的通用方法(下载、保存、在特定设备上运行等),请参阅父类文档。
__call__
< source >( batch_size: int = 1class_labels: typing.Union[torch.Tensor, list[int], int, NoneType] = Nonenum_inference_steps: int = 1timesteps: list = Nonegenerator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Truecallback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = Nonecallback_steps: int = 1 ) → ImagePipelineOutput or tuple
参数
- batch_size (
int,可选,默认为 1) — 要生成的图像数量。 - class_labels (
torch.Tensor或list[int]或int,可选) — 用于条件化类条件一致性模型的可选类别标签。如果模型不是类条件的,则不会使用此参数。 - num_inference_steps (
int, 可选,默认为 1) — 去噪步数。更多的去噪步数通常能生成更高质量的图像,但推理速度会变慢。 - timesteps (
list[int], 可选) — 用于去噪过程的自定义时间步。如果未定义,则使用等间距的num_inference_steps个时间步。必须按降序排列。 - generator (
torch.Generator,可选) — 用于使生成过程确定性的torch.Generator。 - latents (
torch.Tensor, 可选) — 预先生成的、从高斯分布中采样的含噪潜变量,用作图像生成的输入。可用于以不同的提示词微调同一生成过程。如果未提供,则使用提供的随机generator进行采样,生成一个潜变量张量。 - output_type (
str,可选,默认为"pil") — 生成图像的输出格式。可选PIL.Image或np.array。 - return_dict (
bool,可选,默认为True) — 是否返回 ImagePipelineOutput 而不是普通元组。 - callback (
Callable, 可选) — 一个在推理过程中每隔callback_steps步调用一次的函数。该函数以以下参数被调用:callback(step: int, timestep: int, latents: torch.Tensor)。 - callback_steps (
int,可选,默认为 1) — 调用callback函数的频率。如果未指定,则会在每一步调用该回调。
返回值
ImagePipelineOutput 或 tuple
如果 return_dict 为 True,则返回 ImagePipelineOutput;否则返回一个 tuple,其中第一个元素是包含生成图像的列表。
示例:
>>> import torch
>>> from diffusers import ConsistencyModelPipeline
>>> device = "cuda"
>>> # Load the cd_imagenet64_l2 checkpoint.
>>> model_id_or_path = "openai/diffusers-cd_imagenet64_l2"
>>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
>>> pipe.to(device)
>>> # Onestep Sampling
>>> image = pipe(num_inference_steps=1).images[0]
>>> image.save("cd_imagenet64_l2_onestep_sample.png")
>>> # Onestep sampling, class-conditional image generation
>>> # ImageNet-64 class label 145 corresponds to king penguins
>>> image = pipe(num_inference_steps=1, class_labels=145).images[0]
>>> image.save("cd_imagenet64_l2_onestep_sample_penguin.png")
>>> # Multistep sampling, class-conditional image generation
>>> # Timesteps can be explicitly specified; the particular timesteps below are from the original GitHub repo:
>>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77
>>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]
>>> image.save("cd_imagenet64_l2_multistep_sample_penguin.png")ImagePipelineOutput
类 diffusers.ImagePipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
图像流水线输出类。