From 36d2162b0510f1c3d04dda401353de03468e8916 Mon Sep 17 00:00:00 2001 From: zhaoyingzhuo Date: Thu, 20 Mar 2025 09:17:21 +0000 Subject: [PATCH] [feat]: add hijack to comfy.utils.load_torch_file --- README.md | 2 +- __init__.py | 3 ++ hijack/__init__.py | 0 hijack/cache_manager.py | 75 +++++++++++++++++++++++++++++++++++++++++ hijack/hijack_list.py | 5 +++ hijack/utils.py | 70 ++++++++++++++++++++++++++++++++++++++ nodes.py | 7 ++++ prestartup_script.py | 12 +++++++ 8 files changed, 173 insertions(+), 1 deletion(-) create mode 100644 __init__.py create mode 100644 hijack/__init__.py create mode 100644 hijack/cache_manager.py create mode 100644 hijack/hijack_list.py create mode 100644 hijack/utils.py create mode 100644 nodes.py create mode 100644 prestartup_script.py diff --git a/README.md b/README.md index 1daeace..a95c9e1 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,2 @@ # ComfyUI_Model_Cache -A model cached-loader custom node for ComfyUI +A model cached-loader custom node for ComfyUI. Recommand Comfy version: v0.3.26. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3d998ea --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/hijack/__init__.py b/hijack/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hijack/cache_manager.py b/hijack/cache_manager.py new file mode 100644 index 0000000..02fc9bb --- /dev/null +++ b/hijack/cache_manager.py @@ -0,0 +1,75 @@ + +from dataclasses import dataclass, field +from functools import _make_key +from typing import Dict, Tuple +import torch +import logging + +logger = logging.getLogger(__name__) + +class ModelValidChecker: + def __init__(self, result): + self.result = result + # module dict should be the first place + self.module = result[0] if isinstance(result, tuple) else result + self.key_count = len(self.module.keys()) + + def is_valid(self) -> bool: + # if tensors on gpu are released, the module keys are incomplete + if self.key_count == len(self.module.keys()): + return True + return False + + def get_result(self) -> torch.nn.Module | Tuple[torch.nn.Module, ...]: + return self.result + +@dataclass +class ModelCache: + models: Dict[str, ModelValidChecker] = field(default_factory=dict) + + def generate_cache_key(self, *args, **kwargs): + args_key = _make_key(args, kwargs, typed=True) + return args_key + + def cached(self, model_key) -> bool: + checker = self.models.get(model_key, None) + if checker: + return checker.is_valid() + return False + + def register_model( + self, + model_key: str, + result: torch.nn.Module | Tuple[torch.nn.Module, ...], + ) -> None: + """ + Register an cache model result to be used in Comfy. + + :code:`result` can be either: + + - A :class:`torch.nn.Module` class directly referencing the model. + - B :A tuple with `torch.nn.Module` Dict at the first place and other kwargs follow-up + """ + self.models[model_key] = ModelValidChecker(result) + + def get_result(self, model_key): + logger.info(f"\033[92mModel_Cache: Return a cached module result with args:{model_key}\033[0m") + return self.models[model_key].get_result() + +_model_cache = ModelCache() + +def get_model_cache(): + global _model_cache + assert _model_cache is not None, "Model Cache has not been initialized." + return _model_cache + +def cache_model(func): + def wrapper(*args, **kwargs): + model_key = get_model_cache().generate_cache_key(*args, **kwargs) + if get_model_cache().cached(model_key): + return get_model_cache().get_result(model_key) + + result = func(*args, **kwargs) + get_model_cache().register_model(model_key, result) + return result + return wrapper \ No newline at end of file diff --git a/hijack/hijack_list.py b/hijack/hijack_list.py new file mode 100644 index 0000000..15f67bc --- /dev/null +++ b/hijack/hijack_list.py @@ -0,0 +1,5 @@ +from .cache_manager import cache_model + +@cache_model +def load_torch_file_cache_model(origin_func, *args, **kwargs): + return origin_func(*args, **kwargs) \ No newline at end of file diff --git a/hijack/utils.py b/hijack/utils.py new file mode 100644 index 0000000..2da51f3 --- /dev/null +++ b/hijack/utils.py @@ -0,0 +1,70 @@ +import importlib +import inspect +from types import FunctionType +from typing import Callable, Union + + +always_true_func = lambda *args, **kwargs: True + +class CondFunc: + def __new__(cls, orig_func, sub_func, cond_func=always_true_func): + self = super(CondFunc, cls).__new__(cls) + if isinstance(orig_func, str): + func_path = orig_func.split('.') + for i in range(len(func_path)-1, -1, -1): + try: + resolved_obj = importlib.import_module('.'.join(func_path[:i])) + break + except ImportError: + pass + try: + for attr_name in func_path[i:-1]: + resolved_obj = getattr(resolved_obj, attr_name) + orig_func = getattr(resolved_obj, func_path[-1]) + setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: self(*args, **kwargs)) + except AttributeError: + print(f"Warning: Failed to resolve {orig_func} for CondFunc hijack") + pass + self.__init__(orig_func, sub_func, cond_func) + return lambda *args, **kwargs: self(*args, **kwargs) + + def __init__(self, orig_func, sub_func, cond_func): + self.__orig_func = orig_func + self.__sub_func = sub_func + self.__cond_func = cond_func + + def __call__(self, *args, **kwargs): + if not self.__cond_func or self.__cond_func(self.__orig_func, *args, **kwargs): + return self.__sub_func(self.__orig_func, *args, **kwargs) + else: + return self.__orig_func(*args, **kwargs) + +def get_func_full_name(func: FunctionType): + module = inspect.getmodule(func) + if module is None: + raise ValueError(f"Cannot get module of function {func}") + return f"{module.__name__}.{func.__qualname__}" + +def hijack_func( + orig_func: Union[str, Callable], + sub_func: Callable, +): + if isinstance(orig_func, FunctionType): + orig_func = get_func_full_name(orig_func) + else: + raise ValueError(f"origin function is not function type") + return CondFunc(orig_func, sub_func) + +def hijack_class_func( + orig_class, + orig_class_func: str, + sub_func: Callable, +): + try: + func = getattr(orig_class, orig_class_func) + except: + raise ValueError(f"{orig_class} has no function named:{orig_class_func}") + module = inspect.getmodule(orig_class) + orig_class_name = module.__name__ + "." + orig_class.__name__ + orig_func_name = orig_class_name + "." + orig_class_func + return CondFunc(orig_func_name, sub_func) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..4a3a513 --- /dev/null +++ b/nodes.py @@ -0,0 +1,7 @@ +import torch + +NODE_CLASS_MAPPINGS = { +} + +NODE_DISPLAY_NAME_MAPPINGS = { +} diff --git a/prestartup_script.py b/prestartup_script.py new file mode 100644 index 0000000..eaf50eb --- /dev/null +++ b/prestartup_script.py @@ -0,0 +1,12 @@ +import torch +import sys +import os +torch._C._cuda_init() +infinigence_path = os.path.join(os.path.dirname(__file__)) +sys.path.append(infinigence_path) +from hijack.utils import hijack_func +from hijack.hijack_list import load_torch_file_cache_model +from comfy.utils import load_torch_file +sys.path.remove(infinigence_path) + +hijack_func(load_torch_file, load_torch_file_cache_model) \ No newline at end of file