diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..5cf7819 --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .nodes import DynamicRAMCacheControl + +NODE_CLASS_MAPPINGS = { + "DynamicRAMCacheControl": DynamicRAMCacheControl +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DynamicRAMCacheControl": "🔥 Dynamic RAM Cache Control" +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..29a1a49 --- /dev/null +++ b/nodes.py @@ -0,0 +1,187 @@ +import gc +import logging +import time + +caching = None +execution = None + +# Import execution module +try: + import execution +except ImportError: + logging.error("[DynamicRAMCache] Failed to import execution module") + logging.error("[DynamicRAMCache] This may be due to ComfyUI version update to 2025.10.31. Please check if module structure has changed.") + +# Import caching module +try: + from comfy_execution import caching + # Check if RAMPressureCache class exists in the imported caching module + if caching is not None and not hasattr(caching, 'RAMPressureCache'): + logging.error("[DynamicRAMCache] RAMPressureCache class not found in caching module") + logging.error("[DynamicRAMCache] This class may only exist in ComfyUI versions after 2025.10.31") +except ImportError: + logging.error("[DynamicRAMCache] Failed to import caching module") + logging.error("[DynamicRAMCache] This may be due to ComfyUI version update to 2025.10.31. Please check if module structure has changed.") + +# Ensure both modules are successfully imported +if execution is None or caching is None: + logging.error("[DynamicRAMCache] Critical module import failed, plugin may not work correctly") + logging.error("[DynamicRAMCache] Plugin compatibility with ComfyUI 2025.10.31 needs to be verified. Module structure may have changed.") + +class AlwaysEqualProxy(str): + def __eq__(self, _): + return True + + def __ne__(self, _): + return False + +any_type = AlwaysEqualProxy("*") + +class DynamicRAMCacheControl: + def __init__(self): + pass + + @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, {}), + } + } + + 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.") + + if any_input is not None: + return (any_input,) + else: + try: + from comfy_execution.graph import ExecutionBlocker + return (ExecutionBlocker(None),) + except ImportError: + return (None,) + + def _execute_cache_logic(self, mode, cleanup_threshold): + + target_mode_ram = "RAM_PRESSURE" in mode + + executor = None + for obj in gc.get_objects(): + if obj.__class__.__name__ == 'PromptExecutor': + executor = obj + break + + 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 + + if not hasattr(executor, 'caches'): + logging.warning("[DynamicRAMCache] PromptExecutor has no 'caches' attribute.") + return + + cache_set = executor.caches + + if not hasattr(cache_set, 'outputs'): + logging.warning("[DynamicRAMCache] CacheSet has no 'outputs' attribute.") + return + + current_cache = cache_set.outputs + + RAMPressureCacheClass = getattr(caching, 'RAMPressureCache', None) + HierarchicalCacheClass = getattr(caching, 'HierarchicalCache', None) + + if not RAMPressureCacheClass: + logging.error("[DynamicRAMCache] RAMPressureCache class not available in caching module") + logging.error("[DynamicRAMCache] This class is required for RAM_PRESSURE mode and may only exist in ComfyUI versions after 2025.10.31") + logging.error("[DynamicRAMCache] Please check your ComfyUI version or consider switching to CLASSIC mode") + return + + if not HierarchicalCacheClass: + logging.error("[DynamicRAMCache] HierarchicalCache class not available in caching module") + return + + 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 + + def _update_cache_set(self, cache_set, new_cache): + + cache_set.outputs = new_cache + + if hasattr(cache_set, 'all') and isinstance(cache_set.all, list): + for i, item in enumerate(cache_set.all): + if i == 0: + cache_set.all[i] = new_cache + + def _switch_to_ram_pressure(self, cache_set, old_cache, caching_mod): + key_class = getattr(old_cache, 'key_class', None) + if not key_class: + key_class = getattr(caching_mod, 'CacheKeySetInputSignature', None) + + new_cache = caching_mod.RAMPressureCache(key_class) + self._migrate_cache_data(old_cache, new_cache) + + new_cache.timestamps = {} + new_cache.used_generation = {} + new_cache.children = {} + new_cache.generation = 1 + new_cache.min_generation = 0 + + now = time.time() + for key in new_cache.cache: + new_cache.timestamps[key] = now + new_cache.used_generation[key] = 0 + + self._update_cache_set(cache_set, new_cache) + + def _switch_to_classic(self, cache_set, old_cache, caching_mod): + key_class = getattr(old_cache, 'key_class', None) + if not key_class: + key_class = getattr(caching_mod, 'CacheKeySetInputSignature', None) + + new_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): + """迁移缓存核心数据""" + new_cache.cache = old_cache.cache + new_cache.subcaches = old_cache.subcaches + new_cache.dynprompt = getattr(old_cache, 'dynprompt', None) + new_cache.cache_key_set = getattr(old_cache, 'cache_key_set', None) + new_cache.initialized = getattr(old_cache, 'initialized', False) \ No newline at end of file