feat(dynamic_ramcache): 支持新版ComfyUI双RAM缓存阈值，新增inactive参数
重构缓存参数处理逻辑，适配ComfyUI新版`--cache-ram`双参数格式 更新中英文README文档，补充新参数的使用说明与阈值解释 优化缓存清理与切换的执行逻辑，兼容旧版ComfyUI环境 修复缓存实例创建的参数兼容问题
This commit is contained in:
windecay
@@ -1,6 +1,7 @@
|
||||
import gc
|
||||
import logging
|
||||
import time
|
||||
import gc
|
||||
import inspect
|
||||
import logging
|
||||
import time
|
||||
|
||||
caching = None
|
||||
execution = None
|
||||
@@ -44,25 +45,26 @@ class DynamicRAMCacheControl:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mode": (["CLASSIC (No Eviction)", "RAM_PRESSURE (Auto Purge)"], {"default": "RAM_PRESSURE (Auto Purge)"}),
|
||||
"cleanup_threshold": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 256.0, "step": 0.1, "tooltip": "Minimum free RAM to maintain (GB)"}),
|
||||
},
|
||||
"optional": {
|
||||
"any_input": (any_type, {}),
|
||||
}
|
||||
}
|
||||
"required": {
|
||||
"mode": (["CLASSIC (No Eviction)", "RAM_PRESSURE (Auto Purge)"], {"default": "RAM_PRESSURE (Auto Purge)"}),
|
||||
"cleanup_threshold": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 256.0, "step": 0.1, "tooltip": "Active cache free RAM threshold (GB)"}),
|
||||
},
|
||||
"optional": {
|
||||
"inactive_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 256.0, "step": 0.1, "tooltip": "Inactive cache / pinned memory threshold (GB). 0 keeps ComfyUI's current value."}),
|
||||
"any_input": (any_type, {}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (any_type,)
|
||||
RETURN_NAMES = ("output_passthrough",)
|
||||
FUNCTION = "manage_cache"
|
||||
CATEGORY = "utils/dynamic_ramcache"
|
||||
|
||||
def manage_cache(self, mode, cleanup_threshold, any_input=None):
|
||||
if caching is not None and execution is not None:
|
||||
self._execute_cache_logic(mode, cleanup_threshold)
|
||||
else:
|
||||
logging.warning("[DynamicRAMCache] Plugin disabled: Missing internal modules.")
|
||||
def manage_cache(self, mode, cleanup_threshold, inactive_threshold=0.0, any_input=None):
|
||||
if caching is not None and execution is not None:
|
||||
self._execute_cache_logic(mode, cleanup_threshold, inactive_threshold)
|
||||
else:
|
||||
logging.warning("[DynamicRAMCache] Plugin disabled: Missing internal modules.")
|
||||
|
||||
if any_input is not None:
|
||||
return (any_input,)
|
||||
@@ -73,24 +75,24 @@ class DynamicRAMCacheControl:
|
||||
except ImportError:
|
||||
return (None,)
|
||||
|
||||
def _execute_cache_logic(self, mode, cleanup_threshold):
|
||||
|
||||
target_mode_ram = "RAM_PRESSURE" in mode
|
||||
|
||||
executor = self._find_executor()
|
||||
def _execute_cache_logic(self, mode, cleanup_threshold, inactive_threshold=0.0):
|
||||
|
||||
target_mode_ram = "RAM_PRESSURE" in mode
|
||||
|
||||
executor = self._find_executor()
|
||||
|
||||
if executor is None:
|
||||
logging.warning("[DynamicRAMCache] PromptExecutor not found.")
|
||||
return
|
||||
|
||||
if not hasattr(executor, 'cache_args'):
|
||||
executor.cache_args = {}
|
||||
|
||||
old_ram_arg = executor.cache_args.get('ram', 0)
|
||||
executor.cache_args['ram'] = cleanup_threshold
|
||||
|
||||
cache_set = self._get_cache_set(executor)
|
||||
if cache_set is None:
|
||||
old_ram_arg, old_ram_inactive_arg, active_headroom, inactive_headroom, supports_inactive = self._update_cache_args(
|
||||
executor,
|
||||
cleanup_threshold,
|
||||
inactive_threshold,
|
||||
)
|
||||
|
||||
cache_set = self._get_cache_set(executor)
|
||||
if cache_set is None:
|
||||
return
|
||||
|
||||
current_cache = cache_set.outputs
|
||||
@@ -110,23 +112,130 @@ class DynamicRAMCacheControl:
|
||||
|
||||
is_currently_ram = isinstance(current_cache, RAMPressureCacheClass)
|
||||
|
||||
if target_mode_ram and not is_currently_ram:
|
||||
self._switch_to_ram_pressure(cache_set, current_cache, caching)
|
||||
logging.info(f"[DynamicRAMCache] Switched mode: CLASSIC -> RAM_PRESSURE (Headroom: {cleanup_threshold}GB)")
|
||||
|
||||
elif not target_mode_ram and is_currently_ram:
|
||||
self._switch_to_classic(cache_set, current_cache, caching)
|
||||
logging.info(f"[DynamicRAMCache] Switched mode: RAM_PRESSURE -> CLASSIC")
|
||||
|
||||
elif target_mode_ram and is_currently_ram:
|
||||
if old_ram_arg != cleanup_threshold:
|
||||
logging.info(f"[DynamicRAMCache] Updated RAM Headroom: {old_ram_arg}GB -> {cleanup_threshold}GB")
|
||||
|
||||
if target_mode_ram and hasattr(cache_set.outputs, 'poll'):
|
||||
try:
|
||||
cache_set.outputs.poll(cleanup_threshold)
|
||||
except Exception:
|
||||
pass
|
||||
if target_mode_ram and not is_currently_ram:
|
||||
self._switch_to_ram_pressure(cache_set, current_cache, caching)
|
||||
self._set_executor_cache_type(executor, target_mode_ram)
|
||||
logging.info(self._format_headroom_log("[DynamicRAMCache] Switched mode: CLASSIC -> RAM_PRESSURE", active_headroom, inactive_headroom, supports_inactive))
|
||||
|
||||
elif not target_mode_ram and is_currently_ram:
|
||||
self._switch_to_classic(cache_set, current_cache, caching)
|
||||
self._set_executor_cache_type(executor, target_mode_ram)
|
||||
logging.info(f"[DynamicRAMCache] Switched mode: RAM_PRESSURE -> CLASSIC")
|
||||
|
||||
elif target_mode_ram and is_currently_ram:
|
||||
self._set_executor_cache_type(executor, target_mode_ram)
|
||||
if old_ram_arg != active_headroom or old_ram_inactive_arg != inactive_headroom:
|
||||
logging.info(self._format_headroom_update_log(old_ram_arg, old_ram_inactive_arg, active_headroom, inactive_headroom, supports_inactive))
|
||||
|
||||
if target_mode_ram:
|
||||
self._release_ram_cache(cache_set.outputs, active_headroom, inactive_headroom, supports_inactive)
|
||||
|
||||
def _read_threshold(self, value, default_value):
|
||||
try:
|
||||
value = float(value)
|
||||
except (TypeError, ValueError):
|
||||
value = default_value
|
||||
return value
|
||||
|
||||
def _update_cache_args(self, executor, cleanup_threshold, inactive_threshold):
|
||||
if not hasattr(executor, 'cache_args') or executor.cache_args is None:
|
||||
executor.cache_args = {}
|
||||
|
||||
cache_args = executor.cache_args
|
||||
supports_inactive = self._supports_inactive_cache_arg(executor)
|
||||
active_headroom = self._read_threshold(cleanup_threshold, 2.0)
|
||||
old_ram_arg = self._read_threshold(cache_args.get('ram'), active_headroom)
|
||||
|
||||
if supports_inactive:
|
||||
old_ram_inactive_arg = self._read_threshold(cache_args.get('ram_inactive'), old_ram_arg)
|
||||
inactive_value = self._read_threshold(inactive_threshold, 0.0)
|
||||
if inactive_value > 0:
|
||||
inactive_headroom = inactive_value
|
||||
elif old_ram_inactive_arg > 0:
|
||||
inactive_headroom = old_ram_inactive_arg
|
||||
else:
|
||||
inactive_headroom = active_headroom
|
||||
cache_args['ram_inactive'] = inactive_headroom
|
||||
else:
|
||||
old_ram_inactive_arg = None
|
||||
inactive_headroom = None
|
||||
cache_args.pop('ram_inactive', None)
|
||||
|
||||
cache_args.setdefault('lru', 0)
|
||||
cache_args['ram'] = active_headroom
|
||||
return old_ram_arg, old_ram_inactive_arg, active_headroom, inactive_headroom, supports_inactive
|
||||
|
||||
def _supports_inactive_cache_arg(self, executor):
|
||||
cache_args = getattr(executor, 'cache_args', None)
|
||||
if isinstance(cache_args, dict) and 'ram_inactive' in cache_args:
|
||||
return True
|
||||
|
||||
PromptExecutor = getattr(execution, 'PromptExecutor', None)
|
||||
execute_async = getattr(PromptExecutor, 'execute_async', None)
|
||||
if execute_async is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
return 'ram_inactive' in inspect.getsource(execute_async)
|
||||
except (OSError, TypeError):
|
||||
return False
|
||||
|
||||
def _set_executor_cache_type(self, executor, target_mode_ram):
|
||||
CacheType = getattr(execution, 'CacheType', None)
|
||||
if CacheType is None:
|
||||
return
|
||||
|
||||
if target_mode_ram:
|
||||
ram_type = getattr(CacheType, 'RAM_PRESSURE', None)
|
||||
if ram_type is not None:
|
||||
executor.cache_type = ram_type
|
||||
return
|
||||
|
||||
classic_type = getattr(CacheType, 'CLASSIC', None)
|
||||
if classic_type is not None:
|
||||
executor.cache_type = classic_type
|
||||
|
||||
def _format_headroom_log(self, prefix, active_headroom, inactive_headroom, supports_inactive):
|
||||
if supports_inactive:
|
||||
return f"{prefix} (active: {active_headroom}GB, inactive: {inactive_headroom}GB)"
|
||||
return f"{prefix} (headroom: {active_headroom}GB)"
|
||||
|
||||
def _format_headroom_update_log(self, old_ram_arg, old_ram_inactive_arg, active_headroom, inactive_headroom, supports_inactive):
|
||||
if supports_inactive:
|
||||
return f"[DynamicRAMCache] Updated RAM headroom: active {old_ram_arg}GB -> {active_headroom}GB, inactive {old_ram_inactive_arg}GB -> {inactive_headroom}GB"
|
||||
return f"[DynamicRAMCache] Updated RAM headroom: {old_ram_arg}GB -> {active_headroom}GB"
|
||||
|
||||
def _release_ram_cache(self, cache, active_headroom, inactive_headroom, supports_inactive):
|
||||
ram_release = getattr(cache, 'ram_release', None)
|
||||
if callable(ram_release):
|
||||
try:
|
||||
if supports_inactive and inactive_headroom is not None:
|
||||
ram_release(int(inactive_headroom * (1024 ** 3)))
|
||||
ram_release(int(active_headroom * (1024 ** 3)), free_active=True)
|
||||
return
|
||||
except TypeError:
|
||||
try:
|
||||
ram_release(int(active_headroom * (1024 ** 3)))
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logging.warning(f"[DynamicRAMCache] RAM release failed: {e}")
|
||||
|
||||
poll = getattr(cache, 'poll', None)
|
||||
if callable(poll):
|
||||
try:
|
||||
if supports_inactive:
|
||||
poll(ram=active_headroom, ram_inactive=inactive_headroom)
|
||||
else:
|
||||
poll(active_headroom)
|
||||
except TypeError:
|
||||
try:
|
||||
poll(active_headroom)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _find_executor(self):
|
||||
for obj in gc.get_objects():
|
||||
@@ -169,7 +278,7 @@ class DynamicRAMCacheControl:
|
||||
if not key_class:
|
||||
key_class = getattr(caching_mod, 'CacheKeySetInputSignature', None)
|
||||
|
||||
new_cache = caching_mod.RAMPressureCache(key_class)
|
||||
new_cache = self._create_cache(caching_mod.RAMPressureCache, key_class)
|
||||
self._migrate_cache_data(old_cache, new_cache)
|
||||
if getattr(new_cache, 'timestamps', None) is None:
|
||||
new_cache.timestamps = {}
|
||||
@@ -197,12 +306,18 @@ class DynamicRAMCacheControl:
|
||||
if not key_class:
|
||||
key_class = getattr(caching_mod, 'CacheKeySetInputSignature', None)
|
||||
|
||||
new_cache = caching_mod.HierarchicalCache(key_class)
|
||||
new_cache = self._create_cache(caching_mod.HierarchicalCache, key_class)
|
||||
self._migrate_cache_data(old_cache, new_cache)
|
||||
|
||||
self._update_cache_set(cache_set, new_cache)
|
||||
|
||||
def _migrate_cache_data(self, old_cache, new_cache):
|
||||
self._update_cache_set(cache_set, new_cache)
|
||||
|
||||
def _create_cache(self, cache_class, key_class):
|
||||
try:
|
||||
return cache_class(key_class, enable_providers=True)
|
||||
except TypeError:
|
||||
return cache_class(key_class)
|
||||
|
||||
def _migrate_cache_data(self, old_cache, new_cache):
|
||||
"""迁移缓存核心数据"""
|
||||
try:
|
||||
old_dict = getattr(old_cache, '__dict__', None)
|
||||
@@ -262,19 +377,20 @@ class RAMCacheExtremeCleanup(DynamicRAMCacheControl):
|
||||
if executor is None:
|
||||
logging.warning("[DynamicRAMCache] PromptExecutor not found.")
|
||||
else:
|
||||
if not hasattr(executor, 'cache_args'):
|
||||
executor.cache_args = {}
|
||||
old_ram_arg = executor.cache_args.get('ram', 2.0)
|
||||
cache_set = self._get_cache_set(executor)
|
||||
if cache_set is not None:
|
||||
if not hasattr(executor, 'cache_args') or executor.cache_args is None:
|
||||
executor.cache_args = {}
|
||||
old_ram_arg = self._read_threshold(executor.cache_args.get('ram'), 2.0)
|
||||
old_ram_inactive_arg = self._read_threshold(executor.cache_args.get('ram_inactive'), old_ram_arg)
|
||||
cache_set = self._get_cache_set(executor)
|
||||
if cache_set is not None:
|
||||
RAMPressureCacheClass = getattr(caching, 'RAMPressureCache', None)
|
||||
if RAMPressureCacheClass:
|
||||
is_currently_ram = isinstance(cache_set.outputs, RAMPressureCacheClass)
|
||||
old_mode = "RAM_PRESSURE (Auto Purge)" if is_currently_ram else "CLASSIC (No Eviction)"
|
||||
else:
|
||||
old_mode = "CLASSIC (No Eviction)"
|
||||
self._execute_cache_logic("RAM_PRESSURE (Auto Purge)", purge_threshold)
|
||||
self._execute_cache_logic(old_mode, old_ram_arg)
|
||||
self._execute_cache_logic("RAM_PRESSURE (Auto Purge)", purge_threshold, purge_threshold)
|
||||
self._execute_cache_logic(old_mode, old_ram_arg, old_ram_inactive_arg)
|
||||
else:
|
||||
logging.warning("[DynamicRAMCache] Plugin disabled: Missing internal modules.")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user