访问官网
Github
Huggingface
论文DistriFusion：高分辨率扩散模型的分布式并行推理
论文 | 项目 | 博客 | 幻灯片 | YouTube | 海报
[2024年7月29日] DistriFusion 已在 ColossalAI 中得到支持！
[2024年4月4日] DistriFusion 被选为 CVPR 2024 亮点海报展示！
[2024年2月29日] DistriFusion 被 CVPR 2024 接收！我们的代码已公开发布！
我们提出了 DistriFusion，这是一种免训练的算法，可利用多个 GPU 加速扩散模型推理，同时不牺牲图像质量。简单的图像块处理（概览(b)）由于缺乏图像块之间的交互而存在分割问题。展示的示例是使用 SDXL 在 1280×1920 分辨率下生成的，采用 50 步 Euler 采样器，延迟在 A100 GPU 上测量。
DistriFusion：高分辨率扩散模型的分布式并行推理
李牧洋*，蔡天乐*，曹嘉欣，张钦胜，蔡涵，白俊杰，贾扬清，刘明宇，李凯，以及 宋涵
麻省理工学院、普林斯顿大学、Lepton AI 和 NVIDIA
发表于 CVPR 2024。
概述
(a) 在单个设备上运行的原始扩散模型。(b) 简单地将图像分割成 2 个块并分布在 2 个 GPU 上，由于块之间缺乏交互，边界处会出现明显的接缝。(c) 我们的 DistriFusion 在第一步使用同步通信进行块交互。之后，我们通过异步通信重用前一步的激活。通过这种方式，通信开销可以隐藏在计算流水线中。
性能
加速效果
质量
SDXL 的定性结果。FID 是与真实图像相比计算得出的。我们的 DistriFusion 可以根据使用的设备数量减少延迟，同时保持视觉保真度。
参考文献：
- 去噪扩散隐式模型（DDIM），Song 等人，ICLR 2021
- 阐明基于扩散的生成模型的设计空间，Karras 等人，NeurIPS 2022
- 扩散模型的并行采样，Shih 等人，NeurIPS 2023
- SDXL：改进高分辨率图像合成的潜在扩散模型，Podell 等人，ICLR 2024
先决条件
- Python3
- NVIDIA GPU + CUDA >= 12.0 及相应的 CuDNN
- PyTorch = 2.2
入门指南
安装
安装 PyTorch 后，您应该能够通过 PyPI 安装 distrifuser
pip install distrifuser
或通过 GitHub：
pip install git+https://github.com/mit-han-lab/distrifuser.git
或在本地进行开发安装
git clone git@github.com:mit-han-lab/distrifuser.git
cd distrifuser
pip install -e .
使用示例
在 scripts/sdxl_example.py 中，我们提供了一个使用 DistriFusion 运行 SDXL 的最小脚本。
import torch
from distrifuser.pipelines import DistriSDXLPipeline
from distrifuser.utils import DistriConfig
distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)
pipeline = DistriSDXLPipeline.from_pretrained(
distri_config=distri_config,
pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0",
variant="fp16",
use_safetensors=True,
)
pipeline.set_progress_bar_config(disable=distri_config.rank != 0)
image = pipeline(
prompt="宇航员在丛林中，冷色调，柔和的颜色，细节丰富，8k",
generator=torch.Generator(device="cuda").manual_seed(233),
).images[0]
if distri_config.rank == 0:
image.save("astronaut.png")
具体来说，我们的 distrifuser 与 diffusers 共享相同的 API，可以以类似的方式使用。您只需定义一个 DistriFusion 并使用我们封装的 DistriSDXLPipeline 来加载预训练的 SDXL 模型。然后，我们可以像 diffusers 中的 StableDiffusionXLPipeline 一样生成图像。运行命令是
torchrun --nproc_per_node=$N_GPUS scripts/sdxl_example.py
其中 $N_GPUS 是您想要使用的 GPU 数量。
我们还在 scripts/sd_example.py 中提供了一个使用 DistriFusion 运行 SD1.4/2 的最小脚本。使用方法相同。
基准测试
我们的基准测试结果使用 PyTorch 2.2 和 diffusers 0.24.0。首先，您可能需要安装一些额外的依赖项：
pip install git+https://github.com/zhijian-liu/torchprofile datasets torchmetrics dominate clean-fid
COCO 质量
您可以使用 scripts/generate_coco.py 生成带有 COCO 描述的图像。命令是
torchrun --nproc_per_node=$N_GPUS scripts/generate_coco.py --no_split_batch
其中$N_GPUS是您想使用的GPU数量。默认情况下，生成的结果将存储在results/coco中。您也可以使用--output_root自定义它。一些您可能想调整的额外参数：
--num_inference_steps：推理步骤数。我们默认使用50步。--guidance_scale：无分类器引导尺度。我们默认使用5。--scheduler：扩散采样器。我们默认使用DDIM采样器。您也可以使用euler表示Euler采样器和dpm-solver表示DPM求解器。--warmup_steps：额外预热步骤数（默认为4）。--sync_mode：不同的GroupNorm同步模式。默认使用我们修正的异步GroupNorm。--parallelism：您使用的并行范式。默认为补丁并行。您可以使用tensor表示张量并行和naive_patch表示朴素补丁。
生成所有图像后，您可以使用我们的脚本scripts/compute_metrics.py来计算PSNR、LPIPS和FID。用法如下：
python scripts/compute_metrics.py --input_root0 $IMAGE_ROOT0 --input_root1 $IMAGE_ROOT1
其中$IMAGE_ROOT0和$IMAGE_ROOT1是您想比较的图像文件夹路径。如果IMAGE_ROOT0是真实值文件夹，请添加--is_gt标志以进行调整大小。我们还提供了一个脚本scripts/dump_coco.py来导出真实值图像。
延迟
您可以使用scripts/run_sdxl.py来对我们不同方法的延迟进行基准测试。命令如下：
torchrun --nproc_per_node=$N_GPUS scripts/run_sdxl.py --mode benchmark --output_type latent
其中$N_GPUS是您想使用的GPU数量。与scripts/generate_coco.py类似，您也可以更改一些参数：
--num_inference_steps：推理步骤数。我们默认使用50步。--image_size：生成的图像大小。默认为1024×1024。--no_split_batch：禁用无分类器引导的批次分割。--warmup_steps：额外预热步骤数（默认为4）。--sync_mode：不同的GroupNorm同步模式。默认使用我们修正的异步GroupNorm。--parallelism：您使用的并行范式。默认为补丁并行。您可以使用tensor表示张量并行和naive_patch表示朴素补丁。--warmup_times/--test_times：预热/测试运行次数。默认分别为5和20次。
引用
如果您在研究中使用此代码，请引用我们的论文。
@inproceedings{li2023distrifusion,
title={DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models},
author={Li, Muyang and Cai, Tianle and Cao, Jiaxin and Zhang, Qinsheng and Cai, Han and Bai, Junjie and Jia, Yangqing and Liu, Ming-Yu and Li, Kai and Han, Song},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
year={2024}
}
致谢
我们的代码基于huggingface/diffusers和lmxyy/sige开发。我们感谢torchprofile用于MACs测量，clean-fid用于FID计算，以及Lightning-AI/torchmetrics用于PSNR和LPIPS。
我们感谢Jun-Yan Zhu和Ligeng Zhu的有益讨论和宝贵反馈。该项目得到了MIT-IBM Watson AI Lab、亚马逊、MIT Science Hub和美国国家科学基金会的支持。
