xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2025-12-09 22:14:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

convert nodes_sag.py to V3 schema (#9940)

Browse Source
This commit is contained in:
Alexander Piskun 2025-09-27 00:13:52 +03:00 committed by GitHub
parent a061b06321
commit d20576e6a3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 32 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
comfy_extras/nodes_sag.py
View File

@ -2,10 +2,13 @@ import torch
from torch import einsum from torch import einsum
import torch.nn.functional as F import torch.nn.functional as F
import math import math
from typing_extensions import override
from einops import rearrange, repeat from einops import rearrange, repeat
from comfy.ldm.modules.attention import optimized_attention from comfy.ldm.modules.attention import optimized_attention
import comfy.samplers import comfy.samplers
from comfy_api.latest import ComfyExtension, io
# from comfy/ldm/modules/attention.py # from comfy/ldm/modules/attention.py
# but modified to return attention scores as well as output # but modified to return attention scores as well as output
@ -104,19 +107,26 @@ def gaussian_blur_2d(img, kernel_size, sigma):
img = F.conv2d(img, kernel2d, groups=img.shape[-3]) img = F.conv2d(img, kernel2d, groups=img.shape[-3])
return img return img
class SelfAttentionGuidance: class SelfAttentionGuidance(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return {"required": { "model": ("MODEL",), return io.Schema(
"scale": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 5.0, "step": 0.01}), node_id="SelfAttentionGuidance",
"blur_sigma": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), display_name="Self-Attention Guidance",
}} category="_for_testing",
RETURN_TYPES = ("MODEL",) inputs=[
FUNCTION = "patch" io.Model.Input("model"),
io.Float.Input("scale", default=0.5, min=-2.0, max=5.0, step=0.01),
io.Float.Input("blur_sigma", default=2.0, min=0.0, max=10.0, step=0.1),
],
outputs=[
io.Model.Output(),
],
is_experimental=True,
)
CATEGORY = "_for_testing" @classmethod
def execute(cls, model, scale, blur_sigma):
def patch(self, model, scale, blur_sigma):
m = model.clone() m = model.clone()
attn_scores = None attn_scores = None
@ -170,12 +180,16 @@ class SelfAttentionGuidance:
# unet.mid_block.attentions[0].transformer_blocks[0].attn1.patch # unet.mid_block.attentions[0].transformer_blocks[0].attn1.patch
m.set_model_attn1_replace(attn_and_record, "middle", 0, 0) m.set_model_attn1_replace(attn_and_record, "middle", 0, 0)
return (m, ) return io.NodeOutput(m)
NODE_CLASS_MAPPINGS = {
"SelfAttentionGuidance": SelfAttentionGuidance,
}
NODE_DISPLAY_NAME_MAPPINGS = { class SagExtension(ComfyExtension):
"SelfAttentionGuidance": "Self-Attention Guidance", @override
} async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
SelfAttentionGuidance,
]
async def comfy_entrypoint() -> SagExtension:
return SagExtension()