AsyncDiff：异步去噪加速扩散模型，实现多设备并行处理

项目介绍

AsyncDiff 是一个创新的开源项目，旨在通过异步去噪技术加速扩散模型的推理过程。该项目由新加坡国立大学学习与视觉实验室（Learning and Vision Lab）的研究团队开发，核心成员包括 Zigeng Chen、Xinyin Ma、Gongfan Fang、Zhenxiong Tan 和 Xinchao Wang。AsyncDiff 通过将复杂的去噪模型分割成多个组件，并将每个组件分配到不同的设备上进行并行计算，从而显著减少了推理延迟，同时对生成质量的影响极小。

项目技术分析

AsyncDiff 的核心技术在于将传统的顺序去噪过程转变为异步处理。通过利用连续扩散步骤中隐藏状态之间的高度相似性，AsyncDiff 打破了组件之间的依赖链，使得每个组件可以在不同的设备上并行计算。具体来说，AsyncDiff 将去噪模型 εθ 分割成多个组件，并在预热阶段提前准备好每个组件的输入，从而实现并行处理。

项目支持多种扩散模型，包括 Stable Diffusion 3 Medium、Stable Diffusion 2.1、Stable Diffusion 1.5、Stable Diffusion x4 Upscaler、Stable Diffusion XL 1.0、Stable Diffusion XL Inpainting、ControlNet、Stable Video Diffusion 和 AnimateDiff。

项目及技术应用场景

AsyncDiff 适用于需要高效处理大规模扩散模型推理的场景，特别是在以下领域：

• 图像生成与处理：加速 Stable Diffusion 系列模型的图像生成和处理，适用于图像编辑、超分辨率、图像修复等应用。
• 视频生成与处理：加速 Stable Video Diffusion 和 AnimateDiff 模型的视频生成和处理，适用于视频编辑、动画生成等应用。
• 多设备并行计算：适用于拥有多台 GPU 设备的计算集群或数据中心，通过并行计算提高推理效率。

项目特点

• 高效并行处理：通过异步去噪技术，实现多设备并行计算，显著减少推理延迟。
• 广泛兼容性：支持多种主流扩散模型，包括 Stable Diffusion、ControlNet、Stable Video Diffusion 等。
• 易于集成：只需添加两行代码即可为现有的扩散模型启用异步并行推理功能。
• 灵活配置：用户可以根据需求调整模型分割数量、去噪步长、预热步数等参数，以平衡速度和生成质量。

快速开始

安装

1. 环境要求：

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

2. 创建环境：

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

使用示例

以下是一个简单的使用示例，展示如何为 Stable Diffusion 模型启用异步并行推理：

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

pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", 
torch_dtype=torch.float16, use_safetensors=True, low_cpu_mem_usage=True)

async_diff = AsyncDiff(pipeline, model_n=2, stride=1, time_shift=False)

async_diff.reset_state(warm_up=1)
image = pipeline(<prompts>).images[0]
if dist.get_rank() == 0:
  image.save(f"output.jpg")

加速推理

项目提供了详细的脚本，用于加速多种扩散模型的推理，包括 Stable Diffusion XL、Stable Diffusion 2.1、Stable Diffusion 1.5、Stable Diffusion 3 Medium、ControlNet、Stable Diffusion x4 Upscaler、AnimateDiff 和 Stable Video Diffusion。

例如，加速 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

结语

AsyncDiff 通过创新的异步去噪技术，为扩散模型的推理提供了高效的并行处理方案。无论是在图像生成、视频处理还是多设备并行计算领域，AsyncDiff 都能显著提升推理效率，同时保持生成质量。如果你正在寻找一种高效的方式来加速扩散模型的推理，AsyncDiff 绝对值得一试！