Comfy-Org / Comfy-Org/ComfyUI

Implement a callback for memory cleanup by custom nodes that allocate their own memory

Open
#11,283 0 comments 4 reactions 0 assignees View on GitHub
Feature
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Feature Idea

My custom node https://github.com/phaserblast/ComfyUI-DGXSparkSafetensorsLoader uses fastsafetensors to allocate memory and load safetensors models directly into VRAM using GPUDirect, bypassing ComfyUI's memory management entirely. This was needed to get around mmap() used by Hugging Face's safetensors loader, which doesn't work well on shared memory systems like DGX Spark.

The problem with this approach is although models load in only a few seconds (as opposed to over a minute or more using the HF stuff), the allocated memory can't be freed. The reason is ComfyUI's memory management has zero clue about what the node is doing memory-wise, and cannot call the necessary fastsafetensors methods to free the memory.

The solution is to have ComfyUI accept a callback that contains the code necessary to free the memory allocated by the custom implementation. It can check to see if the callback is non-null, and if so use that when it wants to unload a model rather than its own built-in stuff. This will offer custom nodes way more flexibility when loading models.

### Existing Solutions

_No response_

### Other

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.