[feat]: add hijack to comfy.utils.load_torch_file

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.

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

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

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)

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)

nodes.py

@@ -0,0 +1,7 @@
+import torch
+
+NODE_CLASS_MAPPINGS = {
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+}

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)