ComfyUI-CustomModelPatcher

这个仓库并不提供具体功能，这个仓库展示了如何给自定义外部模型或自定义 Tensor 在 ComfyUI 的显存管理器中预留显存。

例如 onnxruntime-gpu、InsightFace 等模型，他们的显存不受 ComfyUI 的显存管理器管理，所以可能出现他们想申请显存时，ComfyUI 此时不自动提前释放显存，然后显存分配失败。

使用 custom_model_patcher.py 你可以在初始化这种外部模型之前，提前让 ComfyUI 释放一部分显存，并且占据一定空间。这部分被占据的空间会持续到当前变量被 gc。

注意：这里有两个概念，ModelPatcher 和 Model。CustomModel 只有两种状态，加载状态(占用显存)和卸载状态(不占用显存)。CustomModelPatcher 只是一个引用计数器。（ComfyUI 内置的 Model 和 ModelPatcher 支持加载卸载单个层。）ModelPatcher 可以被 .clone()。当同一个底层模型的 ModelPatcher 的最后一个实例被 gc 时，它占用的显存将被回收。ComfyUI 的每个节点运行结果按 id 有缓存，当显存不足时，优先删除这些节点运行结果缓存，然后触发 gc，从而释放显存。

使用方法

下载 custom_model_patcher.py

创建 ModelPatcher。指定预留大小和加载模型的函数。

from .custom_model_patcher import CustomModelPatcher, SimpleCustomModel

def model_loader():
    model = torch.zeros(size_gb * 256, 1000, 1000, device="cuda")
    return model

return (CustomModelPatcher(SimpleCustomModel(model_loader, size=size_gb * 1024 * 1024 * 1024)),)

使用 model。必须先 clone，再 load_models_gpu，然后才能使用。

model = model.clone()  # this model variable is CustomModelPatcher
comfy.model_management.load_models_gpu([model])
simple_custom_model = model.model  # this is the actual SimpleCustomModel instance
model_tensor = simple_custom_model.model  # this is the actual model returned by the model_loader function
return (str(model_tensor.shape),)

License

MIT License