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)