ComfyUI 作为基于节点图的 Stable Diffusion 推理界面，在生成高分辨率图像或切换大模型时，经常因为显存占用突增而触发 CUDA out of memory。要稳定使用，就必须清楚每一步节点到底消耗了多少 VRAM，并且在接近上限前主动干预。本文从显存数据获取、前端展示和防护策略三个层面，说明如何给 ComfyUI 增加实用的监控与预防能力。
一、从 PyTorch 接口读取准确的 VRAM 占用数据
很多用户习惯打开系统任务管理器看显存，但那是驱动层统计，包含其他进程开销，且无法区分 ComfyUI 内部张量缓存与碎片。ComfyUI 后端运行在 PyTorch 之上，可以直接调用 torch.cuda.memory_allocated 与 torch.cuda.memory_reserved 拿到当前设备已分配和已预留的字节数。前者是模型权重与中间特征真实占用的显存，后者是 PyTorch 缓存分配器向驱动申请保留的池子，通常比前者大。
我们可以在自定义节点的 execute 方法里，于模型加载前、推理前、推理后分别采样。下面这段代码演示了如何封装一个轻量采集函数，并把结果换算为兆字节返回。注意 device 默认取 cuda:0，多卡环境需自行传入。
import torch
def sample_vram(device="cuda:0"):
if not torch.cuda.is_available():
return None
allocated = torch.cuda.memory_allocated(device) / (1024 * 1024)
reserved = torch.cuda.memory_reserved(device) / (1024 * 1024)
total = torch.cuda.get_device_properties(device).total_memory / (1024 * 1024)
return {
"allocated_mb": round(allocated, 1),
"reserved_mb": round(reserved, 1),
"total_mb": round(total, 1),
"util_percent": round(allocated / total * 100, 1)
}
# 在节点执行中调用
stats = sample_vram()
print(stats)
这种采集方式开销极低，每次只是读取计数器，不会触发设备同步。但如果节点内部频繁调用，建议加一个时间间隔判断，例如每五百毫秒才上报一次，避免日志爆炸。另外，PyTorch 的缓存分配器不会立即把释放的张量还给驱动，所以看到 reserved 居高不下并不一定代表泄漏，可配合 torch.cuda.empty_cache 在空闲节点主动回收。
二、把显存数据实时推送到 ComfyUI 前端界面
ComfyUI 前端通过 WebSocket 与后端通信，我们可以利用已有的通道，在节点执行时向后发送自定义事件。后端在 server.py 的路由或节点回调中调用 prompt_server 的 send_sync 方法，前端监听对应类型消息后更新进度条或文本。相比轮询接口，事件驱动更省资源。
下面示例展示后端如何发送显存消息。这里借用 ComfyUI 内部的 server.PromptServer.instance 单例，消息类型命名为 vram_update，载荷就是前面 sample_vram 的字典。
from server import PromptServer
def push_vram_event(device="cuda:0"):
data = sample_vram(device)
if data is None:
return
PromptServer.instance.send_sync("vram_update", data)
前端扩展通常放在 web/extensions 目录，用 JavaScript 监听并操作 DOM。下面是一段简化逻辑，把利用率写进页面顶部状态栏。实际项目里可以画一个小折线图，记录最近三十秒的趋势，方便判断是不是线性上涨。
app.extensionManager.registerApi("vram_monitor", {
setup() {
const status = document.createElement("div");
status.style.cssText = "position:fixed;top:0;right:0;z-index:999;";
document.body.appendChild(status);
const ws = app.socket;
ws.addEventListener("message", (e) => {
const msg = JSON.parse(e.data);
if (msg.type === "vram_update") {
status.textContent = `VRAM ${msg.data.util_percent}% (${msg.data.allocated_mb}MB)`;
}
});
}
});
这种方案的优点是无需改造核心代码，以扩展形式存在，升级 ComfyUI 时冲突小。但要注意，前端扩展若报错可能让整个界面白屏，因此推送数据要做容错，例如后端字段缺失时前端跳过渲染。同时，显存数值建议保留两位小数即可，过长的小数对肉眼监控没有意义。
三、基于阈值的 OOM 预防与自动降级策略
仅仅看到显存涨上去还不够，真正有价值的是在溢出前拦住请求。我们可以设定两个水位线：警告线例如百分之八十，危险线例如百分之九十。当 util_percent 跨过警告线，后端自动在日志提示并推迟非关键节点；跨过危险线，则中断当前队列，释放缓存，并向用户返回友好错误而不是 CUDA 异常。
自动降级的具体动作包括：把采样步数减半、启用 torch.cuda.amp 的更激进精度、调用 model.to("cpu") 把暂时不用的参考模型移出显存、或者开启注意力切片。下面代码演示一个简易守卫函数，在节点入口判断，如果危险就抛出一个可捕获的业务异常。
class VRAMGuardError(Exception):
pass
def vram_guard(device="cuda:0", warn=80.0, danger=90.0):
data = sample_vram(device)
if data is None:
return
if data["util_percent"] >= danger:
torch.cuda.empty_cache()
raise VRAMGuardError("显存超过危险阈值，已终止生成")
elif data["util_percent"] >= warn:
print("警告：显存偏高，建议降低分辨率")
# 在节点 execute 开头调用
vram_guard()除了被动守卫，还可以在工作流设计阶段预防。比如把大模型加载节点和采样节点放在同一执行分支，避免多个大模型同时驻留；使用 ComfyUI 自带的 VRAM 管理选项，设置为 highvram 或 lowvram 模式。低显存模式下，ComfyUI 会按需把层搬进搬出，虽然慢一点，但能有效防止 OOM。结合前面的实时监控，你就能在界面上直观看到策略生效后的曲线回落。
最后提醒，预防 OOM 不是单点工作，而是采集、展示、干预的闭环。显存监控让你知道发生了什么，阈值守卫让你在崩溃前行动，而合理的节点编排才是根本解法。哪怕只有八显存，只要这三层配合，也能在 ComfyUI 里稳定跑通大多数文生图任务。
ComfyUIVRAM_monitorOOM_prevention修改时间：2026-08-15 03:24:35