提升你的 AIGC 体验：在 ComfyUI Flux.1 中利用 xDiT 实现多 GPU 并行处理
ComfyUI 是最受欢迎的基于网页的扩散模型界面，经过优化以提高工作流程效率。然而，它为单个 GPU 的原生设计使其在面对当今大型 DiT 的需求时显得力不从心，导致像 Flux.1 这样的用户体验到不可接受的高延迟。借助 xDiT 的强大功能，我们成功地在 ComfyUI 中实现了多 GPU 并行处理工作流程，有效解决了 Flux.1 的性能挑战。
图像和视频生成是大型模型多模态应用的关键领域。它们使用基于 DiT 的扩散模型来处理超长输入序列。然而，由于 Transformer 的计算复杂度与序列长度成平方关系，DiT 在实际部署中的延迟问题尤为突出。xDiT 是一个开源的高性能推理框架，专门为 DiT 设计。它可以将 DiT 推理扩展到多卡甚至多机的超大规模设备。结合编译优化技术，确保满足应用的实时需求。
开源项目地址: https://github.com/xdit-project/xDiT
通过应用 xDiT，我们成功地在 ComfyUI 工作流程中部署了多 GPU 并行处理，从而取得了以下成果：
更快且保持质量
在不影响图像生成质量的情况下，XDiT 通过利用多 GPU 并行处理显著加快了 ComfyUI 中 FLUX.1 工作流程的生成效率。以 FLUX.1 开发版本为例，在 20 步的条件下，单张图像的生成时间从 13.87 秒减少到 6.98 秒。
在提高生成速度的同时，我们也确保了生成图像的质量。使用相同的提示输入，生成的图像在质量上没有显著差异，均展现出良好的指令遵循能力。
_提示:___一个宽敞的未来主义教室，墙上装饰着数字植物，天花板上悬挂着太阳系的浮动 3D 全息投影。一侧，一个大型触摸智能白板用霓虹笔写着“XDiT”。中间有几个可变形的智能桌子，学生们通过增强现实眼镜进行互动学习。在角落里，一个小机器人正在为学生提供个性化学习支持。
简单且无缝
在初始阶段，为了将 XDiT 的能力应用于 ComfyUI，我们通过定制实现了两个核心节点，XfuserPipelineLoader 和 XfuserSampler，以管道粒度进行处理。基于 HTTP 服务，我们利用 XDiT 的能力在 ComfyUI 中实现了端到端的生成功能。
然而，这种高度定制的实现与庞大的 ComfyUI 社区生态系统并不兼容；此外，端到端的生成逻辑在很大程度上与 ComfyUI 的高度模块化设计理念相悖，后者旨在实现更大的灵活性。
我们希望以最小的修改实现 ComfyUI 的并行加速支持。为实现这一目标，核心要求是完成整个工作流程中与扩散模型相关部分的转换。以标准的 FLUX.1-dev 官方 ComfyUI 工作流程为例：
扩散模型通过 #12 Load Diffusion Model 节点加载，经过 #30 Model Sampling Flux 和 #22 BasicGuider 节点处理和封装，然后传递给 #13 SamplerCustomAdvanced 进行多步去噪计算。然而，由于 ComfyUI 最初是为单 GPU 的个人计算机设计的，因此在不对现有工作流程进行重大调整的情况下，将属于单个节点的计算任务分配到多个 GPU 进行并行处理仍然是一个重大挑战。
为了解决这个问题，我们对核心的 #12 Load Diffusion Model 加载节点和计算节点 #13 SamplerCustomAdvanced 实施了定制优化。通过集成 Ray 框架的分布式计算能力，我们只需使用 XDiTUNetLoader 和 XDiTSamplerCustomAdvanced 作为替代品，便可自动将模型的计算任务分配到多个计算节点。得益于 XDiT 出色的并行处理性能，我们显著提高了处理速度，有效优化了计算效率，同时不牺牲现有工作流程的稳定性和灵活性。
插件支持
对于一系列扩散模型，仅依靠提示和基本模型远远不足以达到生产级别。为此，ComfyUI 社区还支持一系列插件，包括 Loras、ControlNet 和 IP-Adapter，以更好地实现生成内容的风格化和可控性。
为此，XDiT 成功地将最受欢迎的 Loras 插件集成到 ComfyUI 中。现在，可以通过一个名为 XDiTFluxLoraLoader 的单一节点轻松加载多个 Loras。此外，由于 Loras 模块的模块化，在切换不同 Loras 时，无需更改模型内部的整体权重；只需调整新添加的 Loras 的权重。基于此，XDiT 可以支持在运行时动态切换 Loras，几乎瞬间完成切换，无需等待。
目前，我们正在持续努力完成其他常用社区插件（如 ControlNet 和 IP-Adapter）的集成和适配，这些将会在后续版本中迭代更新。
