标签归档:AI架构

ComfyUI 的缓存架构和实现

围绕 ComfyUI,大家讨论最多的是节点、工作流、算力这些,真正去看缓存细节的人其实不多。但只要你开始在一台机器上堆多个模型、多个 LoRA、多个 workflow,缓存策略就会直接决定这几件事:

  • 你是算力浪费,还是把显存 / 内存用在刀刃上;
  • 容器会不会莫名其妙 OOM;
  • 工作流切换时,是“秒级热身”还是“从头再来”。

这篇文章只做一件事:把 ComfyUI 当前的缓存架构和实现讲清楚,重点是三类策略(CLASSIC / LRU / RAM_PRESSURE)在「键怎么算、什么时候命中、什么时候过期」上的差异,以及在多模型、多 LoRA、多 workflow 场景下应该怎么选择。

1. 整体架构:两路缓存 + 四种策略

先把整体结构捋清楚。

1.1 两类缓存:outputs 和 objects

在执行一个工作流的时候,ComfyUI 维护了两类缓存:

  • outputs
    存中间结果和 UI 输出。命中时可以直接跳过节点执行,这部分是真正省计算的地方。
  • objects
    存节点对象实例(类实例),避免每次重新构造节点对象。

这两类缓存由一个统一的集合类 CacheSet 来管理。核心结构在 execution.py 中:

class CacheType(Enum):
    CLASSIC = 0
    LRU = 1
    NONE = 2
    RAM_PRESSURE = 3

class CacheSet:
    def __init__(self, cache_type=None, cache_args={}):
        if cache_type == CacheType.NONE:
            self.init_null_cache()
        elif cache_type == CacheType.RAM_PRESSURE:
            cache_ram = cache_args.get("ram"16.0)
            self.init_ram_cache(cache_ram)
        elif cache_type == CacheType.LRU:
            cache_size = cache_args.get("lru"0)
            self.init_lru_cache(cache_size)
        else:
            self.init_classic_cache()
        self.all = [self.outputs, self.objects]

不管是哪种策略,结构都是:

  • outputs按输入签名做 key
  • objects按 (node_id, class_type) 做 key

1.2 四种策略:CLASSIC / LRU / RAM_PRESSURE / NONE

从启动参数到缓存策略的选择,大致是这样的优先级(在 main.py 中):

  • 指定 --cache-lru → 用 LRU;
  • 否则指定 --cache-ram → 用 RAM_PRESSURE;
  • 否则指定 --cache-none → 完全关闭缓存;
  • 都没指定 → 默认 CLASSIC。

CacheType 的定义前面已经贴了。不同策略只决定 outputs 的实现方式,而 objects 基本始终使用层级缓存 HierarchicalCache(CacheKeySetID),后面细讲。


2. 缓存键:输入签名 + 变化指纹

缓存是否命中,首先取决于“key 算得是否合理”。ComfyUI 的设计核心是:

  • 输出缓存(outputs):用“输入签名 + is_changed 指纹(+ 某些情况下的 node_id)”作为 key;
  • 对象缓存(objects):用 (node_id, class_type) 作为 key。

2.1 输出缓存:输入签名是怎么来的

输出键的生成由 CacheKeySetInputSignature 负责,核心逻辑在 comfy_execution/caching.py:100-126

async def get_node_signature(self, dynprompt, node_id):
    signature = []
    ancestors, order_mapping = self.get_ordered_ancestry(dynprompt, node_id)
    signature.append(await self.get_immediate_node_signature(dynprompt, node_id,
    order_mapping))
    for ancestor_id in ancestors:
        signature.append(await self.get_immediate_node_signature(dynprompt,
        ancestor_id, order_mapping))
    return to_hashable(signature)

这里有几个点:

  1. 不只看当前节点
    它会按拓扑顺序把“当前节点 + 所有祖先”的签名拼在一起。这保证了只要整个子图的拓扑和输入一致,输出就能命中缓存。

  2. immediate 签名包含什么
    get_immediate_node_signature 会把这些信息打包进去(位置见同文件 108–126):

    • class_type:节点类型;
    • is_changed 指纹(通过 fingerprint_inputs/IS_CHANGED);
    • 必要时的 node_id
    • 有序的输入值/链接。
  3. 什么时候把 node_id 也算进 key
    当节点被声明为非幂等,或者内部有 UNIQUE_ID 这种隐含输入时,会把 node_id 加进签名(见 comfy_execution/caching.py:18-23, 116),避免“看起来一样”的节点被错误复用。

最后通过 to_hashable 转成可哈希结构(tuple/frozenset 等),作为最终键值。

结果是:

  • 输入完全相同 → key 相同 → 直接命中;
  • 任意一个上游节点输入或参数变化 → 指纹变了 → key 不同 → 不会复用旧结果。

2.2 对象缓存:用 (node_id, class_type)

节点对象的键由 CacheKeySetID 构造(comfy_execution/caching.py:66-80),逻辑简单:

  • key = (node_id, class_type)

读取时:

obj = caches.objects.get(unique_id)
if obj is None:
    obj = class_def()
    caches.objects.set(unique_id, obj)

对象缓存存在的目的只有一个:同一个 workflow 执行过程中,不要重复 new 节点实例。


3. 缓存容器:Basic / Hierarchical / LRU / RAM / Null

有了 key,还需要一个合理的「容器」和「驱逐策略」。

主要类在 comfy_execution/caching.py

  • BasicCache:基础容器,提供 set_prompt / clean_unused / get / set 等;
  • HierarchicalCache:按“父节点 → 子图”构建层级缓存;
  • LRUCache:在 Basic + Hierarchical 的基础上增加代际 LRU;
  • RAMPressureCache:在 LRU 的基础上增加 RAM 压力驱逐;
  • NullCache:空实现(禁用缓存)。

核心接口统一在 BasicCache

def clean_unused(self):
    self._clean_cache()
    self._clean_subcaches()

层级相关逻辑在 HierarchicalCache,用来支持“子图单独分区缓存”。

3.1 层级缓存:怎么定位到某个节点的分区

HierarchicalCache 通过 parent 链定位子缓存,代码在 comfy_execution/caching.py:242-269

def _get_cache_for(self, node_id):
    parent_id = self.dynprompt.get_parent_node_id(node_id)
    ...
    for parent_id in reversed(hierarchy):
        cache = cache._get_subcache(parent_id)
        if cache is Nonereturn None
    return cache

def get(self, node_id):
    cache = self._get_cache_for(node_id)
    if cache is Nonereturn None
    return cache._get_immediate(node_id)

def set(self, node_id, value):
    cache = self._get_cache_for(node_id)
    assert cache is not None
    cache._set_immediate(node_id, value)

含义很直接:

  • 根节点存在当前 cache;
  • 某些节点生成子图,会在该节点下挂一个子 cache;
  • 读写时,先通过 parent 链找到对应的子 cache 再读写。

clean_unused() 里除了清除不用的 key,还会删除没用到的子 cache 分区(_clean_subcaches())。

4. 执行循环中的使用路径

缓存不是“挂在那就完事了”,它在执行循环中有比较明确的调用点。

4.1 设置 prompt + 清理

启动一次执行时(execution.py:681-685):

is_changed_cache = IsChangedCache(prompt_id, dynamic_prompt, self.caches.outputs)
for cache in self.caches.all:
    await cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache)
    cache.clean_unused()

这里做了三件事:

  1. 给当前 prompt 生成一个 IsChangedCache,为 key 计算提供 is_changed 结果;
  2. 对 outputs / objects 各自执行一次 set_prompt(不同策略实现不同);
  3. 紧接着执行 clean_unused(),做一次基于“当前 prompt 键集合”的清理。

4.2 节点执行前后:cache 命中与写入

在节点执行路径中:

  • 执行前:优先尝试从 outputs 命中(execution.py:686-701);
  • 执行后:将 (ui, outputs) 作为 CacheEntry 写入 outputsexecution.py:568-571)。

为了简化,这里只看抽象行为:

  • 命中 → 跳过计算,直接拿值;
  • 未命中 → 正常跑一遍,将结果塞回缓存。

4.3 每个节点之后的 RAM 轮询

如果是 RAM_PRESSURE 模式,执行完每个节点都会触发一次内存检查(execution.py:720):

self.caches.outputs.poll(ram_headroom=self.cache_args["ram"])

只有 RAMPressureCache 实现了 poll,其他模式下这个调用等同空操作。


5. CLASSIC:默认层级缓存

先看默认策略:CLASSIC。

5.1 初始化与结构

CacheSet.init_classic_cacheexecution.py:97-126):

class CacheType(Enum):
    CLASSIC = 0
    LRU = 1
    NONE = 2
    RAM_PRESSURE = 3

class CacheSet:
    def init_classic_cache(self):
        self.outputs = HierarchicalCache(CacheKeySetInputSignature)
        self.objects = HierarchicalCache(CacheKeySetID)

可以看到:

  • outputs 用层级缓存 + 输入签名做 key;
  • objects 也用层级缓存 + (node_id, class_type) 做 key;
  • 不涉及 LRU 或 RAM 驱逐。

5.2 CLASSIC 的过期机制:完全由「当前 prompt」驱动

CLASSIC 模式不做容量和时间管理,它只有两种“失效”方式:

  1. 提示切换 / 执行前清理

clean_unused() 的核心逻辑(comfy_execution/caching.py:172-195):

def clean_unused(self):
    self._clean_cache()
    self._clean_subcaches()
  • _clean_cache():把不在“当前 prompt 键集合”的项删掉;
  • _clean_subcaches():把不再需要的子缓存分区删掉。

execution.py:681-685 每次绑定新 prompt 时,都会执行这一步。结果是:

  • 换了一个新的 workflow / prompt,旧 workflow 的 outputs / objects 都会被视为“未使用”,被清理掉;
  • CLASSIC 不会跨不同 prompt 保留旧 workflow 的缓存。
  1. 键不命中(指纹失效)

is_changed 的计算在 execution.py:48-89,当节点输入更新时,指纹会变化;CacheKeySetInputSignature 在构造键时会把这个指纹带进去(comfy_execution/caching.py:115-127)。因此只要:

  • 参数 / 输入 / 上游节点的任一变化 → key 改变 → 旧值自然不命中。

5.3 CLASSIC 明确不会做的事

在 CLASSIC 下:

  • 不做 LRU 容量控制:没有 max_size,也没有代际淘汰逻辑;
  • 不做 RAM 压力驱逐:poll() 是空的,执行循环里即使调用了也什么都不干;
  • 不做 TTL:不看时间,只看 prompt 键集合。

对应的注释已经在参考内容中点得很清楚,这里就不重复堆代码了。

5.4 在频繁切换 workflow / 模型时的表现

结合上面的机制,总结一下 CLASSIC 在多 workflow 场景下的行为:

  • 执行新工作流时:
    set_prompt + clean_unused 会直接把“不在新 prompt 键集合里”的缓存项(包括对象子缓存)全部清掉;
  • 模型 / LoRA 变化:
    即便节点 ID 不变,输入签名和 is_changed 指纹不同,也会生成新键;旧条目先不命中,随后在下一次 prompt 绑定时被清空;
  • 回切旧 workflow:
    因为在上一次切换时已经把旧 workflow 相关缓存清干净了,所以基本等于重新计算。

适用场景

  • workflow 比较固定;
  • 主要想在“一次执行当中”复用中间结果,不在意跨 prompt 的持久化。

6. LRU:代际 LRU 控制 outputs 尺寸

第二种策略是 LRU,主要解决的问题是:在允许跨 prompt 复用输出的前提下,限制缓存总量,避免无限膨胀

6.1 初始化:只作用于 outputs

CacheSet.init_lru_cacheexecution.py:127-135):

def init_lru_cache(self, cache_size):
    self.outputs = LRUCache(CacheKeySetInputSignature, max_size=cache_size)
    self.objects = HierarchicalCache(CacheKeySetID)

注意几点:

  • LRU 只作用于 outputs
  • objects 仍然用 HierarchicalCache,不受 LRU 驱逐;
  • 启动方式:--cache-lru N,且 N > 0

6.2 LRU 的代际设计

LRUCache 的骨架逻辑在 comfy_execution/caching.py:299-337

class LRUCache(BasicCache):
    def __init__(self, key_class, max_size=100):
        self.max_size = max_size
        self.min_generation = 0
        self.generation = 0
        self.used_generation = {}
        self.children = {}

    async def set_prompt(...):
        self.generation += 1
        for node_id in node_ids:
            self._mark_used(node_id)

    def get(self, node_id):
        self._mark_used(node_id)
        return self._get_immediate(node_id)

    def _mark_used(self, node_id):
        cache_key = self.cache_key_set.get_data_key(node_id)
        if cache_key is not None:
            self.used_generation[cache_key] = self.generation

含义:

  • 有一个全局代数 generation,每次绑定新 prompt generation += 1
  • 每次:

    • 绑定 prompt 时,会把该 prompt 内所有节点标记为“在当前代被使用”;
    • get / set 时更新条目的 used_generation[key] 为当前代。

6.3 容量驱逐:按「最老代」逐步清理

clean_unused() 的一部分逻辑在 comfy_execution/caching.py:314-323

def clean_unused(self):
    while len(self.cache) > self.max_size and self.min_generation < self.
    generation:
        self.min_generation += 1
        to_remove = [key for key in self.cache if self.used_generation[key] < self.
        min_generation]
        for key in to_remove:
            del self.cache[key]
            del self.used_generation[key]
            if key in self.children:
                del self.children[key]
    self._clean_subcaches()

简单归纳一下:

  • 只要 len(cache) > max_size,就逐步提升 min_generation
  • 每提升一代,就删除“最近使用代 < min_generation”的条目;
  • 同步清除掉和这些 key 绑定的子缓存引用;
  • 清理完再执行 _clean_subcaches() 做层级清扫。

6.4 子图分区与代际配合

子缓存创建时会显式标记父节点和子节点“被使用”,避免刚刚生成的子图被误删。代码在 comfy_execution/caching.py:338-349

async def ensure_subcache_for(self, node_id, children_ids):
    await super()._ensure_subcache(node_id, children_ids)
    await self.cache_key_set.add_keys(children_ids)
    self._mark_used(node_id)
    cache_key = self.cache_key_set.get_data_key(node_id)
    self.children[cache_key] = []
    for child_id in children_ids:
        self._mark_used(child_id)
        self.children[cache_key].append(self.cache_key_set.get_data_key(child_id))
    return self

配合前面提到的层级结构,就形成了“按 workflow 子图分区 + LRU 按代际清理”的整体行为。

6.5 触发时机和行为总结

在 LRU 模式下:

  • 每次绑定 prompt:
    generation += 1,标记当前 prompt 的节点使用代为当前代;
  • 每次 get/set
    更新条目的 used_generation 为当前代;
  • 每次 clean_unused()

    • len(cache) > max_size 时,通过提升 min_generation 清除旧代条目;
    • 额外清理无用子缓存。

特点

  • 可以跨 prompt 保留一部分中间结果;
  • max_size 控制缓存上限;
  • 没有 RAM 压力感知:poll() 依然不做事。

适用场景

  • 希望在多 workflow 之间部分复用缓存;
  • 但机器内存有限,需要给 outputs 一个明确的容量上限;
  • 对 RAM 细粒度控制没有强需求,或使用的是物理机 / 内存足够的环境。

7. RAM_PRESSURE:按可用内存压力驱逐

第三种策略是 RAM_PRESSURE,对应类是 RAMPressureCache。它继承自 LRUCache,但不按 max_size 做驱逐,而是:

  • 通过 poll(ram_headroom),在可用内存不足时按“OOM 评分”驱逐条目。

7.1 初始化:objects 仍然是层级缓存

CacheSet.init_ram_cacheexecution.py:131-133):

def init_ram_cache(self, min_headroom):
    self.outputs = RAMPressureCache(CacheKeySetInputSignature)
    self.objects = HierarchicalCache(CacheKeySetID)

注意两个点:

  • RAM 模式下,只有 outputs 会按 RAM 压力驱逐
  • objects 不参与 RAM 驱逐,逻辑完全和 CLASSIC/LRU 下相同。

7.2 poll:可用 RAM 检测 + OOM 评分驱逐

poll 的主逻辑在 comfy_execution/caching.py:384-454

def poll(self, ram_headroom):
    def _ram_gb():
        # 优先 cgroup v2/v1,失败回退 psutil
        ...
    if _ram_gb() > ram_headroom: return
    gc.collect()
    if _ram_gb() > ram_headroom: return
    clean_list = []
    for key, (outputs, _), in self.cache.items():
        oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation -
        self.used_generation[key])
        ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
        def scan_list_for_ram_usage(outputs):
            nonlocal ram_usage
            ...
        scan_list_for_ram_usage(outputs)
        oom_score *= ram_usage
        bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
    while _ram_gb() < ram_headroom * RAM_CACHE_HYSTERESIS and clean_list:
        _, _, key = clean_list.pop()
        del self.cache[key]
        gc.collect()

流程拆一下:

  1. 获取可用 RAM

    _ram_gb() 的实现优先读取 cgroup 的限制:

    • cgroup v2:memory.max / memory.current
    • cgroup v1:memory.limit_in_bytes / memory.usage_in_bytes
    • 都失败才回退 psutil.virtual_memory().available

    这解决了容器环境下“宿主机内存大,容器实际被限制”的常见问题。

  2. 阈值和 GC

    • 如果可用 RAM > ram_headroom,直接返回;
    • 否则先跑一次 gc.collect()
    • 再测一次 RAM,如果还是不足,进入驱逐流程。
  3. 为每个条目计算 OOM 评分

    • 初始 oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (generation - used_generation[key])

      • 大概意思是:越久没被用,分数指数级放大(默认倍数为 1.3,见 comfy_execution/caching.py:365);
    • 初始 ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE(0.1,见 360);
    • 递归遍历 outputs 列表:

      • CPU tensor:numel * element_size * 0.5(认为 CPU 上的 tensor 价值更高,折半);
      • 自定义对象:如果实现了 get_ram_usage() 就加上它;
    • 最后 oom_score *= ram_usage,得到综合评分。

    所有条目按 (oom_score, timestamp, key) 排序,放入 clean_list

  4. 按迟滞阈值逐个删除

    • 只要 _ram_gb() < ram_headroom * RAM_CACHE_HYSTERESIS,就从 clean_list 末尾 pop 一个 key 并删除;
    • 每删一个都跑一次 gc.collect()
    • 迟滞倍数 RAM_CACHE_HYSTERESIS 默认 1.1,避免“刚删完又马上触发清理”的抖动。
  5. 访问时间戳

    访问时会更新 timestamps(comfy_execution/caching.py:376-382):

   def set(self, node_id, value):
       self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
       super().set(node_id, value)

   def get(self, node_id):
       self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
       return super().get(node_id)

在 oom_score 一样时,timestamp 起到“最近访问优先保留”的作用。

7.3 提示绑定下的行为:不清理 outputs

RAM 模式下,clean_unused() 的行为与 CLASSIC 不同(见参考说明):

  • RAM 模式:
    clean_unused() 只做子缓存分区清理,不会删掉“当前 prompt 未使用”的 outputs 条目;
  • CLASSIC 模式:
    clean_unused() 会同时删掉当前 prompt 未用到的 outputs 条目。

结果是:

  • RAM 模式可以跨多个 workflow 长时间保留中间结果;
  • 只有在 RAM 不够时才做清退。

7.4 容器环境下需要注意的点

的 AutoDL 中,用 psutil.virtual_memory().available,在容器里看到的是宿主机内存,而不是容器的限额,导致永远“不触发回收”,最后 OOM。

适用场景

  • 多 workflow / 多模型 / 多 LoRA 同时存在,且希望尽可能长时间复用输出结果;
  • 机器内存有限,但更关心“不 OOM”,而不是一个固定的 max_size
  • 特别适合容器环境(K8s / AutoDL 等),配合 --cache-ram <GB>

8. 一些落地建议

最后,用一段比较直接的建议收尾。

8.1 单模型 / workflow 稳定:用 CLASSIC 即可

特点:

  • 工作流基本不换;
  • 主要希望避免一次执行中的重复计算(比如多次 preview)。

用默认 CLASSIC:

  • 结构简单;
  • 不参与复杂的跨 prompt 保留;
  • 不需要担心 LRU 尺寸和 RAM 阈值调参。

8.2 多 workflow + 控制缓存尺寸:用 LRU

场景:

  • 有多套 workflow,在它们之间来回切;
  • 机器内存不是特别大,希望 outputs 不要无限膨胀;
  • 又希望某些常用子图能被复用。

做法:

  • 启动时加 --cache-lru N(N 先给一个相对保守的值,比如几百到几千条,看内存曲线再调);
  • LRUCache 用代际 + max_size 帮你自动做“近期常用保留、早期冷门清理”。

8.3 多模型 / 多 LoRA / 容器环境:优先 RAM_PRESSURE

场景:

  • 容器 / 云平台(K8s、AutoDL 等);
  • 有较多模型、LoRA 和 workflow 混用;
  • 内存被容器限制,容易因爆 RAM 掉进 OOM。

做法:

  • 启动时用 --cache-ram <GB> 配一个“希望保留的 RAM 余量”;

    • 比如容器给了 32GB,就设在 16–24GB 看情况;
  • RAMPressureCache

    • 最大程度保留各个 workflow 的中间结果(包括应用 LoRA 后的模型对象等);
    • 在内存不足时,根据 OOM 评分优先清旧代、大内存条目。

注意一点:即使有容器优化,如果平台本身的 cgroup 挂得不标准,_ram_gb() 的结果还是有可能偏离实际,这一点要结合平台文档确认。

8.4 完全不想折腾:直接关缓存

如果你:

  • 环境受限;
  • 或者调试阶段不想被缓存影响行为理解;

可以直接用:

  • --cache-none
    对应 CacheType.NONECacheSet.init_null_cache()NullCache,所有 get/set 都是 no-op。

代价就是:每次执行都完全重新算一遍。

9. 小结一下

把上面的内容压缩成几句话:

  • ComfyUI 有两路缓存:

    • outputs 存中间结果,真正用来省算力;
    • objects 存节点实例,只减少 Python 对象构造开销。
  • 键体系:

    • 输出:输入签名 + is_changed 指纹(+ 条件下的 node_id)
    • 对象:**(node_id, class_type)**。
  • 四种策略:

    • CLASSIC:默认层级缓存,按当前 prompt 键集合清理,不做 LRU / RAM 驱逐;
    • LRU:只对 outputs 做代际 LRU,配合 max_size 控制容量;
    • RAM_PRESSURE:在 LRU 基础上加 RAM 压力驱逐,在内存不足时按 OOM 评分清理;
    • NONE:彻底关掉缓存。
  • 在多模型 / 多 workflow / 多 LoRA 场景下:

    • 对象缓存 objects 始终是层级缓存,不参与 LRU / RAM 驱逐,在每次绑定 prompt 时按键集合清理;
    • 模型权重 / LoRA 的真实驻留由模型管理层控制;
    • 真正要关心的是:如何选择 outputs 的策略,让中间结果既能有效复用,又不会把内存打爆。

如果我们在一个复杂的工作流环境里跑 ComfyUI,建议先搞清楚自己处于哪种场景,再结合上面的策略选项,把缓存调成我们能控制的状态,而不是让它在后台「自动长草」。

以上。

深入源码解析 ComfyUI 的模块化节点设计架构

ComfyUI 是一个基于 Stable Diffusion 的开源 AI 绘图工具,采用了模块化的节点式工作流设计。它通过将 Stable Diffusion 的各个组件和处理步骤抽象为独立的节点,使得用户可以通过直观的拖拽、连接操作来构建复杂的图像生成流程。

ComfyUI 解决了传统 AI 绘图工具易用性差、扩展性低的问题。其模块化设计和直观的 Web 界面大大降低了用户的使用门槛,无需深入了解底层技术细节,即可快速构建和调整工作流。同时,ComfyUI 还提供了强大的自定义节点机制,允许开发者轻松扩展新的功能和模型,使其能够适应不断发展的AI绘图领域。

ComfyUI 最初由开发者 Comfyanonymous 在 2022 年末发起,旨在提供一个简单、直观的 Stable Diffusion Web UI。早期版本实现了基本的节点类型和 Web 界面,展示了其模块化设计的优势,吸引了一批 AI 绘图爱好者的关注。

在 2023 年春夏,ComfyUI 进入了快速发展阶段。项目不断增加新的节点类型,如 ControlNet、Inpaint、Upscale等,支持更多的图像控制和后处理功能。同时,ComfyUI 引入了自定义节点机制,大大扩展了其功能和适用范围。项目也集成了更多 Stable Diffusion 衍生模型,为用户提供了更多选择。

随着用户社区的不断壮大,ComfyUI 的生态也日益丰富。社区成员积极贡献工作流、节点脚本、训练模型等资源,推动项目的发展。ComfyUI 举办了一系列社区活动,促进了用户间的交流和创作。项目代码库也迎来了更多贡献者,社区力量成为 ComfyUI 发展的重要推动力。

2023 年冬开始,ComfyUI 开始着眼于生态融合和应用拓展。项目与其他 AI 绘图工具建立了联系,支持工作流的导入导出和 API 集成。ComfyUI 也开始探索更多应用场景,如虚拟主播、游戏 mod 等,拓宽了 AI绘图的应用范围。越来越多的开发者和公司开始关注和使用 ComfyUI,其发展前景备受看好。

前两周,Comfy 推出跨平台的 ComfyUI 安装包,现在我们可以一键安装 ComfyUI 了,这次更新不仅带来了全新的桌面版应用,还对用户界面进行了全面升级,并新增了自定义按键绑定、自动资源导入等一系列实用功能,让工作流程更加流畅。

今天我们深入到 ComyUI 的源码去看一下其实现原理和过程。

ComfyUI 执行的大概流程如下:

用户界面 -> 工作流定义 -> PromptQueue
   ↓
PromptExecutor 开始执行
   ↓
验证输入 (validate_prompt)
   ↓
构建执行图
   ↓
按顺序执行节点
   ↓
缓存结果
   ↓
返回结果到界面

1. ComfyUI 的初始化与执行流程详解

ComfyUI 的一个明显的优点是有着灵活的图形用户界面,可以用于处理复杂的图像生成和处理工作流。

它具有精良的架构设计,通过模块化设计、缓存优化、资源管理以及错误处理机制,确保了系统的高效性和可靠性。

1.1 系统初始化流程

ComfyUI 的启动过程分为几个主要阶段:预启动脚本的执行、节点系统的初始化、服务器的启动与 WebSocket 的连接。

1. 启动前准备

在系统启动前,ComfyUI 会首先执行预启动脚本,确保自定义节点的环境准备就绪。这一过程允许在加载节点之前执行一些必要的自定义操作。

def execute_prestartup_script():
    # 执行自定义节点的预启动脚本
    for custom_node_path in node_paths:
        if os.path.exists(script_path):
            execute_script(script_path)

2. 节点系统初始化

节点系统是 ComfyUI 的核心组件。此阶段会加载内置节点以及用户自定义的节点,并注册到系统中供后续使用。

def init_extra_nodes():
    # 加载内置节点
    import_failed = init_builtin_extra_nodes()
    # 加载自定义节点
    init_external_custom_nodes()
  • 内置节点: 位于 comfy_extras 目录下,定义了基本的图像处理功能。
  • 自定义节点: 用户可以通过 custom_nodes 目录添加自定义节点,扩展系统的功能。

3. 服务器初始化

服务器初始化是启动 ComfyUI 的 Web 服务器的过程。它包括 WebSocket 的初始化,允许前端和后端实时通信。此外,执行队列也会在此阶段创建,用于管理节点的执行顺序和任务调度。

class PromptServer:
    def __init__(self, loop):
        # 初始化 Web 服务器
        self.app = web.Application()
        # 初始化 WebSocket
        self.sockets = dict()
        # 初始化执行队列
        self.prompt_queue = None

1.2 工作流执行流程

工作流执行是 ComfyUI 的核心功能,它包括从提交工作流到执行节点的整个过程。以下是工作流执行的几个关键步骤。

1. 工作流验证

首先,系统会对提交的工作流进行验证,确保节点的类型存在、节点连接有效,并且每个节点的输入符合要求。

def validate_prompt(prompt):
    # 1. 验证节点类型是否存在
    # 2. 验证是否有输出节点
    # 3. 验证节点输入
    return (valid, error, good_outputs, node_errors)

2. 执行准备

在验证通过后,系统会初始化执行环境。这包括创建动态的提示(DynamicPrompt),以及初始化缓存系统,以避免重复计算并提高执行效率。

def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
    # 1. 初始化执行环境
    with torch.inference_mode():
        # 2. 创建动态提示
        dynamic_prompt = DynamicPrompt(prompt)
        # 3. 初始化缓存
        is_changed_cache = IsChangedCache(dynamic_prompt, self.caches.outputs)

3. 节点执行

每个节点的执行流程包括获取节点的输入数据、检查是否有缓存的数据可以复用、执行节点逻辑、并缓存执行结果。节点执行是系统的核心环节,其过程如下:

def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results):
    # 1. 获取节点信息
    unique_id = current_item
    inputs = dynprompt.get_node(unique_id)['inputs']
    class_type = dynprompt.get_node(unique_id)['class_type']
    
    # 2. 检查缓存
    if caches.outputs.get(unique_id) is not None:
        return (ExecutionResult.SUCCESS, NoneNone)
    
    # 3. 获取输入数据
    input_data_all, missing_keys = get_input_data(inputs, class_def, unique_id, caches.outputs, dynprompt, extra_data)
    
    # 4. 执行节点
    output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all)
    
    # 5. 缓存结果
    caches.ui.set(unique_id, {...})
  1. 获取节点信息: 获取当前节点的输入和类型信息。
  2. 检查缓存: 如果节点的输出已经缓存,则直接返回缓存结果,避免重复执行。
  3. 获取输入数据: 从上一个节点或缓存中获取需要的输入数据。
  4. 执行节点: 调用节点的执行函数,处理输入并生成输出数据。
  5. 缓存结果: 将执行结果缓存,以便后续节点使用。

1.3 执行队列管理

ComfyUI 通过执行队列管理工作流中的节点执行顺序。每个节点的执行任务会被放入队列中,系统按顺序处理这些任务。

def prompt_worker(q, server):
    e = execution.PromptExecutor(server, lru_size=args.cache_lru)
    
    while True:
        # 1. 获取队列任务
        queue_item = q.get(timeout=timeout)
        
        # 2. 执行提示
        e.execute(item[2], prompt_id, item[3], item[4])
        
        # 3. 资源管理
        if need_gc:
            comfy.model_management.cleanup_models()
            gc.collect()
  • 获取队列任务: 从队列中取出下一个需要执行的节点任务。
  • 执行节点: 调用执行器执行当前节点。
  • 资源管理: 在必要时触发模型清理和垃圾回收,确保系统资源不被过度占用。

1.4 缓存系统

ComfyUI 的缓存系统采用层次化设计,可以缓存节点的输出、对象和 UI 相关的数据,极大地提高了执行效率。

class HierarchicalCache:
    def __init__(self):
        self.outputs = {}  # 节点输出缓存
        self.ui = {}  # UI 相关缓存
        self.objects = {}  # 节点对象缓存
  • 输出缓存(outputs): 缓存节点的执行结果,避免重复计算。
  • 对象缓存(objects): 缓存大数据对象,如模型和图像,以减少加载时间。
  • UI 缓存(ui): 缓存与前端界面相关的信息,如预览图像和执行状态。

1.5 资源管理与错误处理

为了确保系统的稳定性,ComfyUI 提供了完善的资源管理和错误处理机制。在执行工作流的过程中,系统会自动清理未使用的模型和缓存,并在必要时触发内存回收。

资源管理

资源管理包括内存清理、模型卸载以及缓存清理。系统会根据内存使用情况自动卸载不必要的模型,并定期触发垃圾回收。

# 1. 内存清理
comfy.model_management.cleanup_models()
gc.collect()

# 2. 模型卸载
comfy.model_management.unload_all_models()

# 3. 缓存清理
cache.clean_unused()

错误处理

ComfyUI 实现了详细的错误处理机制,能够捕获并处理执行过程中发生的各种异常。对于节点执行中的错误,系统会记录错误信息并通知用户。

def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, error, ex):
    # 1. 处理中断异常
    if isinstance(ex, comfy.model_management.InterruptProcessingException):
        self.add_message("execution_interrupted", mes)
    # 2. 处理其他错误
    else:
        self.add_message("execution_error", mes)
  • 处理中断异常: 当执行被中断时,系统会捕获异常并记录中断信息。
  • 处理其他错误: 处理其他执行错误,并通过 UI 向用户报告错误详情。

2. 节点系统架构

节点系统是 ComfyUI 的核心系统,其节点系统架构设计精巧,支持动态节点的加载、执行和扩展。今天我们详细介绍 ComfyUI 的节点系统架构,涵盖节点定义、执行流程、缓存机制、扩展性和系统特性等方面。

2.1 节点系统的基础架构

ComfyUI 的节点系统基于 Python 模块化设计,所有节点及其行为都通过类的形式进行定义。这些节点在启动时会进行注册,允许系统灵活地加载和使用内置节点与自定义节点。

核心节点定义与注册

ComfyUI 的节点系统在 nodes.py 中定义,并通过以下映射存储所有节点类及其显示名称:

NODE_CLASS_MAPPINGS = {}  # 存储所有节点类的映射
NODE_DISPLAY_NAME_MAPPINGS = {}  # 节点显示名称映射

节点通过类定义,并包含以下几个关键属性:

  • INPUT_TYPES: 输入参数的类型定义。
  • RETURN_TYPES: 返回数据的类型定义。
  • FUNCTION: 节点的具体执行函数。
  • CATEGORY: 节点的类别,用于在 UI 中分类显示。

节点类型

ComfyUI 支持两种类型的节点:

  • 内置节点: 在系统启动时加载,存储在 comfy_extras 目录下。内置节点提供了常见的图像操作、模型加载和处理功能。
  • 自定义节点: 用户可以在 custom_nodes 目录中添加自定义节点,系统在启动时自动加载这些节点。

节点加载机制

ComfyUI 提供了灵活的节点加载机制,允许动态加载内置节点和自定义节点:

def init_builtin_extra_nodes():
    extras_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras")
    extras_files = ["nodes_latent.py""nodes_hypernetwork.py"]

对于自定义节点,ComfyUI 使用动态模块导入的方式加载:

def load_custom_node(module_path: str, ignore=set(), module_parent="custom_nodes") -> bool:
    module_spec = importlib.util.spec_from_file_location(module_name, module_path)
    module = importlib.util.module_from_spec(module_spec)
    if hasattr(module, "NODE_CLASS_MAPPINGS"):
        for name, node_cls in module.NODE_CLASS_MAPPINGS.items():
            NODE_CLASS_MAPPINGS[name] = node_cls

这种设计使得 ComfyUI 可以方便地扩展和加载新节点,用户可以根据需求自定义节点功能并动态加载。

2.2 节点执行系统

节点的执行逻辑由 execution.py 中的 PromptExecutor 类负责。该类管理了节点的执行流程,包括输入验证、节点函数执行、结果缓存和输出返回等。

执行流程

节点的执行流程如下:

  1. 输入验证: 系统首先验证节点的输入是否符合预定义的输入类型。
  2. 获取输入数据: 从上一个节点或缓存中获取节点的输入数据。
  3. 执行节点函数: 根据定义的 FUNCTION 执行节点逻辑。
  4. 缓存结果: 执行结果会缓存,避免重复计算。
  5. 返回输出: 将执行结果返回给下游节点或 UI。

执行器的核心代码如下:

class PromptExecutor:
    def execute(self, prompt, prompt_id, extra_data=None, execute_outputs=[]):
        # 节点执行的主要逻辑

ComfyUI 通过此执行器确保节点按顺序执行,并管理节点间的数据流动。

2.3 缓存机制

为了提高性能,减少重复计算,ComfyUI 实现了多层缓存系统,缓存节点的输出、对象和 UI 相关数据。缓存的实现类为 HierarchicalCache,并且系统支持 LRU(最近最少使用) 缓存策略。后续章节会详细讲一下缓存,这里先略过。

2.4 数据流与依赖管理

ComfyUI 的节点系统依赖于图形化的数据流管理。节点之间通过输入和输出相互连接,系统会自动分析节点间的依赖关系,确保数据流在节点间正确传递。

节点图解析与执行顺序

  1. 节点图解析: ComfyUI 解析 JSON 格式的节点图,识别节点之间的连接关系。
  2. 依赖管理: 系统自动分析节点间的依赖,确保每个节点在其依赖的节点完成后执行。
  3. 执行顺序构建: 系统基于依赖关系构建节点的执行顺序,防止循环依赖和执行死锁的发生。

2.5 错误处理与资源管理

ComfyUI 实现了完善的错误处理机制,确保节点执行过程中出现错误时,系统能够及时捕获并反馈给用户,避免系统崩溃。同时,ComfyUI 通过自动垃圾回收和内存管理功能,定期清理未使用的模型和缓存,优化系统资源使用。

常见的错误处理机制包括:

  • 输入验证失败: 如果输入数据类型不匹配或数据缺失,系统会抛出异常。
  • 执行错误: 如果节点在执行过程中遇到错误,系统会捕获并将错误信息反馈到前端 UI。

垃圾回收机制由以下代码管理:

gc_collect_interval = 10.0  # 每10秒进行一次垃圾回收

2.6 前端接口与用户交互

ComfyUI 提供了 WebSocket 和 REST API 来管理前端与后端的通信。通过这些接口,前端 UI 可以实时监控节点的执行状态,并获取节点的执行结果。

WebSocket 通信

WebSocket 负责处理前端与后端的实时通信,包括节点执行状态的推送和执行命令的接收。

@routes.get('/ws')
async def websocket_handler(request):
    # 处理前端连接
    # 发送执行状态
    # 处理节点执行

REST API

REST API 用于获取节点信息和处理其他非实时请求。例如,可以通过以下 API 获取节点的详细信息:

@routes.get("/object_info/{node_class}")
async def get_object_info_node(request):
    # 获取节点信息

2.7 系统扩展与自定义

ComfyUI 的节点系统设计非常灵活,支持用户根据需求扩展新功能。用户可以通过编写自定义节点来扩展系统的功能,而无需修改核心代码。系统支持热插拔的节点加载机制,使得自定义节点的开发和调试更加便捷。

动态类型系统

ComfyUI 支持动态类型系统,可以根据运行时情况自动调整节点的输入输出类型,确保系统的灵活性和可扩展性。

插件与自定义节点

自定义节点通过 custom_nodes 目录加载,用户可以编写自己的节点并通过 NODE_CLASS_MAPPINGS 注册到系统中。

3. 缓存系统

在 ComfyUI 项目中,缓存系统的设计主要由以下几个部分组成:

3.1 缓存类型

ComfyUI 的缓存系统主要包括三种缓存类型:

  • outputs: 用于缓存节点的输出结果,减少不必要的重复计算。
  • ui: 缓存与用户界面相关的数据,比如预览图像、状态信息等。
  • objects: 专门用于存储大型对象,如模型等大体积数据。

这三种缓存通过 CacheSet 类进行初始化和管理,具体实现如下:

class CacheSet:
    def __init__(self, lru_size=None):
        if lru_size is None or lru_size == 0:
            self.init_classic_cache() 
        else:
            self.init_lru_cache(lru_size)
        self.all = [self.outputs, self.ui, self.objects]

3.2 缓存实现方式

缓存系统可以根据不同的需求选择不同的实现方式:

Classic Cache(传统缓存)

当没有设置 LRU 缓存大小时,缓存系统会初始化为经典的层级缓存。具体实现如下:

def init_classic_cache(self):
    self.outputs = HierarchicalCache(CacheKeySetInputSignature)
    self.ui = HierarchicalCache(CacheKeySetInputSignature)
    self.objects = HierarchicalCache(CacheKeySetID)
  • outputs 和 ui 都使用 CacheKeySetInputSignature 作为缓存键,用于基于输入签名进行缓存。
  • objects 使用的是 CacheKeySetID 作为缓存键,主要是为了存储大型对象,如模型数据。

LRU Cache(最近最少使用缓存)

如果设置了 LRU 缓存大小,系统会初始化为 LRU 缓存,适用于内存较充足的情况下,以优化性能。

def init_lru_cache(self, cache_size):
    self.outputs = LRUCache(CacheKeySetInputSignature, max_size=cache_size)
    self.ui = LRUCache(CacheKeySetInputSignature, max_size=cache_size)
    self.objects = HierarchicalCache(CacheKeySetID)
  • outputs 和 ui 使用 LRU 缓存,能够根据使用频率自动淘汰最少使用的缓存项,保证内存得到最优管理。
  • objects 仍然使用 HierarchicalCache,因为模型等大型对象的加载和卸载代价较高,不适合频繁淘汰。

3.3 缓存键策略

为了确保缓存的正确性,缓存系统使用两种不同的缓存键策略:

CacheKeySetInputSignature

  • 用于 outputs 和 ui 缓存。
  • 该缓存键基于节点的输入签名,包含节点类型、输入参数及祖先节点关系等,可以确保相同的输入得到相同的输出。

CacheKeySetID

  • 用于 objects 缓存。
  • 这是一种基于节点 ID 和类型的简单缓存键,用于存储与输入无关的大型对象,如模型等。

3.4 缓存清理机制

无论是传统缓存还是 LRU 缓存,ComfyUI 都提供了缓存清理机制,避免过多的缓存占用内存资源。

经典缓存清理

经典缓存会及时释放不再需要的数据,防止内存溢出。

LRU 缓存清理

LRU 缓存通过一个 generation 计数器和 used_generation 字典记录每个缓存项的使用时间。当缓存超出预设大小时,LRU 算法会移除最久未使用的项。

3.5 扩展性和调试

ComfyUI 的缓存系统支持扩展和调试:

  • 扩展性: 缓存系统支持自定义节点和模型的缓存策略,可以根据需要调整缓存大小和键生成逻辑。
  • 调试: 提供了 recursive_debug_dump 函数,能够以递归方式输出缓存的调试信息,方便开发者进行问题排查。
def recursive_debug_dump(self):
    result = {
        "outputs": self.outputs.recursive_debug_dump(),
        "ui": self.ui.recursive_debug_dump(),
    }
    return result

3.6 缓存小结

ComfyUI 的缓存系统设计非常灵活,它通过 CacheSet 类将不同类型的数据(节点输出、UI 数据、大型对象)分离管理,并支持经典缓存和 LRU 缓存两种策略。依靠层次化的缓存结构和精确的缓存键策略,ComfyUI 能够有效地减少重复计算,并优化内存使用。

4. 使用限制

在使用 ComfyUI 时,虽然系统本身没有硬性限制节点数量,但仍然存在一些与性能、资源管理和系统稳定性相关的限制。这些限制大多与历史记录、图像分辨率、缓存、内存管理等方面有关。

4.1 历史记录限制

ComfyUI 会保存工作流中的历史记录,用于回溯和调试。系统中对历史记录的最大保存数量有如下限制:

MAXIMUM_HISTORY_SIZE = 10000
  • 限制描述: 系统最多保存 10000 条历史记录。
  • 影响: 当历史记录达到上限时,旧的记录会被移除,无法继续回溯到更早的操作。
  • 建议: 如果不需要长时间保存历史记录,定期清理历史记录可以释放内存资源,提升系统的运行效率。

4.2 图像分辨率限制

ComfyUI 中对单个图像处理的最大分辨率有限制,防止超大图像导致系统崩溃或内存不足:

MAX_RESOLUTION = 16384
  • 限制描述: 系统允许的最大图像分辨率为 16384×16384 像素。
  • 影响: 超过此分辨率的图像无法处理,可能会导致内存溢出或显存不足的情况。
  • 建议: 尽量避免处理过于高分辨率的图像,尤其是在显存较小的 GPU 上运行时。对于需要处理大图像的工作流,可以考虑将图像分割成多个较小的瓦片。

4.3 VAE 解码瓦片大小限制

在图像生成过程中,VAE 解码器对瓦片大小进行了限制,以确保解码过程的效率与内存管理:

"tile_size": ("INT", {"default"512"min"128"max"4096"step"32})
  • 限制描述: VAE 解码时,允许的瓦片大小在 128 到 4096 像素之间。
  • 影响: 如果瓦片大小设置过大,系统可能会因为内存不足而运行缓慢或崩溃;瓦片大小过小则可能影响解码效率。
  • 建议: 根据 GPU 显存大小合理调整瓦片大小,找到性能和内存占用之间的平衡点。

4.4 文件上传大小限制

在使用 ComfyUI 时,文件上传的大小受限于系统的配置,特别是通过命令行参数 max_upload_size 来控制:

max_upload_size = round(args.max_upload_size * 1024 * 1024)
  • 限制描述: 上传文件的最大尺寸由命令行参数 max_upload_size 控制,单位为 MB。
  • 影响: 超过上传文件大小限制的文件将无法上传或处理。
  • 建议: 如果需要上传较大的文件,可以通过调整命令行参数来提高上传文件的大小限制。

4.5 缓存限制

ComfyUI 使用缓存系统来优化计算效率,减少重复计算。缓存的大小和管理方式可以通过 LRU(最近最少使用)策略进行控制:

def __init__(self, lru_size=None):
    if lru_size is None or lru_size == 0:
        self.init_classic_cache() 
    else:
        self.init_lru_cache(lru_size)
  • 限制描述: 缓存的大小受 LRU 策略控制,缓存过多时,系统会淘汰最少使用的缓存项。
  • 影响: 当缓存大小设置过小,可能会频繁清除缓存,导致重复计算;缓存过大则可能占用大量内存。
  • 建议: 根据内存资源合理设置缓存大小。对于内存充足的系统,可以增加缓存大小,以减少重复计算;对于内存紧张的系统,建议定期清理缓存。

4.6 执行队列限制

节点的执行通过队列进行管理,系统按顺序执行节点,避免同时执行过多节点造成性能瓶颈:

  • 限制描述: 系统使用队列调度节点执行,包括当前运行的队列和等待执行的队列。
  • 影响: 如果节点过多,执行速度会受到队列调度的影响,可能出现等待时间过长的情况。
  • 建议: 尽量简化工作流,避免过多的节点同时排队执行。如果遇到性能瓶颈,可以将大规模的工作流分解为多个子工作流逐步执行。

4.7 Tokenizer 限制

在文本处理方面,CLIP 模型的 Tokenizer 有一个最大长度限制:

"model_max_length"77
  • 限制描述: CLIP 模型的 Tokenizer 最多支持 77 个 token。可修改配置。
  • 影响: 超过 77 个 token 的输入文本将被截断,可能会影响文本生成的精度。
  • 建议: 尽量保持输入文本简洁,避免过长的描述。如果必须使用长文本,可以通过分段输入的方式进行处理。

4.8 限制小结

虽然 ComfyUI 对节点数量没有明确的硬性限制,但在使用过程中仍然受到一些系统资源和配置的限制。这些限制大多是为了确保系统的稳定性、优化性能以及合理使用内存资源。为了避免因这些限制导致的性能瓶颈或崩溃,建议在使用时遵循以下最佳实践:

  • 保持工作流简洁: 避免不必要的节点,定期清理历史记录和缓存。
  • 监控系统资源: 注意监控内存和显存的使用情况,避免资源耗尽。
  • 分解大型工作流: 对于复杂的工作流,可以分成多个子工作流来逐步执行,减少单次执行的节点数量。
  • 调整系统配置: 根据实际需求调整文件上传大小、缓存设置等参数。

通过合理地管理工作流和系统资源,ComfyUI 可以在大型工作流中保持高效运行,避免因资源限制导致的性能问题。

模块化设计带来的无限可能

ComfyUI 的模块化节点系统不仅提升了用户的易用性,还通过灵活的扩展机制和高效的缓存管理提供了强大的自定义能力。

它的图形化工作流设计大大降低了 AI 绘图的技术门槛,使得更多用户能够轻松上手,并根据自己的需求定制不同的图像生成方案。

随着社区的不断壮大和功能的持续扩展,ComfyUI 有望成为 AI 绘图领域的重要基础设施之一,为创作与开发者提供无限的可能性。

以上。