ComfyUI Extension: ComfyUI-mem-safe-wrapper

Authored by MiddleKD

Created

Updated

3 stars

ComfyUI's Smart Memory Management efficiently manages RAM, GPU memory, and garbage collection. This feature keeps frequently used models in memory to increase inference speed, and dynamically releases less important models when memory is low to optimize resources. However, not all ComfyUI custom node developers create nodes that are compatible with Smart memory management. This includes several impressive models. Mem-safe-wrapper is a custom node that wraps these model nodes to enable ComfyUI's Smart memory management capabilities.

Custom Nodes (0)

    README

    🎨 ComfyUI-mem-safe-wrapper

    git_header Load된 모델을 ComfyUI에서 Smart memory management를 가능하도록 바꾸는 노드입니다.

    📌 Index

    🚀 Introduction

    ComfyUI의 Smart Memory Management는 RAM, GPU 메모리 및 가비지 컬렉션을 효율적으로 관리합니다. 이 기능은 자주 사용되는 모델을 메모리에 유지하여 추론 속도를 높이고, 메모리가 부족할 경우 덜 중요한 모델을 동적으로 해제하여 자원을 최적화합니다.

    하지만 모든 ComfyUI 커스텀 노드 개발자가 Smart memory management가 가능하도록 노드를 개발하지는 않습니다. 그 중에서 멋진 모델들도 다수 포함됩니다.

    Mem-safe-wrapper는 이런 모델 노드를 wrapping하여 ComfyUI가 Smart memory management를 가능하도록 하는 커스텀 노드입니다.

    💡 Features

    Memory safe wrap

    torch.nn.Module을 상속받은 모델을 input으로 받아 comfyui의 ModelPatcher로 wrapping합니다. 기존 model의 attribute들은 그대로 사용할 수 있습니다.

    Reset model patcher

    다른 node에서 Modelpatcher의 calculate weight 함수를 injection했다면 오류가 발생하는 custom node가 다수 존재합니다.(ex: comfyui-inpaint-nodes <a href="https://github.com/Acly/comfyui-inpaint-nodes">https://github.com/Acly/comfyui-inpaint-nodes</a>). 이를 원본 Modelpatcher의 calculate weight로 reset하여 문제를 해결하는 노드입니다.

    📥 Install

    cd {your_comfyui_path}/custom_nodes
    git clone {this_repository}
    

    🖥 How to use

    해당 코드를 main.prompt_worker에 추가하고 실행한 결과입니다. (사용시 코드를 추가할 필요는 없습니다. 단순 테스트 용도입니다.)

    from comfy.model_management import current_loaded_models
    print(f"""
    ----------
    Current loaded model in memory:
    {current_loaded_models}
    ----------
    """)
    

    Without mem-safe-wrapper

    project_header project_header

    With mem-safe-wrapper

    project_header project_header

    📞 Contact

    middlek - [email protected]