xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2026-04-10 16:46:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

smarter sage patch

Browse Source
This commit is contained in:
kijai 2025-05-07 19:08:25 +03:00
parent ca07b9dadc
commit bec42252c6
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
nodes/model_optimization_nodes.py
View File

@ -92,6 +92,7 @@ class BaseLoaderKJ:
comfy.ldm.wan.model.optimized_attention = attention_sage
else:
print("Restoring initial comfy attention")
comfy_attention.optimized_attention = _original_functions.get("orig_attention")
comfy.ldm.hunyuan_video.model.optimized_attention = _original_functions.get("orig_attention")
comfy.ldm.flux.math.optimized_attention = _original_functions.get("orig_attention")
@ -145,9 +146,13 @@ class PathchSageAttentionKJ(BaseLoaderKJ):
def patch(self, model, sage_attention):
model_clone = model.clone()
def patch_attention(model):
def patch_attention_enable(model):
self._patch_modules(False, sage_attention)
model_clone.add_callback(CallbacksMP.ON_PRE_RUN,patch_attention)
def patch_attention_disable(model):
self._patch_modules(False, "disabled")
model_clone.add_callback(CallbacksMP.ON_PRE_RUN, patch_attention_enable)
model_clone.add_callback(CallbacksMP.ON_CLEANUP, patch_attention_disable)
return model_clone,