consistency_models

一致性模型

该存储库包含一致性模型的代码库，使用 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 上的 CD，具有 l2 指标：cd_imagenet64_l2.pt

• ImageNet-64 上的 CD，具有 LPIPS 指标：cd_imagenet64_lpips.pt

• ImageNet-64 上的 CT：ct_imagenet64.pt

• LSUN Bedroom-256 上的 EDM：edm_bedroom256_ema.pt

• LSUN Bedroom-256 上的 CD，具有 l2 指标：cd_bedroom256_l2.pt

• LSUN Bedroom-256 上的 CD，具有 LPIPS 指标：cd_bedroom256_lpips.pt

• LSUN Bedroom-256 上的 CT：ct_bedroom256.pt

• LSUN Cat-256 上的 EDM：edm_cat256_ema.pt

• LSUN Cat-256 上具有 l2 指标的 CD：cd_cat256_l2.pt

• LSUN Cat-256 上的 CD，具有 LPIPS 指标：cd_cat256_lpips.pt

• LSUN Cat-256 上的 CT：ct_cat256.pt

依赖关系

要安装此代码库中的所有包及其依赖项，请运行

pip install -e 。

要使用 Docker 安装，请运行以下命令：

 cd docker && make build && make run

模型训练和采样

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

评价

为了比较不同的生成模型，我们使用 FID、精度、召回率和初始分数。这些指标都可以使用.npz (numpy) 文件中存储的批量样本来计算。人们可以使用 cm/evaluations/evaluator.py 以与 openai/guided-diffusion 中描述的相同方式评估样本，其中提供了参考数据集批次。

用在？扩散器

支持一致性模型？通过ConsistencyModelPipeline类的扩散器。下面我们提供一个例子：

 import torchfrom differs import ConsistencyModelPipelinedevice = "cuda"# 加载 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 Samplingimage = pipeline(num_inference_steps=1).images[0]image.save("consistency_model_onestep_sample.png")#一步采样，类条件图像生成#ImageNet-64类标签145对应王企鹅class_id = 145class_id = torch.tensor(class_id, dtype=torch.long)图像=管道(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#L77image = pipeline(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 的扩散器。

引文

如果您发现此方法和/或代码有用，请考虑引用

@article{song2023consistency，title={一致性模型}，作者={Song，Yang和Dhariwal，Prafulla和Chen，Mark和Sutskever，Ilya}，journal={arXiv预印本arXiv：2303.01469}，year={2023}，
}