Skip to content

Commit 542b4b3

Browse files
Prevent custom nodes from hooking certain functions. (Comfy-Org#7825)
1 parent ac10a0d commit 542b4b3

2 files changed

Lines changed: 21 additions & 1 deletion

File tree

hook_breaker_ac10a0.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# Prevent custom nodes from hooking anything important
2+
import comfy.model_management
3+
4+
HOOK_BREAK = [(comfy.model_management, "cast_to")]
5+
6+
7+
SAVED_FUNCTIONS = []
8+
9+
10+
def save_functions():
11+
for f in HOOK_BREAK:
12+
SAVED_FUNCTIONS.append((f[0], f[1], getattr(f[0], f[1])))
13+
14+
15+
def restore_functions():
16+
for f in SAVED_FUNCTIONS:
17+
setattr(f[0], f[1], f[2])

main.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ def execute_script(script_path):
141141
import comfy.model_management
142142
import comfyui_version
143143
import app.logger
144-
144+
import hook_breaker_ac10a0
145145

146146
def cuda_malloc_warning():
147147
device = comfy.model_management.get_torch_device()
@@ -215,6 +215,7 @@ def prompt_worker(q, server_instance):
215215
comfy.model_management.soft_empty_cache()
216216
last_gc_collect = current_time
217217
need_gc = False
218+
hook_breaker_ac10a0.restore_functions()
218219

219220

220221
async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None):
@@ -268,7 +269,9 @@ def start_comfyui(asyncio_loop=None):
268269
prompt_server = server.PromptServer(asyncio_loop)
269270
q = execution.PromptQueue(prompt_server)
270271

272+
hook_breaker_ac10a0.save_functions()
271273
nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)
274+
hook_breaker_ac10a0.restore_functions()
272275

273276
cuda_malloc_warning()
274277

0 commit comments

Comments
 (0)