295 lines
11 KiB
Python
295 lines
11 KiB
Python
import psutil
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
import time
|
|
import platform
|
|
import gc
|
|
from server import PromptServer
|
|
import comfy.model_management
|
|
|
|
try:
|
|
import torch
|
|
HAS_TORCH = True
|
|
except ImportError:
|
|
HAS_TORCH = False
|
|
|
|
class AnyType(str):
|
|
def __eq__(self, _) -> bool:
|
|
return True
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
|
|
|
|
def safe_int(value, default=0):
|
|
if isinstance(value, int):
|
|
return value
|
|
try:
|
|
return int(float(str(value).strip()))
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def safe_bool(value, default=False):
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, (int, float)):
|
|
return bool(value)
|
|
if isinstance(value, str):
|
|
return value.strip().lower() in ("true", "1", "yes")
|
|
return default
|
|
|
|
|
|
def check_memory_critical(threshold_mb=500):
|
|
try:
|
|
memory = psutil.virtual_memory()
|
|
available_mb = memory.available / (1024 * 1024)
|
|
return available_mb < threshold_mb, available_mb
|
|
except:
|
|
return False, 0
|
|
|
|
|
|
def force_memory_cleanup():
|
|
gc.collect()
|
|
gc.collect()
|
|
gc.collect()
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
try:
|
|
comfy.model_management.soft_empty_cache()
|
|
except:
|
|
pass
|
|
|
|
|
|
class VRAMCleanup:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"batch_index": ("INT", {
|
|
"default": 60, "min": 0, "step": 1,
|
|
"label": "批次索引"
|
|
}),
|
|
"batch_clear": ("INT", {
|
|
"default": 60, "min": 1, "step": 1,
|
|
"label": "清理间隔"
|
|
}),
|
|
"bypass_all": ("BOOLEAN", {"default": False, "label": "忽略全部"}),
|
|
"offload_model": ("BOOLEAN", {"default": True}),
|
|
"offload_cache": ("BOOLEAN", {"default": True}),
|
|
"critical_threshold_mb": ("INT", {
|
|
"default": 500, "min": 100, "max": 8192, "step": 100,
|
|
"label": "危险阈值(MB)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"anything": (any, {}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (any, "INT",)
|
|
RETURN_NAMES = ("anything", "batch_index",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "empty_cache"
|
|
CATEGORY = "Memory Management"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, **kwargs):
|
|
return float(time.time())
|
|
|
|
def empty_cache(self, bypass_all, offload_model, offload_cache, batch_index, batch_clear, critical_threshold_mb=500, anything=None, unique_id=None, extra_pnginfo=None):
|
|
try:
|
|
bypass_all = safe_bool(bypass_all, False)
|
|
offload_model = safe_bool(offload_model, True)
|
|
offload_cache = safe_bool(offload_cache, True)
|
|
batch_index = safe_int(batch_index, 0)
|
|
batch_clear = safe_int(batch_clear, 1)
|
|
critical_threshold_mb = safe_int(critical_threshold_mb, 500)
|
|
|
|
if batch_clear <= 0:
|
|
batch_clear = 1
|
|
|
|
is_critical, available_mb = check_memory_critical(critical_threshold_mb)
|
|
should_execute = (batch_index % batch_clear == 0)
|
|
|
|
if bypass_all and not is_critical:
|
|
print(f"VRAM清理跳过 (bypass_all=True, 可用: {available_mb:.0f}MB)")
|
|
elif not should_execute and not is_critical:
|
|
print(f"VRAM清理跳过 (index={batch_index} % batch={batch_clear} = {batch_index % batch_clear}, 可用: {available_mb:.0f}MB)")
|
|
else:
|
|
if is_critical:
|
|
print(f"⚠️ 内存危险! 可用仅 {available_mb:.0f}MB < {critical_threshold_mb}MB, 强制清理!")
|
|
|
|
if offload_model:
|
|
comfy.model_management.unload_all_models()
|
|
|
|
if offload_cache:
|
|
force_memory_cleanup()
|
|
PromptServer.instance.prompt_queue.set_flag("free_memory", True)
|
|
|
|
_, after_mb = check_memory_critical(critical_threshold_mb)
|
|
print(f"VRAM清理完成 [卸载模型: {offload_model}, 清空缓存: {offload_cache}, index: {batch_index}, 可用: {available_mb:.0f}MB -> {after_mb:.0f}MB]")
|
|
|
|
except Exception as e:
|
|
print(f"VRAM清理失败: {type(e).__name__}: {e}")
|
|
try:
|
|
force_memory_cleanup()
|
|
except:
|
|
pass
|
|
|
|
return (anything, batch_index,)
|
|
|
|
|
|
class RAMCleanup:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"batch_index": ("INT", {
|
|
"default": 30, "min": 0, "step": 1,
|
|
"label": "当前索引"
|
|
}),
|
|
"batch_clear": ("INT", {
|
|
"default": 30, "min": 1, "step": 1,
|
|
"label": "批次间隔"
|
|
}),
|
|
"bypass_all": ("BOOLEAN", {"default": False, "label": "忽略全部"}),
|
|
"clean_file_cache": ("BOOLEAN", {"default": True, "label": "清理文件缓存"}),
|
|
"clean_processes": ("BOOLEAN", {"default": True, "label": "清理进程内存"}),
|
|
"clean_dlls": ("BOOLEAN", {"default": True, "label": "清理未使用DLL"}),
|
|
"retry_times": ("INT", {
|
|
"default": 3, "min": 1, "max": 10, "step": 1,
|
|
"label": "重试次数"
|
|
}),
|
|
"critical_threshold_mb": ("INT", {
|
|
"default": 500, "min": 100, "max": 8192, "step": 100,
|
|
"label": "危险阈值(MB)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"anything": (any, {}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
}
|
|
}
|
|
|
|
# ★★★ 两个输出：anything + batch_index ★★★
|
|
RETURN_TYPES = (any, "INT",)
|
|
RETURN_NAMES = ("anything", "batch_index",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "clean_ram"
|
|
CATEGORY = "Memory Management"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, **kwargs):
|
|
return float(time.time())
|
|
|
|
def get_ram_usage(self):
|
|
memory = psutil.virtual_memory()
|
|
return memory.percent, memory.available / (1024 * 1024)
|
|
|
|
def clean_ram(self, bypass_all, clean_file_cache, clean_processes, clean_dlls, retry_times, batch_index, batch_clear, critical_threshold_mb=500, anything=None, unique_id=None, extra_pnginfo=None):
|
|
try:
|
|
bypass_all = safe_bool(bypass_all, False)
|
|
clean_file_cache = safe_bool(clean_file_cache, True)
|
|
clean_processes = safe_bool(clean_processes, True)
|
|
clean_dlls = safe_bool(clean_dlls, True)
|
|
retry_times = safe_int(retry_times, 3)
|
|
batch_index = safe_int(batch_index, 0)
|
|
batch_clear = safe_int(batch_clear, 1)
|
|
critical_threshold_mb = safe_int(critical_threshold_mb, 500)
|
|
|
|
if batch_clear <= 0:
|
|
batch_clear = 1
|
|
if retry_times <= 0:
|
|
retry_times = 1
|
|
|
|
before_usage, before_available = self.get_ram_usage()
|
|
system = platform.system()
|
|
|
|
is_critical, available_mb = check_memory_critical(critical_threshold_mb)
|
|
should_execute = (batch_index % batch_clear == 0)
|
|
|
|
if bypass_all and not is_critical:
|
|
print(f"RAM清理跳过 (bypass_all=True, 可用: {available_mb:.0f}MB)")
|
|
elif not should_execute and not is_critical:
|
|
print(f"RAM清理跳过 (index={batch_index} % batch={batch_clear} = {batch_index % batch_clear}, 可用: {available_mb:.0f}MB)")
|
|
else:
|
|
if is_critical:
|
|
print(f"⚠️ 内存危险! 可用仅 {available_mb:.0f}MB < {critical_threshold_mb}MB, 强制清理!")
|
|
|
|
force_memory_cleanup()
|
|
|
|
for attempt in range(retry_times):
|
|
if clean_file_cache:
|
|
try:
|
|
if system == "Windows":
|
|
ctypes.windll.kernel32.SetSystemFileCacheSize(-1, -1, 0)
|
|
elif system == "Linux":
|
|
libc = ctypes.CDLL("libc.so.6")
|
|
libc.malloc_trim(0)
|
|
except:
|
|
pass
|
|
|
|
if clean_processes:
|
|
if system == "Windows":
|
|
for process in psutil.process_iter(['pid', 'name']):
|
|
try:
|
|
handle = ctypes.windll.kernel32.OpenProcess(
|
|
wintypes.DWORD(0x001F0FFF),
|
|
wintypes.BOOL(False),
|
|
wintypes.DWORD(process.info['pid'])
|
|
)
|
|
ctypes.windll.psapi.EmptyWorkingSet(handle)
|
|
ctypes.windll.kernel32.CloseHandle(handle)
|
|
except:
|
|
continue
|
|
|
|
if clean_dlls:
|
|
try:
|
|
if system == "Windows":
|
|
ctypes.windll.kernel32.SetProcessWorkingSetSize(-1, -1, -1)
|
|
except:
|
|
pass
|
|
|
|
time.sleep(1)
|
|
|
|
_, current_mb = check_memory_critical(critical_threshold_mb)
|
|
if current_mb > critical_threshold_mb * 2:
|
|
print(f" 第{attempt + 1}次清理后内存充足 ({current_mb:.0f}MB), 停止重试")
|
|
break
|
|
|
|
after_usage, after_available = self.get_ram_usage()
|
|
freed_mb = after_available - before_available
|
|
print(f"RAM清理完成 [{before_usage:.1f}% -> {after_usage:.1f}%, 释放: {freed_mb:.0f}MB, index: {batch_index}, batch: {batch_clear}]")
|
|
|
|
except Exception as e:
|
|
print(f"RAM清理失败: {type(e).__name__}: {e}")
|
|
print(f" 调试 -> batch_index={repr(batch_index)}(type={type(batch_index).__name__}), batch_clear={repr(batch_clear)}(type={type(batch_clear).__name__})")
|
|
try:
|
|
force_memory_cleanup()
|
|
except:
|
|
pass
|
|
|
|
# ★★★ 返回 anything + batch_index ★★★
|
|
return (anything, batch_index,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DocKr.VRAM-Cleanup": VRAMCleanup,
|
|
"DocKr.RAM-Cleanup": RAMCleanup,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DocKr.VRAM-Cleanup": "DocKr / VRAM-Cleanup",
|
|
"DocKr.RAM-Cleanup": "DocKr / RAM-Cleanup",
|
|
}
|