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

文章目录Consistency Models单步生成高质量图像的扩散模型Consistency Models单步生成高质量图像的扩散模型OpenAI 开源的 Consistency Models在 GitHub 上获得了 6,488 个 StarConsistency Models 是一个基于 PyTorch 的代码库用于在 ImageNet-64、LSUN Bedroom-256 和 LSUN Cat-256 数据集上进行大规模图像生成实验。该项目的核心目标是将传统扩散模型中需要多步迭代的采样过程压缩到单步完成同时保持生成质量。该代码库在 openai/guided-diffusion 的基础上进行了扩展添加了三种训练模式和多种采样策略。Consistency distillation 从预训练的扩散模型中提取知识Consistency training 直接从头训练一致性模型这两种方法都能在推理时将采样步骤大幅减少。项目提供了多组预训练模型覆盖了三种数据集和不同训练配置。EDM 模型基于扩散模型CD 模型通过蒸馏获得CT 模型通过直接训练获得。每种配置都有对应的 checkpoint 可供下载用户可以根据任务需求选择合适的模型。安装方式有两种。通过 pip 直接安装pipinstall-e.或者使用 Dockercddockermakebuildmakerun项目代码中已经包含了 EDM 训练、一致性蒸馏、一致性训练、单步生成和多步生成的完整示例脚本位于 scripts/launch.sh 文件中。评估部分支持 FID、Precision、Recall 和 Inception Score 等指标与 guided-diffusion 的评估流程保持一致。Consistency Models 也已经集成到了 Hugging Face 的 diffusers 库中通过 ConsistencyModelPipeline 类即可调用importtorchfromdiffusersimportConsistencyModelPipeline devicecudamodel_id_or_pathopenai/diffusers-cd_imagenet64_l2pipeConsistencyModelPipeline.from_pretrained(model_id_or_path,torch_dtypetorch.float16)pipe.to(device)# 单步采样imagepipe(num_inference_steps1).images[0]image.save(consistency_model_onestep_sample.png)# 多步采样imagepipe(timesteps[22,0]).images[0]image.save(consistency_model_multistep_sample.png)单步采样即可生成图像多步采样则能进一步提升图像质量。项目还提供了类别条件生成的示例可以指定 ImageNet 类别标签来生成特定类型的图像。对于需要快速图像生成的应用场景Consistency Models 提供了一个可行的技术路径。相比传统扩散模型几十甚至上百步的采样过程单步生成在速度上有明显优势适合对延迟敏感的生产环境。路径。相比传统扩散模型几十甚至上百步的采样过程单步生成在速度上有明显优势适合对延迟敏感的生产环境。