一致性模型
一致性模型由杨松、普拉弗拉·达里瓦尔、马克·陈和伊利亚·苏茨克弗在 Consistency Models 中提出。
论文摘要如下：
扩散模型极大地推动了图像、音频和视频生成领域的发展，但它们依赖于迭代采样过程，导致生成速度缓慢。为了克服这一限制，我们提出了**一致性模型**，这是一种新的模型家族，通过将噪声直接映射到数据来生成高质量样本。它们通过设计支持快速一步生成，同时仍然允许多步采样以权衡计算量和样本质量。它们还支持零样本数据编辑，例如图像修复、着色和超分辨率，而无需在这些任务上进行显式训练。一致性模型可以通过蒸馏预训练的扩散模型或作为独立的生成模型进行训练。通过大量实验，我们证明它们在一步和几步采样方面优于现有的扩散模型蒸馏技术，在 CIFAR-10 上实现了 3.55 的新最先进 FID，在 ImageNet 64x64 上实现了 6.20 的一步生成 FID。当单独训练时，一致性模型成为一个新的生成模型家族，在 CIFAR-10、ImageNet 64x64 和 LSUN 256x256 等标准基准上，其性能优于现有的单步非对抗生成模型。
原始代码库可在 openai/consistency_models 中找到，其他检查点可在 openai 中获取。
提示
为了进一步加速，可以使用 torch.compile 在小于 1 秒的时间内生成多个图像。
import torch
from diffusers import ConsistencyModelPipeline
device = "cuda"
# Load the cd_bedroom256_lpips checkpoint.
model_id_or_path = "openai/diffusers-cd_bedroom256_lpips"
pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_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
< 源代码 >( unet: UNet2DModel scheduler: CMStochasticIterativeScheduler )
参数
- unet (UNet2DModel) — 用于对编码图像潜在变量进行去噪的
UNet2DModel。 - scheduler (SchedulerMixin) — 与
unet结合使用以对编码图像潜在变量进行去噪的调度器。目前仅兼容 CMStochasticIterativeScheduler。
用于无条件或类别条件图像生成的管道。
此模型继承自 DiffusionPipeline。有关为所有管道实现的通用方法（下载、保存、在特定设备上运行等），请查看超类文档。
__call__
< 源代码 >( batch_size: int = 1 class_labels: Union = None num_inference_steps: int = 1 timesteps: List = None generator: Union = None latents: Optional = None output_type: Optional = 'pil' return_dict: bool = True callback: Optional = None callback_steps: int = 1 ) → ImagePipelineOutput 或 元组
参数
- 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")