使用xDiT实现ComfyUI Flux.1工作流的多GPU并行加速！

增强您的AIGC体验：利用xDiT在ComfyUI Flux.1工作流中实现多GPU并行！

在AIGC领域快速发展的今天，用户对高质量图像生成的需求日益增长，同时也对使用的便捷性和成本效益提出了更高的要求。以ComfyUI为代表的工作流生图场景，随着社区的持续建设，俨然已成为目前最为主流的用户平台。然而，ComfyUI作为原生单卡设计的平台，虽然在个人PC玩家中广受欢迎，但其处理能力仍受限于单GPU的性能。对于如Flux.1-dev这样的大型模型，其庞大的参数量和计算量使得在本地单卡上运行变得充满挑战。使用xDiT让ComfyUI工作流中使用多GPU并行计算，显著降低用户等待延迟，提供足够流畅的用户体验。我们展示了一个例子使Flux.1-dev单图生成时间由14秒降低至7秒。

阅读英文版本

图片与视频生成是大模型多模态应用的关键领域，它们使用DiTs为架构的扩散模型来处理超长输入序列。然而，由于Transformer的计算复杂度随序列长度呈平方级增长，导致DiTs在实际部署中的延迟问题尤为突出。xDiT作为一款开源的专为DiTs设计的高性能推理框架，它可以将DiT推理并行扩展至多卡乃至多机超大规模设备，结合编译优化技术，确保应用的实时性需求得以满足。

开源项目地址： https://github.com/xdit-project/xDiT

通过应用xDiT，我们成功在ComfyUI工作流中部署了多GPU并行处理，从而实现了以下成效：

速度更快&质量无损

在不影响图片生成质量的情况下，xDiT通过利用多GPU并行，显著地加速了ComfyUI中FLUX.1工作流的生成效率。以FLUX.1 dev版本为例，在20步的情况下，单图生成时间由13.87秒降低至6.98秒：

同时在提高生成速度的同时，我们也保证了生成图像的质量。使用相同的提示输入时，生成的图像在质量上没有明显差异，都表现出了明显的指令遵循能力：

提示： 一个宽敞的未来教室，墙上装饰着数字植物，天花板上悬挂着浮动的3D全息太阳系投影。一侧，一个大型触摸敏感的智能黑板用霓虹笔写着“XDiT”。中间有几个可变形的智能桌子，学生们通过增强现实眼镜进行互动学习。在角落里，一个小机器人正在为学生提供个性化学习支持。

简单&无缝

在初期，为了将xDiT的能力应用到ComfyUI中，我们以Pipeline为粒度，通过定制化实现XfuserPipelineLoader和XfuserSampler两个核心节点，基于HTTP服务，实现了在ComfyUI中利用xDiT的能力完成端到端的生成功能。

然而，这样高度定制化的实现，并不能很好地兼容庞大的ComfyUI的社区生态；并且端到端的生成逻辑，也在很大程度上与ComfyUI高度模块化的设计理念相悖。我们希望能够以最小的改动代价，对ComfyUI完成并行化的加速支持。为了实现这个目标，核心是需要完成对整个工作流中扩散模型（Diffusion Model）相关部分的改造，以标准的FLUX.1-dev官方工作流为例：

扩散模型通过#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等的集成适配，将陆续在后续的版本中迭代更新。

联系我们！

目前，我们这项功能仍处于演示开发和试验阶段。如果您对ComfyUI的xDiT并行版本感兴趣，我们欢迎您通过电子邮件jiaruifang@tencent.com与我们取得联系。