Ubuntu下ComfyUI显存占用过高如何优化？

Question

在Ubuntu系统下运行ComfyUI时，常出现显存占用过高的问题，尤其在加载大型Stable Diffusion模型或多节点工作流时，GPU显存易接近满载，导致页面响应卡顿甚至崩溃。常见表现为`nvidia-smi`显示显存使用率超90%，推理过程中出现CUDA out of memory错误。该问题可能源于默认未启用显存优化选项、模型加载精度过高或缓存管理不当。如何在不影响生成质量的前提下，通过配置参数或修改执行环境有效降低ComfyUI的显存占用？

Answer 1

一、问题背景与现象分析

在Ubuntu系统下运行ComfyUI时，显存占用过高是常见瓶颈。尤其在加载大型Stable Diffusion模型（如SDXL、SD 1.5大参数版本）或多节点复杂工作流（如ControlNet + LoRA + Upscaler级联）时，GPU显存极易达到90%以上。

通过 nvidia-smi 可观察到显存使用迅速飙升，典型表现为：

• CUDA out of memory 错误频繁触发
• Web UI响应延迟或完全卡死
• 生成任务中途崩溃，日志提示显存分配失败

根本原因通常包括：默认FP32精度加载模型、未启用显存优化策略、缓存机制冗余、以及多节点并行执行导致中间张量堆积。

二、显存占用来源深度剖析

来源	说明	典型占用比例
模型权重	Stable Diffusion UNet、VAE、CLIP等组件以FP32载入	~40%
中间激活张量	前向传播过程中产生的特征图缓存	~35%
优化器状态	训练模式下保留梯度与动量	N/A（推理中可关闭）
节点缓存	ComfyUI默认缓存部分节点输出以防重复计算	~15%
PyTorch内部开销	内存池、CUDA上下文管理等底层开销	~10%

三、基础优化配置方案

1. 启用 FP16混合精度 模式，在启动脚本中添加参数：

2. python main.py --use_fp16

3. 强制启用 Torch的内存节省模式：

4. torch.backends.cuda.matmul.allow_tf32 = True
   torch.set_float32_matmul_precision('medium')

5. 设置最大批处理尺寸限制，避免一次性加载过多图像：

6. "max_batch_size": 4

7. 禁用不必要的VAE解码缓存：

8. # 在custom nodes中插入
   node_vae.decode = torch.no_grad()(node_vae.decode)

四、高级显存管理技术

对于具备深度调优能力的开发者，可采用以下进阶手段：

1. 梯度检查点（Gradient Checkpointing）：牺牲少量计算时间换取显存压缩，适用于UNet主干：

2. from torch.utils.checkpoint import checkpoint
   model.enable_gradient_checkpointing()

3. 模型分片加载（Model Offloading）：将不活跃模块临时移至CPU：

4. from accelerate import cpu_offload
   cpu_offload(model.unet, device)

5. 动态张量释放：在ComfyUI后端注册清理钩子：

6. def clear_cache_hook(module, input, output):
       torch.cuda.empty_cache()
   layer.register_forward_hook(clear_cache_hook)

五、执行环境与运行时调优

结合Ubuntu系统特性，调整运行时环境参数：

# 设置CUDA可见设备（多卡场景）
export CUDA_VISIBLE_DEVICES=0

# 启用PyTorch CUDA内存碎片整理
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

# 使用NVIDIA Compute Mode提升调度效率
nvidia-smi -c 1

推荐使用如下启动命令整合所有优化项：

python main.py \
  --disable-xformers \
  --gpu-only \
  --highvram \
  --dont-upcast-attention \
  --fp16

六、可视化流程与决策路径

graph TD A[开始] --> B{是否使用大型模型?} B -- 是 --> C[启用FP16] B -- 否 --> D[保持FP32] C --> E{是否存在OOM?} E -- 是 --> F[开启梯度检查点] E -- 否 --> G[正常运行] F --> H{仍超限?} H -- 是 --> I[启用CPU卸载] H -- 否 --> G I --> J[监控性能下降幅度] J --> K[平衡质量与资源消耗]