Consistency Models：单步生成高质量图像的扩散模型

OpenAI 开源的 Consistency Models，在 GitHub 上获得了 6,488 个 Star：

Consistency Models 是一个基于 PyTorch 的代码库，用于在 ImageNet-64、LSUN Bedroom-256 和 LSUN Cat-256 数据集上进行大规模图像生成实验。该项目的核心目标是将传统扩散模型中需要多步迭代的采样过程压缩到单步完成，同时保持生成质量。

该代码库在 openai/guided-diffusion 的基础上进行了扩展，添加了三种训练模式和多种采样策略。Consistency distillation 从预训练的扩散模型中提取知识，Consistency training 直接从头训练一致性模型，这两种方法都能在推理时将采样步骤大幅减少。

项目提供了多组预训练模型，覆盖了三种数据集和不同训练配置。EDM 模型基于扩散模型，CD 模型通过蒸馏获得，CT 模型通过直接训练获得。每种配置都有对应的 checkpoint 可供下载，用户可以根据任务需求选择合适的模型。

安装方式有两种。通过 pip 直接安装：

pip install -e .

或者使用 Docker：

cd docker && make build && make run

项目代码中已经包含了 EDM 训练、一致性蒸馏、一致性训练、单步生成和多步生成的完整示例脚本，位于 scripts/launch.sh 文件中。评估部分支持 FID、Precision、Recall 和 Inception Score 等指标，与 guided-diffusion 的评估流程保持一致。

Consistency Models 也已经集成到了 Hugging Face 的 diffusers 库中，通过 ConsistencyModelPipeline 类即可调用：

import torch
from diffusers import ConsistencyModelPipeline

device = "cuda"
model_id_or_path = "openai/diffusers-cd_imagenet64_l2"
pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.to(device)

# 单步采样
image = pipe(num_inference_steps=1).images[0]
image.save("consistency_model_onestep_sample.png")

# 多步采样
image = pipe(timesteps=[22, 0]).images[0]
image.save("consistency_model_multistep_sample.png")

单步采样即可生成图像，多步采样则能进一步提升图像质量。项目还提供了类别条件生成的示例，可以指定 ImageNet 类别标签来生成特定类型的图像。

对于需要快速图像生成的应用场景，Consistency Models 提供了一个可行的技术路径。相比传统扩散模型几十甚至上百步的采样过程，单步生成在速度上有明显优势，适合对延迟敏感的生产环境。

路径。相比传统扩散模型几十甚至上百步的采样过程，单步生成在速度上有明显优势，适合对延迟敏感的生产环境。