AsyncDiff: 异步去噪加速扩散模型的实现

1. 项目介绍

AsyncDiff 是一个通用且即插即用的扩散模型加速方案，它通过在多个设备上实现模型并行性来提高去噪效率。该方案通过将复杂的噪声预测模型分解为多个组件，并将每个组件分配给不同的设备，从而打破组件间的依赖链。利用连续扩散步骤中隐藏状态的高相似性，将传统的顺序去噪转变为异步过程，从而显著减少推理延迟，同时最小化对生成质量的影响。

2. 项目快速启动

在开始之前，请确保您的系统满足以下要求：

• NVIDIA GPU + CUDA >= 12.0 和相应的 CuDNN

创建环境

conda create -n asyncdiff python=3.10
conda activate asyncdiff
pip install -r requirements.txt

使用示例

以下是一个使用 Stable Diffusion pipeline 的示例，以启用扩散模型的异步并行推理：

import torch
from diffusers import StableDiffusionPipeline
from asyncdiff.async_sd import AsyncDiff

# 加载预训练的 Stable Diffusion 模型
pipeline = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1",
    torch_dtype=torch.float16,
    use_safetensors=True,
    low_cpu_mem_usage=True
)

# 初始化 AsyncDiff
async_diff = AsyncDiff(pipeline, model_n=2, stride=1, time_shift=False)
async_diff.reset_state(warm_up=1)

# 生成图像
image = pipeline("<prompts>")[0]

# 如果是分布式训练的第一节点，保存图像
if dist.get_rank() == 0:
    image.save("output.jpg")

您可以替换 pipeline 为其他版本的 Stable Diffusion pipeline，如 SD 2.1、SD 1.5、SDXL 或 SVD。

3. 应用案例和最佳实践

AsyncDiff 支持多种扩散模型的加速推理，以下是一些案例：

• 加速 Stable Diffusion XL 推理：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.run --nproc_per_node=4 --run-path examples/run_sdxl.py

• 加速 Stable Diffusion 2.1 或 1.5 推理：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.run --nproc_per_node=4 --run-path examples/run_sd.py

• 加速 Stable Diffusion 3 Medium 推理：

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --nproc_per_node=2 --run-path examples/run_sd3.py

更多案例和最佳实践可以在项目的 examples 目录中找到。

4. 典型生态项目

AsyncDiff 的生态项目包括但不限于以下几种：

• Stable Diffusion 2.1、1.5、3 Medium、SDXL、SDXL Inpainting
• ControlNet、SD Upscaler、AnimateDiff、Stable Video Diffusion

这些项目都可以在 AsyncDiff 的官方仓库中找到相应的脚本和示例。