mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-09 19:17:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove AttentionOverrideTest node, that's something to cook up for later

Browse Source
This commit is contained in:
Jedrzej Kosinski
2025-08-30 15:19:36 -07:00
parent dd0a5093f6
commit 66c4eb006b
1 changed files with 0 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
comfy_extras/nodes_wan.py
View File

@@ -1058,51 +1058,6 @@ class Wan22ImageToVideoLatent(io.ComfyNode):
out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))
return io.NodeOutput(out_latent)
import comfy.patcher_extension
import comfy.ldm.modules.attention
import logging
class AttentionOverrideTest(io.ComfyNode):
@classmethod
def define_schema(cls):
attention_function_names = list(comfy.ldm.modules.attention.REGISTERED_ATTENTION_FUNCTIONS.keys())
return io.Schema(
node_id="AttentionOverrideTest",
category="devtools",
inputs=[
io.Model.Input("model"),
io.Combo.Input("attention", options=attention_function_names),
],
outputs=[
io.Model.Output(),
],
)
@staticmethod
def attention_override_factory(attention_func):
def attention_override(func, *args, **kwargs):
return attention_func(*args, **kwargs)
return attention_override
@staticmethod
def sampler_sampler_wrapper(executor, *args, **kwargs):
try:
# extra_args = args[2]
return executor(*args, **kwargs)
finally:
pass
@classmethod
def execute(cls, model: io.Model.Type, attention: str) -> io.NodeOutput:
attention_func = comfy.ldm.modules.attention.get_attention_function(attention, None)
if attention_func is None:
logging.info(f"Attention type '{attention}' not found, using default optimized attention for your hardware.")
return model
model = model.clone()
model.model_options["transformer_options"]["optimized_attention_override"] = cls.attention_override_factory(attention_func)
model.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE, "attention_override_test", cls.sampler_sampler_wrapper)
return io.NodeOutput(model)
class WanExtension(ComfyExtension):
@override
@@ -1121,7 +1076,6 @@ class WanExtension(ComfyExtension):
WanSoundImageToVideo,
WanSoundImageToVideoExtend,
Wan22ImageToVideoLatent,
AttentionOverrideTest,
]
async def comfy_entrypoint() -> WanExtension: