convert nodes_sag.py to V3 schema (#9940)

comfy_extras/nodes_sag.py

@@ -2,10 +2,13 @@ import torch
 from torch import einsum
 import torch.nn.functional as F
 import math
+from typing_extensions import override
 
 from einops import rearrange, repeat
 from comfy.ldm.modules.attention import optimized_attention
 import comfy.samplers
+from comfy_api.latest import ComfyExtension, io
+
 
 # from comfy/ldm/modules/attention.py
 # 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])
     return img
 
-class SelfAttentionGuidance:
+class SelfAttentionGuidance(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "scale": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 5.0, "step": 0.01}),
-                             "blur_sigma": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SelfAttentionGuidance",
+            display_name="Self-Attention Guidance",
+            category="_for_testing",
+            inputs=[
+                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"
-
-    def patch(self, model, scale, blur_sigma):
+    @classmethod
+    def execute(cls, model, scale, blur_sigma):
         m = model.clone()
 
         attn_scores = None
@@ -170,12 +180,16 @@ class SelfAttentionGuidance:
         # unet.mid_block.attentions[0].transformer_blocks[0].attn1.patch
         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 = {
-    "SelfAttentionGuidance": "Self-Attention Guidance",
-}
+class SagExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            SelfAttentionGuidance,
+        ]
+
+
+async def comfy_entrypoint() -> SagExtension:
+    return SagExtension()