147 lines
5.2 KiB
Python
147 lines
5.2 KiB
Python
"""
|
|
VRAMFree Node for ComfyUI
|
|
MIT License - Independent implementation without GPL code reuse
|
|
"""
|
|
|
|
import logging
|
|
|
|
ANY_SOCKET_TYPES = "IMAGE,LATENT,MODEL,CONDITIONING,CLIP,VAE,MASK,INT,FLOAT,STRING,BOOLEAN"
|
|
|
|
try:
|
|
from .memory_utils import MemoryManager
|
|
except ImportError:
|
|
from memory_utils import MemoryManager
|
|
|
|
|
|
class VRAMFreeNode:
|
|
"""ComfyUI custom node that releases memory by unloading models and clearing caches."""
|
|
|
|
def __init__(self):
|
|
self.memory_manager = MemoryManager()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define node input configuration."""
|
|
return {
|
|
"required": {},
|
|
"optional": {
|
|
"Any": (ANY_SOCKET_TYPES,),
|
|
"MinFreeVRAMGB": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"step": 0.1,
|
|
"tooltip": "Skip cleanup when free VRAM is already above this threshold (GB)",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (ANY_SOCKET_TYPES,)
|
|
RETURN_NAMES = ("Bypass(Relay)",)
|
|
|
|
FUNCTION = "execute_cleanup"
|
|
CATEGORY = "u5"
|
|
OUTPUT_NODE = True
|
|
DESCRIPTION = "Unload models and free memory without disrupting workflow"
|
|
|
|
def execute_cleanup(self, Any=None, MinFreeVRAMGB=0.0, **legacy_kwargs):
|
|
"""Run the memory cleanup routine and return status details."""
|
|
legacy_detail_log = legacy_kwargs.pop("detail_log", None)
|
|
legacy_verbose = legacy_kwargs.pop("verbose", None)
|
|
|
|
# Legacy compatibility for removed flags (always True now)
|
|
legacy_kwargs.pop("unload_models", None)
|
|
legacy_kwargs.pop("clear_cache", None)
|
|
legacy_kwargs.pop("force_gc", None)
|
|
legacy_min_vram = legacy_kwargs.pop("min_free_vram_gb", None)
|
|
if legacy_min_vram is not None:
|
|
MinFreeVRAMGB = legacy_min_vram
|
|
|
|
detail_debug_log = bool(legacy_detail_log or legacy_verbose)
|
|
|
|
try:
|
|
initial_status = self.memory_manager.get_memory_status()
|
|
initial_summary = self.memory_manager.get_memory_summary()
|
|
|
|
if detail_debug_log:
|
|
print("\n=== VRAMFree: Starting Cleanup ===")
|
|
print(f"Initial Memory Status:\n{initial_summary}")
|
|
print("Options: unload_models=True, clear_cache=True, force_gc=True")
|
|
|
|
result = self.memory_manager.cleanup_memory(
|
|
unload_models=True,
|
|
clear_cache=True,
|
|
force_gc=True,
|
|
min_free_vram_gb=MinFreeVRAMGB
|
|
)
|
|
|
|
final_status = result.get("final_status") or {}
|
|
freed_vram = float(result.get("freed_vram_gb", 0.0))
|
|
freed_ram = float(result.get("freed_ram_gb", 0.0))
|
|
total_freed = max(0.0, freed_vram + freed_ram)
|
|
|
|
status_parts = ["Models", "Cache", "GC"]
|
|
actions_label = ", ".join(status_parts)
|
|
freed_str = self.memory_manager.format_memory_size(total_freed)
|
|
status_line = f"Cleared: {actions_label} | Freed: {freed_str}" if status_parts else "No actions performed"
|
|
|
|
final_summary = self.memory_manager.get_memory_summary(result.get("final_status"))
|
|
|
|
report_lines = [
|
|
"=== Memory Cleanup Report ===",
|
|
"",
|
|
f"Status: {'SUCCESS' if result.get('success') else 'FAILED'}",
|
|
f"Actions: {status_line}",
|
|
"",
|
|
"--- Memory Changes ---",
|
|
f"VRAM Freed: {self.memory_manager.format_memory_size(freed_vram)}",
|
|
f"RAM Freed: {self.memory_manager.format_memory_size(freed_ram)}",
|
|
f"Total Freed: {self.memory_manager.format_memory_size(total_freed)}",
|
|
"",
|
|
"--- Current Memory Status ---",
|
|
final_summary,
|
|
]
|
|
|
|
if detail_debug_log:
|
|
report_lines.append("")
|
|
report_lines.append("--- Detailed Actions ---")
|
|
for action in result.get("actions_performed", []):
|
|
report_lines.append(f" - {action}")
|
|
|
|
if not result.get("success"):
|
|
report_lines.append("")
|
|
report_lines.append(f"Error: {result.get('message', 'Unknown error')}")
|
|
|
|
report = "\n".join(report_lines)
|
|
|
|
logging.info(f"VRAMFree: {status_line}")
|
|
|
|
if detail_debug_log:
|
|
print(report)
|
|
|
|
return (Any,)
|
|
|
|
except Exception as exc: # noqa: BLE001
|
|
status_line = f"Exception: {exc}"
|
|
report_lines = [
|
|
"=== Memory Cleanup Report ===",
|
|
"",
|
|
"Status: EXCEPTION",
|
|
status_line,
|
|
"",
|
|
"Please check the console for detailed error information.",
|
|
]
|
|
report = "\n".join(report_lines)
|
|
|
|
logging.error(f"VRAMFree: {status_line}")
|
|
|
|
if detail_debug_log:
|
|
print(f"\nEXCEPTION: {exc}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
return (Any,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **_kwargs):
|
|
"""Force reevaluation so the node executes every time."""
|
|
return float("NaN") |