From ac10a0d69e9905662296c5280bcea61945c39762 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Sat, 26 Apr 2025 16:56:22 -0700
Subject: [PATCH 1/2] Make loras work with --async-offload (#7824)

---
 comfy/ops.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/comfy/ops.py b/comfy/ops.py
index 62daf447b..032787915 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -22,6 +22,7 @@ import comfy.model_management
 from comfy.cli_args import args, PerformanceFeature
 import comfy.float
 import comfy.rmsnorm
+import contextlib
 
 cast_to = comfy.model_management.cast_to #TODO: remove once no more references
 
@@ -38,20 +39,28 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
             device = input.device
 
     offload_stream = comfy.model_management.get_offload_stream(device)
+    if offload_stream is not None:
+        wf_context = offload_stream
+    else:
+        wf_context = contextlib.nullcontext()
+
     bias = None
     non_blocking = comfy.model_management.device_supports_non_blocking(device)
     if s.bias is not None:
         has_function = len(s.bias_function) > 0
         bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function, stream=offload_stream)
+
         if has_function:
-            for f in s.bias_function:
-                bias = f(bias)
+            with wf_context:
+                for f in s.bias_function:
+                    bias = f(bias)
 
     has_function = len(s.weight_function) > 0
     weight = comfy.model_management.cast_to(s.weight, dtype, device, non_blocking=non_blocking, copy=has_function, stream=offload_stream)
     if has_function:
-        for f in s.weight_function:
-            weight = f(weight)
+        with wf_context:
+            for f in s.weight_function:
+                weight = f(weight)
 
     comfy.model_management.sync_stream(device, offload_stream)
     return weight, bias

From 542b4b36b694148504656ad54433b8ddf0c38c4d Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Sat, 26 Apr 2025 17:52:56 -0700
Subject: [PATCH 2/2] Prevent custom nodes from hooking certain functions.
 (#7825)

---
 hook_breaker_ac10a0.py | 17 +++++++++++++++++
 main.py                |  5 ++++-
 2 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 hook_breaker_ac10a0.py

diff --git a/hook_breaker_ac10a0.py b/hook_breaker_ac10a0.py
new file mode 100644
index 000000000..c3e1c0633
--- /dev/null
+++ b/hook_breaker_ac10a0.py
@@ -0,0 +1,17 @@
+# Prevent custom nodes from hooking anything important
+import comfy.model_management
+
+HOOK_BREAK = [(comfy.model_management, "cast_to")]
+
+
+SAVED_FUNCTIONS = []
+
+
+def save_functions():
+    for f in HOOK_BREAK:
+        SAVED_FUNCTIONS.append((f[0], f[1], getattr(f[0], f[1])))
+
+
+def restore_functions():
+    for f in SAVED_FUNCTIONS:
+        setattr(f[0], f[1], f[2])
diff --git a/main.py b/main.py
index ac9d24b7b..f3f56597a 100644
--- a/main.py
+++ b/main.py
@@ -141,7 +141,7 @@ import nodes
 import comfy.model_management
 import comfyui_version
 import app.logger
-
+import hook_breaker_ac10a0
 
 def cuda_malloc_warning():
     device = comfy.model_management.get_torch_device()
@@ -215,6 +215,7 @@ def prompt_worker(q, server_instance):
                 comfy.model_management.soft_empty_cache()
                 last_gc_collect = current_time
                 need_gc = False
+                hook_breaker_ac10a0.restore_functions()
 
 
 async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None):
@@ -268,7 +269,9 @@ def start_comfyui(asyncio_loop=None):
     prompt_server = server.PromptServer(asyncio_loop)
     q = execution.PromptQueue(prompt_server)
 
+    hook_breaker_ac10a0.save_functions()
     nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)
+    hook_breaker_ac10a0.restore_functions()
 
     cuda_malloc_warning()