MMDiT 和 Transformer 模型

相关源文件

• comfy/latent_formats.py
• comfy/ldm/wan/model.py
• comfy/model_base.py
• comfy/model_detection.py
• comfy/sd.py
• comfy/supported_models.py
• comfy/supported_models_base.py
• comfy/text_encoders/wan.py
• comfy/utils.py
• comfy_extras/nodes_camera_trajectory.py
• comfy_extras/nodes_wan.py
• nodes.py

概述

多模态扩散 Transformer (MMDiT) 是一种专门为扩散模型设计的、基于 Transformer 的架构，它可以处理多种输入模态，例如文本和图像。它扩展了扩散 Transformer (DiT) 架构，以支持跨模态交互，从而实现不同类型数据之间更有效的条件约束。

本文档记录了 ComfyUI 中 MMDiT 架构和 Transformer 模型，重点介绍了它们的实现、组件以及它们如何与整个系统集成。

来源: comfy/ldm/modules/diffusionmodules/mmdit.py745-1013

架构概述

MMDiT 在 Transformer 架构原理的基础上进行扩展，但专门针对扩散模型进行了优化。其核心是一个类似 UNet 的去噪模型，它使用 Transformer 块而不是卷积层来处理潜在表示。

该架构由几个关键组件组成

1. 输入处理:

   • 将输入潜在表示分割成 token
   • 对 timestep 进行嵌入以进行扩散条件约束
   • 处理上下文（例如文本嵌入）
   • 应用位置嵌入

2. Transformer 块:

   • 一系列块，处理上下文和潜在特征
   • 每个块使用自注意力机制和自适应层归一化
   • 联合处理支持模态间的交叉注意力

3. 输出处理:

   • 最后一层将处理后的特征转换回潜在空间
   • 将输出“反分割”回原始潜在格式

来源: comfy/ldm/modules/diffusionmodules/mmdit.py745-780 comfy/ldm/modules/diffusionmodules/mmdit.py982-1013

核心组件

Patch 嵌入与位置编码

MMDiT 模型首先使用 PatchEmbed 模块将输入潜在表示分割成 patch

位置编码系统使用正弦嵌入，可以根据输入尺寸动态裁剪和缩放，从而使模型能够处理不同分辨率的输入。

来源: comfy/ldm/modules/diffusionmodules/mmdit.py57-107 comfy/ldm/modules/diffusionmodules/mmdit.py120-201

Transformer 块

MMDiT 的核心是其基于 Transformer 的处理块

这些块同时处理上下文信息和输入特征

1. JointBlock: 管理上下文和输入特征的配对块
2. DismantledBlock: 核心 Transformer 单元，包含注意力和 MLP 路径
3. SelfAttention: 实现多头自注意力机制

每个块都使用自适应层归一化 (adaLN)，该归一化根据 timestep 嵌入来调节特征，使模型能够根据扩散过程调整其行为。

来源: comfy/ldm/modules/diffusionmodules/mmdit.py386-595 comfy/ldm/modules/diffusionmodules/mmdit.py645-667

块混合与注意力

MMDiT 的一个关键创新是其块混合机制，它支持上下文特征和输入特征之间的交叉注意力

过程如下：

1. 从上下文和输入特征中提取查询、键、值 (QKV)
2. 将它们连接起来以进行联合注意力计算
3. 使用组合特征计算注意力
4. 拆分结果并单独更新每个特征路径

这允许信息在模态之间双向流动。

来源: comfy/ldm/modules/diffusionmodules/mmdit.py598-642

调制技术

MMDiT 使用多种调制技术来控制特征转换

1. 自适应层归一化 (adaLN):

2. RMS 归一化：使用均方根归一化而非传统的层归一化来归一化特征

3. 自适应调制：每个块接收源自 timestep 嵌入的调制信号，这些信号控制特征转换

来源: comfy/ldm/modules/diffusionmodules/mmdit.py109-112 comfy/rmsnorm.py14-25

实现变体

ComfyUI 包含了 MMDiT 类似架构的多种实现

标准 MMDiT

标准实现侧重于上下文特征和输入特征的联合处理，并具有基于 timestep 的自适应层归一化。

主要特性

• 支持不同输入尺寸的灵活位置嵌入
• 可拆卸的块以实现高效处理
• 支持自定义扩散条件约束

来源: comfy/ldm/modules/diffusionmodules/mmdit.py745-1013

AuraFlow MMDiT

AuraFlow 的实现提供了一种不同的方法，具有不同的块结构

关键区别

• 独立的 MMDiTBlock 和 DiTBlock 类，用于不同的处理阶段
• 注意力的替代实现（SingleAttention/DoubleAttention）
• 位置嵌入和特征调制的方法不同

来源: comfy/ldm/aura/mmdit.py179-207 comfy/ldm/aura/mmdit.py305-498

GenMo 非对称联合模型

GenMo 实现提供了一种非对称注意力机制，该机制专门针对不同模态之间的不同维度进行设计

主要特性

• 文本和图像特征之间的非对称注意力
• 文本和图像模态的不同维度
• 旋转位置嵌入 (RoPE) 用于空间感知
• 专为视频生成任务设计

来源: comfy/ldm/genmo/joint_model/asymm_models_joint.py160-286 comfy/ldm/genmo/joint_model/asymm_models_joint.py319-557

集成到 ComfyUI

MMDiT 模型通过标准化的接口与 ComfyUI 的架构集成

集成通过以下方式实现

1. 模型接口：标准的 forward 方法，接收 latents、timesteps 和 context
2. Transformer 选项：配置对象，允许自定义模型行为
3. 控制系统：通过 Transformer 选项支持控制机制

示例 forward 接口

来源: comfy/ldm/modules/diffusionmodules/mmdit.py982-1013

技术细节

位置嵌入

MMDiT 实现了几种位置嵌入方案

1. 2D 正弦嵌入：用于空间维度的标准正余弦嵌入

2. 动态裁剪：将位置嵌入适应输入尺寸

3. 旋转位置嵌入 (RoPE)：在 GenMo 的实现中用于旋转感知

这些嵌入系统允许模型在分块处理的 Transformer 过程中保持空间感知。

来源: comfy/ldm/modules/diffusionmodules/mmdit.py120-201 comfy/ldm/genmo/joint_model/temporal_rope.py7-34

Patch 处理

模型通过以下方式在连续潜在空间和 token 空间之间进行转换

1. Patchify：将潜在表示转换为 token

2. Unpatchify：将 token 转换回潜在表示

这种双向转换使模型能够利用 Transformer 架构，同时保持与整体扩散流程的兼容性。

来源: comfy/ldm/modules/diffusionmodules/mmdit.py913-931

结论

MMDiT 模型代表了 ComfyUI 中一种先进的、基于 Transformer 的扩散模型方法。通过利用 Transformer 架构并针对扩散过程进行专门修改，它们提供了有效的多模态条件约束和特征处理。

不同的实现（标准 MMDiT、AuraFlow 和 GenMo）对同一基本概念提供了不同的方法，每种方法都有针对不同用例和模态的特定优化。这些模型是 ComfyUI 高级模型生态系统的关键组成部分，实现了超越标准 Stable Diffusion 模型的复杂生成能力。