77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
import torch
|
|
import gc
|
|
import psutil
|
|
import os
|
|
import ctypes
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any_typ = AnyType("*")
|
|
|
|
|
|
class DeleteModelPassthrough:
|
|
"""
|
|
ComfyUI Custom Node that:
|
|
- Accepts any input (data) and a MODEL (`model`).
|
|
- Deletes the MODEL completely (VRAM + RAM).
|
|
- Passes through the data unchanged.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"data": (any_typ,),
|
|
"model": (any_typ,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (any_typ,)
|
|
RETURN_NAMES = ("output",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "Memory Management"
|
|
|
|
def _free_gpu_vram(self):
|
|
if not torch.cuda.is_available():
|
|
return "CUDA not available."
|
|
|
|
initial = torch.cuda.memory_allocated()
|
|
torch.cuda.empty_cache()
|
|
final = torch.cuda.memory_allocated()
|
|
freed = max(0, initial - final)
|
|
return f"GPU VRAM freed: {freed/1e9:.2f} GB"
|
|
|
|
def _free_system_ram(self):
|
|
before = psutil.virtual_memory().percent
|
|
collected = gc.collect()
|
|
after = psutil.virtual_memory().percent
|
|
freed = before - after
|
|
return f"System RAM freed: {freed:.2f}%, GC collected {collected} objects."
|
|
|
|
def run(self, data, model):
|
|
logs = []
|
|
|
|
try:
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
logs.append("Deleted MODEL completely from VRAM and RAM.")
|
|
except Exception as e:
|
|
logs.append(f"Failed to delete model: {e}")
|
|
|
|
logs.append(self._free_gpu_vram())
|
|
logs.append(self._free_system_ram())
|
|
|
|
return (data,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DeleteModelPassthrough": DeleteModelPassthrough,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DeleteModelPassthrough": "Delete Model (Passthrough Any)",
|
|
}
|