nodes/infinigence-ComfyUI_Model_Cache
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat]: update Readme, refactor code and add AutoModel hijack (#1)

Browse Source 
* [doc]: fix readme

* [feat]: support torch.nn.Module cache and add test wf

* [refactor]: optimize directory structure
This commit is contained in:
Zhao Yingzhuo
2025-03-21 18:19:23 +08:00
committed by GitHub
parent d001f233ca
commit ec199c2a4f
7 changed files with 1100 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+9 -1
View File
@@ -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
```
```
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.
hijack/cache_manager.py
+13 -2
View File
@@ -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
hijack/hijack_list.py
+25 -1
View File
@@ -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)
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])
nodes.py
-2
View File
@@ -1,5 +1,3 @@
import torch
NODE_CLASS_MAPPINGS = {
}
prestartup_script.py
+5 -7
View File
@@ -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)
hijack_all()
test/test_apply_fooocus.json
+522
View File
@@ -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
}
test/test_qwenimage2prompt.json
+526
View File
@@ -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
}