UniPC: 用于快速采样扩散模型的统一预测-校正框架

由 Wenliang Zhao*、Lujia Bai*、Yongming Rao、Jie Zhou、Jiwen Lu 创建

本代码包含 UniPC (NeurIPS 2023) 的 Pytorch 实现。

UniPC 与 stable-diffusion 结合的在线演示。非常感谢 HuggingFace 🤗 提供的帮助和硬件资源支持。

UniPC 是一个无需训练的框架，专为扩散模型的快速采样而设计，由校正器 (UniC) 和预测器 (UniP) 组成，它们共享统一的分析形式并支持任意阶数。

[项目页面] [arXiv]

新闻

• 2023.3.11: UniPC 已在 stable-diffusion-webui 中使用。
• 2023.2.20: UniPC 已集成到 🧨 Diffusers 中，可通过 UniPCMultistepScheduler 访问。

UniPC 在设计上是与模型无关的，支持像素空间/潜在空间 DPM 的无条件/条件采样。它也可以应用于噪声预测模型和数据预测模型。

与之前的方法相比，由于精度阶数的提高，UniPC 收敛更快。定量和定性结果都表明 UniPC 可以显著提高采样质量，尤其是在极少步骤（5~10）的情况下。

代码示例

我们在 example 文件夹中提供了基于 ScoreSDE 和 Stable-Diffusion 的代码示例。请按照相应示例中的 README.md 文件进行进一步操作以使用我们的 UniPC。

使用 UniPC 的 ScoreSDE

我们在 example/score_sde_pytorch 中提供了一个 pytorch 示例，展示了如何使用我们的 UniPC 从预训练的 CIFAR10 DPM 中进行采样。

使用 UniPC 的 Stable-Diffusion

我们在 example/stable-diffusion 中提供了将 UniPC 应用于 stable-diffusion 的示例。我们的 UniPC 可以加速条件和无条件采样。

与 🤗 Diffusers 库的集成

UniPC 现已在 🧨 Diffusers 中可用，可通过 UniPCMultistepScheduler 访问。 Diffusers 允许您仅用几行代码就可以在 PyTorch 中测试 UniPC。

您可以按如下方式安装 diffusers：

pip install diffusers accelerate transformers

然后只需几行代码就可以尝试该模型：

from diffusers import StableDiffusionPipeline, UniPCMultistepScheduler
import torch

path = "CompVis/stable-diffusion-v1-4"

pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)

# 更改为 UniPC 调度器
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

prompt = "a highly realistic photo of green turtle"
generator = torch.manual_seed(0)

# 只需 15 步就能获得良好结果 => GPU 上 2-4 秒
image = pipe(prompt, generator=generator, num_inference_steps=15).images[0]

# 保存图像
image.save("turtle.png")

有关 UniPC 和 diffusers 的更多信息，请查看这里和这里。

致谢

我们的代码基于 ScoreSDE、Stable-Diffusion 和 DPM-Solver。

引用

如果您发现我们的工作对您的研究有用，请考虑引用：

@article{zhao2023unipc,
  title={UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models},
  author={Zhao, Wenliang and Bai, Lujia and Rao, Yongming and Zhou, Jie and Lu, Jiwen},
  journal={NeurIPS},
  year={2023}
}