Skip to content

一致性模型

一致性模型由Yang Song、Prafulla Dhariwal、Mark Chen和Ilya Sutskever在一致性模型中提出。

论文的摘要如下:

扩散模型显著推动了图像、音频和视频生成领域的发展,但它们依赖于一个迭代采样过程,导致生成速度较慢。为了克服这一限制,我们提出了一致性模型,这是一种新的模型家族,通过直接将噪声映射到数据来生成高质量样本。它们从设计上支持快速的一步生成,同时仍然允许多步采样以换取计算资源来提高样本质量。它们还支持零样本数据编辑,如图像修复、着色和超分辨率,而无需在这些任务上进行显式训练。一致性模型可以通过蒸馏预训练的扩散模型进行训练,或者作为独立的生成模型进行训练。通过广泛的实验,我们证明它们在一步和少步采样中优于现有的扩散模型蒸馏技术,在CIFAR-10上实现了3.55的新FID记录,在ImageNet 64x64上实现了6.20的新FID记录。当独立训练时,一致性模型成为一种新的生成模型家族,可以在CIFAR-10、ImageNet 64x64和LSUN 256x256等标准基准上优于现有的单步非对抗生成模型。

原始代码库可以在openai/consistency_models找到,额外的检查点可以在openai获取。

该流水线由dg845ayushtues贡献。❤️

提示

为了进一步加速,使用torch.compile可以在不到1秒的时间内生成多张图像:

diff
  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

[[autodoc]] ConsistencyModelPipeline - all - call

ImagePipelineOutput

[[autodoc]] pipelines.ImagePipelineOutput