Comfy-WaveSpeed开发指南：如何为自定义模型添加FBCache支持

Comfy-WaveSpeed是ComfyUI的一站式推理优化解决方案，提供通用、灵活且快速的性能增强功能。本文将详细介绍如何为自定义模型添加FBCache（First Block Cache）支持，以实现推理速度的显著提升。## 什么是FBCache？FBCache是Comfy-WaveSpeed中的核心优化技术，通过缓存模型第一层的计算结果并在后续步骤中复用，有效减少重复计算，从而加速

富艾霏

297人浏览 · 2026-03-14 03:20:47

富艾霏 · 2026-03-14 03:20:47 发布

Comfy-WaveSpeed开发指南：如何为自定义模型添加FBCache支持

【免费下载链接】Comfy-WaveSpeed [WIP] The all in one inference optimization solution for ComfyUI, universal, flexible, and fast. 项目地址: https://gitcode.com/gh_mirrors/co/Comfy-WaveSpeed

Comfy-WaveSpeed是ComfyUI的一站式推理优化解决方案，提供通用、灵活且快速的性能增强功能。本文将详细介绍如何为自定义模型添加FBCache（First Block Cache）支持，以实现推理速度的显著提升。

什么是FBCache？

FBCache是Comfy-WaveSpeed中的核心优化技术，通过缓存模型第一层的计算结果并在后续步骤中复用，有效减少重复计算，从而加速推理过程。实验数据显示，启用FBCache后，Flux模型的推理速度从1.63it/s提升至4.33it/s，效率提升超过165%。

图1：未启用FBCache的原始推理速度（1.63it/s）

图2：启用FBCache后的推理速度（4.33it/s）

FBCache工作原理

FBCache通过以下机制实现性能优化：

缓存第一层输出：在模型推理过程中，缓存Transformer或UNet的第一层输出
残差比较：通过比较连续步骤间的残差差异判断是否可以复用缓存
状态管理：维护缓存上下文，处理不同采样阶段的缓存有效性

核心实现位于first_block_cache.py和fbcache_nodes.py文件中，通过模块化设计支持多种模型架构。

为自定义模型添加FBCache支持的步骤

1. 准备工作

首先确保你的模型满足以下条件：

基于PyTorch框架实现
具有明确的层级结构（如Transformer blocks或UNet输入/输出块）
支持模型输入/输出的缓存机制

2. 实现模型适配代码

根据模型类型，你需要实现对应的FBCache适配代码：

对于UNet类模型

# 参考 first_block_cache.py 中的 create_patch_unet_model__forward 函数
def create_patch_your_unet_forward(model, residual_diff_threshold, validate_can_use_cache_function):
    # 实现自定义UNet模型的前向传播补丁
    # 1. 提取第一层输出作为缓存基准
    # 2. 实现剩余块的调用逻辑
    # 3. 添加残差比较和缓存逻辑
    pass

对于Transformer类模型

# 参考 CachedTransformerBlocks 类
class YourModelCachedBlocks(torch.nn.Module):
    def __init__(self, transformer_blocks, residual_diff_threshold):
        super().__init__()
        self.transformer_blocks = transformer_blocks
        self.residual_diff_threshold = residual_diff_threshold
        
    def forward(self, hidden_states, encoder_hidden_states):
        # 实现缓存逻辑
        # 1. 处理第一层并缓存结果
        # 2. 比较残差决定是否使用缓存
        # 3. 复用缓存或继续计算剩余块
        pass

3. 添加节点定义

在fbcache_nodes.py中添加你的模型支持：

class ApplyFBCacheOnYourModel(ApplyFBCacheOnModel):
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("YOUR_MODEL_TYPE",),
                # 添加你的模型特定参数
                "residual_diff_threshold": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 1.0, "step": 0.001}),
            }
        }
    
    def patch(self, model, residual_diff_threshold):
        # 应用你的模型补丁函数
        patch_forward = create_patch_your_model_forward(model, residual_diff_threshold)
        # 设置模型包装器
        model.set_model_unet_function_wrapper(your_model_function_wrapper)
        return (model,)