nodes.py

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 = 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:
            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 _find_executor(self):
        for obj in gc.get_objects():
            try:
                if obj.__class__.__name__ == 'PromptExecutor':
                    return obj
            except (ReferenceError, AttributeError):
                continue
            except Exception:
                continue
        return None

    def _get_cache_set(self, executor):
        if not hasattr(executor, 'caches'):
            logging.warning("[DynamicRAMCache] PromptExecutor has no 'caches' attribute.")
            return None
        
        cache_set = executor.caches

        if not hasattr(cache_set, 'outputs'):
            logging.warning("[DynamicRAMCache] CacheSet has no 'outputs' attribute.")
            return None
        return cache_set

    def _update_cache_set(self, cache_set, new_cache):
        try:
            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
        except (ReferenceError, AttributeError):
            logging.warning("[DynamicRAMCache] Failed to update cache_set: object no longer exists.")
        except Exception as e:
            logging.warning(f"[DynamicRAMCache] Unexpected error updating cache_set: {e}")

    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)
        if getattr(new_cache, 'timestamps', None) is None:
            new_cache.timestamps = {}
        if getattr(new_cache, 'used_generation', None) is None:
            new_cache.used_generation = {}
        if getattr(new_cache, 'children', None) is None:
            new_cache.children = {}
        if getattr(new_cache, 'generation', None) is None:
            new_cache.generation = 1
        if getattr(new_cache, 'min_generation', None) is None:
            new_cache.min_generation = 0

        if isinstance(getattr(new_cache, 'cache', None), dict) and isinstance(new_cache.timestamps, dict) and isinstance(new_cache.used_generation, dict):
            now = time.time()
            for key in new_cache.cache:
                if key not in new_cache.timestamps:
                    new_cache.timestamps[key] = now
                if key not in new_cache.used_generation:
                    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):
        """迁移缓存核心数据"""
        try:
            old_dict = getattr(old_cache, '__dict__', None)
            new_dict = getattr(new_cache, '__dict__', None)
            if isinstance(old_dict, dict) and isinstance(new_dict, dict):
                new_dict.update(old_dict)
            else:
                old_cache_data = getattr(old_cache, 'cache', None)
                if old_cache_data is not None:
                    new_cache.cache = old_cache_data

                old_subcaches = getattr(old_cache, 'subcaches', None)
                if old_subcaches is not None:
                    new_cache.subcaches = old_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)

            if getattr(new_cache, 'cache', None) is None:
                new_cache.cache = {}
            if getattr(new_cache, 'subcaches', None) is None:
                new_cache.subcaches = {}

            if hasattr(old_cache, 'is_changed_cache'):
                new_cache.is_changed_cache = old_cache.is_changed_cache
            if getattr(new_cache, 'is_changed_cache', None) is None:
                new_cache.is_changed_cache = {}
        except (ReferenceError, AttributeError):
            logging.warning("[DynamicRAMCache] Failed to migrate cache data: source object no longer exists.")
        except Exception as e:
            logging.warning(f"[DynamicRAMCache] Unexpected error migrating cache data: {e}")

class RAMCacheExtremeCleanup(DynamicRAMCacheControl):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "purge_threshold": ("FLOAT", {"default": 256.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 = "extreme_cleanup"
    CATEGORY = "utils/dynamic_ramcache"

    def extreme_cleanup(self, purge_threshold, any_input=None):
        if caching is not None and execution is not None:
            executor = self._find_executor()
            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:
                    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)
        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,)