算力平台:
一致性模型
一致性模型由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获取。
提示
为了进一步加速,使用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