97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
"""Force memory cleanup nodes for ComfyUI.
|
|
|
|
This extension deliberately asks ComfyUI's prompt worker to clean up *after*
|
|
the current prompt has finished. Cleaning synchronously inside a node is too
|
|
early because the executor still owns references to the current graph.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from server import PromptServer
|
|
|
|
|
|
LOGGER = logging.getLogger("ComfyUI-Force-Cleanup")
|
|
|
|
|
|
class ForceCleanupAfterPrompt:
|
|
"""Request a model/cache cleanup after the current prompt completes."""
|
|
|
|
CATEGORY = "utils/memory"
|
|
FUNCTION = "request_cleanup"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
DESCRIPTION = (
|
|
"Requests cleanup after the entire prompt finishes. Full cleanup unloads "
|
|
"models, clears ComfyUI's execution cache, runs Python GC, and releases "
|
|
"the PyTorch device allocator cache."
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"cleanup_level": (
|
|
[
|
|
"full: models + execution cache + allocator",
|
|
"models: unload models + allocator",
|
|
],
|
|
{
|
|
"default": "full: models + execution cache + allocator",
|
|
"tooltip": (
|
|
"Full frees the most memory but the next run must "
|
|
"reload/recompute everything."
|
|
),
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **_: Any) -> float:
|
|
# Output nodes can otherwise be satisfied from the execution cache.
|
|
# NaN never compares equal, so the cleanup request is made on every run.
|
|
return float("nan")
|
|
|
|
def request_cleanup(self, cleanup_level: str):
|
|
prompt_server = getattr(PromptServer, "instance", None)
|
|
prompt_queue = getattr(prompt_server, "prompt_queue", None)
|
|
if prompt_queue is None:
|
|
raise RuntimeError(
|
|
"ComfyUI prompt queue is unavailable; cleanup was not scheduled."
|
|
)
|
|
|
|
if cleanup_level.startswith("full:"):
|
|
# ComfyUI's worker handles this after execution:
|
|
# unload_all_models() -> executor.reset() -> gc.collect()
|
|
# -> soft_empty_cache().
|
|
prompt_queue.set_flag("free_memory", True)
|
|
message = (
|
|
"Full cleanup scheduled after this prompt: models, execution "
|
|
"cache, Python GC, and device allocator cache."
|
|
)
|
|
else:
|
|
# The worker unloads models, runs GC, and empties the device allocator,
|
|
# while retaining the executor's reusable graph cache.
|
|
prompt_queue.set_flag("unload_models", True)
|
|
message = (
|
|
"Model cleanup scheduled after this prompt; reusable execution "
|
|
"cache will be kept."
|
|
)
|
|
|
|
LOGGER.info(message)
|
|
return {"ui": {"text": [message]}, "result": ()}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ForceCleanupAfterPrompt": ForceCleanupAfterPrompt,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ForceCleanupAfterPrompt": "Force Cleanup After Prompt",
|
|
}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|