diff --git a/README.md b/README.md index c3c25c0..12945fd 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,15 @@ # ComfyUI_Model_Cache A model cached-loader custom node for ComfyUI. Recommended Comfy version: v0.3.26. +# How does it work? +Some custom nodes call comfy.utils.load_torch_file whey they are executed. However, it is waste of time for these frequently executed workflows. Here is a easy implementation for ComfyUI to cache torch file and avoid from duplicated loading torch files. By Using symbols hijack, many functions that read disk freqently can be cached in host memory. So when they are called next time, it will take no time. In order to trap tensor status after user clicks **`Clear Execution Cache`**(After user clicks that button, the tensors on gpu will be released and the caches are dirty), I add decorator "cache_model" to get the tensor status. If the tensors are released, they will load them again instead of use the dirty cache. + # How to install? +If you have installed comfy manager, you can use command below: ```bash python3 $CUSTOM_MANAGER_CLI/cm-cli.py install https://github.com/infinigence/ComfyUI_Model_Cache.git -``` \ No newline at end of file +``` +If you don't, you can easily use this custom node by moving the whole project to `ComfyUI/custom_nodes`. + +# Examples +In `test/` directory, there are some sample workflows. They demonstrate the acceleration for some custom nodes. \ No newline at end of file diff --git a/hijack/cache_manager.py b/hijack/cache_manager.py index 02fc9bb..e6a700e 100644 --- a/hijack/cache_manager.py +++ b/hijack/cache_manager.py @@ -12,13 +12,24 @@ class ModelValidChecker: 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()) + self.key_count = self.get_latest_key_count() def is_valid(self) -> bool: # if tensors on gpu are released, the module keys are incomplete - if self.key_count == len(self.module.keys()): + if self.key_count == self.get_latest_key_count(): return True return False + + def get_latest_key_count(self) -> int: + if isinstance(self.module, torch.nn.Module): + return len(self.module.state_dict().keys()) + elif isinstance(self.module, Dict): + return len(self.module.keys()) + + logger.warning(f"\033[92mModel_Cache: result is not torch.nn.Module or dict, but {type(self.module)}\033[0m") + logger.warning(f"\033[92mModel_Cache: So cache will never happen for this result! \033[0m") + # will nevel equal and never cache the result + return -1 def get_result(self) -> torch.nn.Module | Tuple[torch.nn.Module, ...]: return self.result diff --git a/hijack/hijack_list.py b/hijack/hijack_list.py index 15f67bc..301b2b5 100644 --- a/hijack/hijack_list.py +++ b/hijack/hijack_list.py @@ -1,5 +1,29 @@ from .cache_manager import cache_model +from .utils import hijack_func, hijack_class_func + +from comfy.utils import load_torch_file +from transformers import AutoModel @cache_model def load_torch_file_cache_model(origin_func, *args, **kwargs): - return origin_func(*args, **kwargs) \ No newline at end of file + return origin_func(*args, **kwargs) + +@cache_model +def hijack_from_pretrained(origin_func, *args, **kwargs): + return origin_func(*args, **kwargs) + +HIJACK_FUNC_MAP = { + load_torch_file: load_torch_file_cache_model, +} + +HIJACK_CLASS_FUNC_MAP = { + AutoModel: ("from_pretrained", hijack_from_pretrained), +} + +def hijack_all(): + for origin_func in HIJACK_FUNC_MAP: + hijack_func(origin_func, HIJACK_FUNC_MAP[origin_func]) + + for origin_class in HIJACK_CLASS_FUNC_MAP: + hijack_class_func(origin_class, *HIJACK_CLASS_FUNC_MAP[origin_class]) + diff --git a/nodes.py b/nodes.py index 4a3a513..d576eb7 100644 --- a/nodes.py +++ b/nodes.py @@ -1,5 +1,3 @@ -import torch - NODE_CLASS_MAPPINGS = { } diff --git a/prestartup_script.py b/prestartup_script.py index eaf50eb..04c19d6 100644 --- a/prestartup_script.py +++ b/prestartup_script.py @@ -2,11 +2,9 @@ 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) +model_cache_custom_path = os.path.join(os.path.dirname(__file__)) +sys.path.append(model_cache_custom_path) +from hijack.hijack_list import hijack_all +sys.path.remove(model_cache_custom_path) -hijack_func(load_torch_file, load_torch_file_cache_model) \ No newline at end of file +hijack_all() \ No newline at end of file diff --git a/test/test_apply_fooocus.json b/test/test_apply_fooocus.json new file mode 100644 index 0000000..2921e91 --- /dev/null +++ b/test/test_apply_fooocus.json @@ -0,0 +1,522 @@ +{ + "last_node_id": 378, + "last_link_id": 137, + "nodes": [ + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [ + 3666.205322265625, + -664.9154052734375 + ], + "size": [ + 315, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 130 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 85 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 68, + 74 + ] + } + ], + "title": "Checkpoint加载器(简易)", + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "sdxl/LEOSAM AIArt 兔狲插画 SDXL大模型_v1.safetensors" + ] + }, + { + "id": 131, + "type": "LoadImage", + "pos": [ + 4205.1982421875, + -288.0040588378906 + ], + "size": [ + 315, + 314 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 127 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 129 + ], + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "1-watch.jpeg", + "image" + ] + }, + { + "id": 296, + "type": "INPAINT_VAEEncodeInpaintConditioning", + "pos": [ + 4774.0859375, + -594.1358032226562 + ], + "size": [ + 415.8000183105469, + 106 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 136 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 137 + }, + { + "name": "vae", + "type": "VAE", + "link": 68 + }, + { + "name": "pixels", + "type": "IMAGE", + "link": 127 + }, + { + "name": "mask", + "type": "MASK", + "link": 129 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 80 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 81 + ] + }, + { + "name": "latent_inpaint", + "type": "LATENT", + "links": [ + 113 + ], + "slot_index": 2 + }, + { + "name": "latent_samples", + "type": "LATENT", + "links": [ + 125 + ] + } + ], + "title": "VAE编码局部重绘条件", + "properties": { + "Node name for S&R": "INPAINT_VAEEncodeInpaintConditioning" + }, + "widgets_values": [] + }, + { + "id": 302, + "type": "VAEDecode", + "pos": [ + 6053.39697265625, + -312.6009216308594 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 73 + }, + { + "name": "vae", + "type": "VAE", + "link": 74 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 131 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 310, + "type": "KSampler", + "pos": [ + 5652.11767578125, + -445.38739013671875 + ], + "size": [ + 315, + 262 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 124 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 80 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 81 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 125 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 73 + ] + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 305966591281514, + "randomize", + 20, + 1, + "euler", + "normal", + 1 + ] + }, + { + "id": 314, + "type": "CLIPTextEncode", + "pos": [ + 4157.634765625, + -587.5955200195312 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 85 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 136, + 137 + ], + "slot_index": 0 + } + ], + "title": "CLIP文本编码器", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "test", + [ + false, + true + ] + ] + }, + { + "id": 367, + "type": "easy applyFooocusInpaint", + "pos": [ + 5310.4697265625, + -624.2901000976562 + ], + "size": [ + 327.5999755859375, + 102 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 130 + }, + { + "name": "latent", + "type": "LATENT", + "link": 113 + } + ], + "outputs": [ + { + "name": "model", + "type": "MODEL", + "links": [ + 124 + ] + } + ], + "title": "应用Fooocus局部重绘", + "properties": { + "Node name for S&R": "easy applyFooocusInpaint" + }, + "widgets_values": [ + "fooocus_inpaint_head", + "inpaint_v25 (2.58GB)" + ] + }, + { + "id": 370, + "type": "PreviewImage", + "pos": [ + 6319.49853515625, + -312.55975341796875 + ], + "size": [ + 210, + 246 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 131 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + } + ], + "links": [ + [ + 68, + 1, + 2, + 296, + 2, + "VAE" + ], + [ + 73, + 310, + 0, + 302, + 0, + "LATENT" + ], + [ + 74, + 1, + 2, + 302, + 1, + "VAE" + ], + [ + 80, + 296, + 0, + 310, + 1, + "CONDITIONING" + ], + [ + 81, + 296, + 1, + 310, + 2, + "CONDITIONING" + ], + [ + 85, + 1, + 1, + 314, + 0, + "CLIP" + ], + [ + 113, + 296, + 2, + 367, + 1, + "LATENT" + ], + [ + 124, + 367, + 0, + 310, + 0, + "MODEL" + ], + [ + 125, + 296, + 3, + 310, + 3, + "LATENT" + ], + [ + 127, + 131, + 0, + 296, + 3, + "IMAGE" + ], + [ + 129, + 131, + 1, + 296, + 4, + "MASK" + ], + [ + 130, + 1, + 0, + 367, + 0, + "MODEL" + ], + [ + 131, + 302, + 0, + 370, + 0, + "IMAGE" + ], + [ + 136, + 314, + 0, + 296, + 0, + "CONDITIONING" + ], + [ + 137, + 314, + 0, + 296, + 1, + "CONDITIONING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.672749994932563, + "offset": [ + -3563.7964746082107, + 1103.680551891574 + ] + }, + "ue_links": [] + }, + "version": 0.4 +} \ No newline at end of file diff --git a/test/test_qwenimage2prompt.json b/test/test_qwenimage2prompt.json new file mode 100644 index 0000000..56a9a1b --- /dev/null +++ b/test/test_qwenimage2prompt.json @@ -0,0 +1,526 @@ +{ + "last_node_id": 384, + "last_link_id": 143, + "nodes": [ + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [ + 4024.78759765625, + -654.6642456054688 + ], + "size": [ + 315, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 143 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 141 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 68, + 74 + ] + } + ], + "title": "Checkpoint加载器(简易)", + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "sdxl/LEOSAM AIArt 兔狲插画 SDXL大模型_v1.safetensors" + ] + }, + { + "id": 131, + "type": "LoadImage", + "pos": [ + 3479.22216796875, + -542.3040771484375 + ], + "size": [ + 315, + 314 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 127, + 138 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 129 + ], + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "1-watch.jpeg", + "image" + ] + }, + { + "id": 296, + "type": "INPAINT_VAEEncodeInpaintConditioning", + "pos": [ + 4968.38330078125, + -434.7489318847656 + ], + "size": [ + 415.8000183105469, + 106 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 142 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 140 + }, + { + "name": "vae", + "type": "VAE", + "link": 68 + }, + { + "name": "pixels", + "type": "IMAGE", + "link": 127 + }, + { + "name": "mask", + "type": "MASK", + "link": 129 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 80 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 81 + ] + }, + { + "name": "latent_inpaint", + "type": "LATENT", + "links": [], + "slot_index": 2 + }, + { + "name": "latent_samples", + "type": "LATENT", + "links": [ + 125 + ] + } + ], + "title": "VAE编码局部重绘条件", + "properties": { + "Node name for S&R": "INPAINT_VAEEncodeInpaintConditioning" + }, + "widgets_values": [] + }, + { + "id": 302, + "type": "VAEDecode", + "pos": [ + 6053.39697265625, + -312.6009216308594 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 73 + }, + { + "name": "vae", + "type": "VAE", + "link": 74 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 131 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 310, + "type": "KSampler", + "pos": [ + 5652.11767578125, + -445.38739013671875 + ], + "size": [ + 315, + 262 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 143 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 80 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 81 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 125 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 73 + ] + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 638693406916504, + "randomize", + 20, + 1, + "euler", + "normal", + 1 + ] + }, + { + "id": 370, + "type": "PreviewImage", + "pos": [ + 6319.49853515625, + -312.55975341796875 + ], + "size": [ + 210, + 246 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 131 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 379, + "type": "CLIPTextEncodeSDXLRefiner", + "pos": [ + 4458.5185546875, + -460.5770568847656 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 141 + }, + { + "name": "text", + "type": "STRING", + "widget": { + "name": "text" + }, + "link": 139 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 140, + 142 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncodeSDXLRefiner" + }, + "widgets_values": [ + 6, + 1024, + 1024, + "", + [ + false, + true + ] + ] + }, + { + "id": 380, + "type": "LayerUtility: QWenImage2Prompt", + "pos": [ + 3883.260009765625, + -326.6363830566406 + ], + "size": [ + 491.4000244140625, + 58 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 138 + } + ], + "outputs": [ + { + "name": "text", + "type": "STRING", + "links": [ + 139 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LayerUtility: QWenImage2Prompt" + }, + "widgets_values": [ + "describe this image" + ], + "color": "rgba(38, 73, 116, 0.7)" + } + ], + "links": [ + [ + 68, + 1, + 2, + 296, + 2, + "VAE" + ], + [ + 73, + 310, + 0, + 302, + 0, + "LATENT" + ], + [ + 74, + 1, + 2, + 302, + 1, + "VAE" + ], + [ + 80, + 296, + 0, + 310, + 1, + "CONDITIONING" + ], + [ + 81, + 296, + 1, + 310, + 2, + "CONDITIONING" + ], + [ + 125, + 296, + 3, + 310, + 3, + "LATENT" + ], + [ + 127, + 131, + 0, + 296, + 3, + "IMAGE" + ], + [ + 129, + 131, + 1, + 296, + 4, + "MASK" + ], + [ + 131, + 302, + 0, + 370, + 0, + "IMAGE" + ], + [ + 138, + 131, + 0, + 380, + 0, + "IMAGE" + ], + [ + 139, + 380, + 0, + 379, + 1, + "STRING" + ], + [ + 140, + 379, + 0, + 296, + 1, + "CONDITIONING" + ], + [ + 141, + 1, + 1, + 379, + 0, + "CLIP" + ], + [ + 142, + 379, + 0, + 296, + 0, + "CONDITIONING" + ], + [ + 143, + 1, + 0, + 310, + 0, + "MODEL" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.6588450000000007, + "offset": [ + -3627.859493579418, + 897.3765880646366 + ] + }, + "ue_links": [] + }, + "version": 0.4 +} \ No newline at end of file