diff --git a/README.md b/README.md index 7f52af8..90ee648 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,8 @@ A custom node for ComfyUI that dynamically manages RAM usage by intelligently co - **Dual Cache Management Modes**: - CLASSIC (No Eviction): Standard cache behavior with no automatic purging - RAM_PRESSURE (Auto Purge): Smart cache purging when system memory falls below the threshold -- **RAM Cache Extreme Cleanup (🧹)**: Performs a one-time aggressive purge and automatically restores the previous mode and threshold. -- **Configurable Memory Threshold**: Fine-tune the minimum amount of free RAM to maintain +- **RAM Cache Extreme Cleanup (🧹)**: Performs a one-time aggressive purge and automatically restores the previous mode and thresholds. +- **Configurable Memory Thresholds**: Supports the newer `--cache-ram active inactive` RAM threshold pair - **Seamless Cache Migration**: Preserves essential cache data when switching between modes - **Real-time Memory Monitoring**: Continuously checks system RAM and virtual memory availability @@ -34,7 +34,13 @@ A custom node for ComfyUI that dynamically manages RAM usage by intelligently co - **cleanup_threshold**: - Range: 0.1 - 256.0 GB - Default: 2.0 GB - - Determines the minimum amount of free RAM to maintain (in gigabytes) + - Active cache free RAM threshold, matching the first `--cache-ram` value + +- **inactive_threshold**: + - Range: 0 - 256.0 GB + - Default: 0 GB + - Inactive cache / pinned memory threshold, matching the optional second `--cache-ram` value + - Optional for old workflows. `0` keeps ComfyUI's current value, and older ComfyUI builds ignore it ### Extreme Cleanup Parameters @@ -44,7 +50,7 @@ A custom node for ComfyUI that dynamically manages RAM usage by intelligently co - Temporary headroom used for the one-shot purge - **restore behavior**: - - Automatically restores the mode and threshold that were active before cleanup + - Automatically restores the mode and both thresholds that were active before cleanup 3. Connect the node anywhere in your workflow - it will continuously monitor and manage cache based on your settings 4. Optionally add the "RAMCacheExtremeCleanup" node at the end of a workflow to trigger a one-shot cleanup and restore the previous state @@ -62,7 +68,7 @@ In CLASSIC mode, the node maintains standard cache behavior without automatic pu In RAM_PRESSURE mode, the node actively monitors system memory and intelligently purges cache when: 1. Available system RAM falls below the specified cleanup_threshold -2. The system is experiencing memory pressure +2. Inactive cache / pinned memory exceeds the inactive_threshold behavior in newer ComfyUI builds When purging, it prioritizes removing older, less recently used cache items while attempting to preserve critical data. diff --git a/README_zh.md b/README_zh.md index 5a2cd99..87c2469 100644 --- a/README_zh.md +++ b/README_zh.md @@ -8,7 +8,7 @@ - **动态缓存模式切换**：支持在CLASSIC（无内存回收）和RAM_PRESSURE（自动内存清理）模式之间切换 - **智能内存管理**：自动监控和清理 RAM 缓存，当内存不足时释放不必要的缓存数据 -- **自定义清理阈值**：允许用户设置最小空闲内存保持量，灵活适应不同硬件配置 +- **自定义清理阈值**：支持新版 `--cache-ram active inactive` 两个 RAM 阈值 - **无缝数据迁移**：在切换缓存模式时保留现有缓存数据，避免重复计算 - **直观的节点接口**：简单易用的参数设置，适用于不同水平的用户 - **极限清理节点**：一次性强清理（包含虚拟内存占用）并自动恢复清理前的模式与阈值 @@ -25,7 +25,8 @@ 1. 在ComfyUI中，从`utils/dynamic_ramcache`类别中添加`DynamicRAMCacheControl`节点 2. 配置以下参数： - **mode**：选择缓存模式（CLASSIC或RAM_PRESSURE） - - **cleanup_threshold**：设置最小空闲内存保持量（GB） + - **cleanup_threshold**：设置 active cache 空闲内存阈值（GB） + - **inactive_threshold**：可选。设置 inactive cache / pinned memory 阈值（GB），填 `0` 或旧工作流缺少此项时沿用 ComfyUI 当前值 3. 可以选择性地连接任意输入到`any_input`端口（节点会透传此输入） 4. 可选：在工作流末尾添加`RAMCacheExtremeCleanup`节点，执行一次性清理并恢复之前状态 5. 运行工作流 @@ -37,16 +38,22 @@ - **CLASSIC (No Eviction)**：传统缓存模式，不会自动清理缓存，可能导致内存使用持续增长 - **RAM_PRESSURE (Auto Purge)**：自动内存清理模式，当可用内存低于设定阈值时会自动清理缓存 -### cleanup_threshold（清理阈值） +### cleanup_threshold（active 阈值） - 类型：浮点数（0.1-256.0 GB） - 默认值：2.0 GB -- 说明：系统尝试维持的最小空闲内存量。当可用内存低于此值时，RAM_PRESSURE模式会触发缓存清理 +- 说明：对应新版 `--cache-ram` 的第一个值。可用内存低于此值时，RAM_PRESSURE模式会清理 active cache + +### inactive_threshold（inactive 阈值） + +- 类型：浮点数（0-256.0 GB） +- 默认值：0 +- 说明：对应新版 `--cache-ram` 的第二个值。旧版 ComfyUI 不使用该参数时会自动忽略；`0` 表示不改 ComfyUI 当前值；大于 `0` 时会更新 inactive cache / pinned memory 阈值 ### 极限清理参数 - **purge_threshold**：一次性清理时使用的临时阈值（默认 256.0 GB） -- **恢复行为**：自动恢复清理前的模式与阈值 +- **恢复行为**：自动恢复清理前的模式、active 阈值和 inactive 阈值 ### 输出 @@ -59,7 +66,7 @@ 1. 检测ComfyUI的`PromptExecutor`实例 2. 根据所选模式在`RAMPressureCache`和`HierarchicalCache`之间切换 3. 在切换过程中保留现有的缓存数据 -4. 当处于RAM_PRESSURE模式时，监控系统内存并在需要时触发缓存清理 +4. 当处于RAM_PRESSURE模式时，按 active / inactive 两个阈值触发缓存清理 ## 兼容性要求 diff --git a/nodes.py b/nodes.py index c8e3e71..b686525 100644 --- a/nodes.py +++ b/nodes.py @@ -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.")