一致性模型

本仓库包含 一致性模型 的代码库，使用 PyTorch 实现，用于在 ImageNet-64、LSUN Bedroom-256 和 LSUN Cat-256 数据集上进行大规模实验。我们的仓库基于 openai/guided-diffusion，该仓库最初以 MIT 许可证发布。我们对代码进行了修改，支持一致性蒸馏、一致性训练，以及论文中讨论的多种采样和编辑算法。

CIFAR-10 数据集上的实验代码库采用 JAX 编写，位于 openai/consistency_models_cifar10。

预训练模型

我们发布了论文中主要模型的检查点。在使用这些模型之前，请查阅相应的 模型卡片，以了解这些模型的预期用途及局限性。

以下是各模型检查点的下载链接：

• ImageNet-64 上的 EDM：edm_imagenet64_ema.pt
• ImageNet-64 上使用 l2 度量的一致性蒸馏：cd_imagenet64_l2.pt
• ImageNet-64 上使用 LPIPS 度量的一致性蒸馏：cd_imagenet64_lpips.pt
• ImageNet-64 上的一致性训练：ct_imagenet64.pt
• LSUN Bedroom-256 上的 EDM：edm_bedroom256_ema.pt
• LSUN Bedroom-256 上使用 l2 度量的一致性蒸馏：cd_bedroom256_l2.pt
• LSUN Bedroom-256 上使用 LPIPS 度量的一致性蒸馏：cd_bedroom256_lpips.pt
• LSUN Bedroom-256 上的一致性训练：ct_bedroom256.pt
• LSUN Cat-256 上的 EDM：edm_cat256_ema.pt
• LSUN Cat-256 上使用 l2 度量的一致性蒸馏：cd_cat256_l2.pt
• LSUN Cat-256 上使用 LPIPS 度量的一致性蒸馏：cd_cat256_lpips.pt
• LSUN Cat-256 上的一致性训练：ct_cat256.pt

依赖项

要安装本代码库及其所有依赖项，请运行：

pip install -e .

若使用 Docker 安装，请执行以下命令：

cd docker && make build && make run

模型训练与采样

我们在 scripts/launch.sh 中提供了 EDM 训练、一致性蒸馏、一致性训练、单步生成和多步生成的示例。

评估

为了比较不同的生成模型，我们使用 FID、Precision、Recall 和 Inception Score 等指标。这些指标均可通过存储在 .npz（numpy）文件中的样本批次来计算。可以使用 cm/evaluations/evaluator.py 对样本进行评估，方法与 openai/guided-diffusion 中描述的一致，并参考其中提供的参考数据集批次。

在 🧨 diffusers 中的使用

一致性模型已在 🧨 diffusers 中通过 ConsistencyModelPipeline 类 得到支持。以下是一个示例：

import torch

from diffusers import ConsistencyModelPipeline

device = "cuda"
# 加载 cd_imagenet64_l2 检查点。
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")

# 单步采样，条件式图像生成
# ImageNet-64 类标签 145 对应帝企鹅

class_id = 145
class_id = torch.tensor(class_id, dtype=torch.long)

image = pipe(num_inference_steps=1, class_labels=class_id).images[0]
image.save("consistency_model_onestep_sample_penguin.png")

# 多步采样，条件式图像生成
# 步骤数可以显式指定；以下步骤来自原始 GitHub 仓库。
# https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77
image = pipe(timesteps=[22, 0], class_labels=class_id).images[0]
image.save("consistency_model_multistep_sample_penguin.png")

您还可以通过在 pipe.unet 上使用 torch.compile() 来进一步加速推理过程（仅支持 PyTorch 2.0 及以上版本）。更多详情请参阅 官方文档。此支持由 dg845 和 ayushtues 贡献给 🧨 diffusers。

引用

如果您认为本方法和/或代码有用，请考虑引用以下文献：

@article{song2023consistency,
  title={Consistency Models},
  author={Song, Yang and Dhariwal, Prafulla and Chen, Mark and Sutskever, Ilya},
  journal={arXiv preprint arXiv:2303.01469},
  year={2023},
}

Consistency Models 快速上手指南

Consistency Models 是由 OpenAI 提出的一种新型生成模型，支持单步或多步快速采样。本指南基于官方 PyTorch 实现，帮助开发者快速部署和使用预训练模型。

环境准备

• 系统要求：Linux 或 macOS（推荐 Linux），具备 NVIDIA GPU 以加速推理和训练。
• 前置依赖：
  • Python 3.8+
  • PyTorch 1.10+ (推荐使用最新稳定版)
  • CUDA Toolkit (如需 GPU 加速)
• 国内加速建议：
  • 安装 Python 包时，建议使用清华或阿里镜像源加速下载：

    pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>
    

  • 下载预训练模型权重时，若官方链接访问缓慢，可尝试使用国内云存储中转或代理工具。

安装步骤

方法一：源码安装（推荐）

克隆仓库并安装依赖：

git clone https://github.com/openai/consistency_models.git
cd consistency_models
pip install -e .

注：如遇网络问题，请配合上述国内镜像源使用。

方法二：Docker 安装

如果希望隔离环境，可使用 Docker：

cd docker
make build
make run

基本使用

最便捷的使用方式是通过 Hugging Face 🧨 diffusers 库加载预训练模型，无需手动下载权重文件。

1. 安装 diffusers

pip install diffusers transformers accelerate

2. 单步图像生成示例

以下代码演示了如何加载 ImageNet-64 的预训练模型并进行单步采样（One-step Sampling）：

import torch
from diffusers import ConsistencyModelPipeline

# 设置设备
device = "cuda"

# 加载预训练模型 (自动从 Hugging Face 下载)
# 此处使用 L2 度量的 Consistency Distillation 模型
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")

3. 类别条件生成（可选）

若需生成特定类别的图像（例如 ImageNet 中的“王企鹅”，类别 ID 145）：

# 定义类别标签
class_id = 145
class_id = torch.tensor(class_id, dtype=torch.long).to(device)

# 单步类别条件生成
image = pipe(num_inference_steps=1, class_labels=class_id).images[0]
image.save("consistency_model_onestep_sample_penguin.png")

# 多步类别条件生成 (指定时间步)
image = pipe(timesteps=[22, 0], class_labels=class_id).images[0]
image.save("consistency_model_multistep_sample_penguin.png")

4. 推理加速（可选）

如果您使用的是 PyTorch 2.0 及以上版本，可以使用 torch.compile() 进一步加速推理：

pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

提示：更多预训练模型（如 LSUN Bedroom/Cat 数据集）可在 Hugging Face Model Hub 搜索 openai 组织下的相关模型，或直接参考官方 README 中的权重下载链接手动加载。