convert nodes_differential_diffusion.py to V3 schema (#10056)

comfy_extras/nodes_differential_diffusion.py

@@ -1,34 +1,41 @@
 # code adapted from https://github.com/exx8/differential-diffusion
 
+from typing_extensions import override
+
 import torch
+from comfy_api.latest import ComfyExtension, io
 
-class DifferentialDiffusion():
+
+class DifferentialDiffusion(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {
+        return io.Schema(
-            "required": {
+            node_id="DifferentialDiffusion",
-                "model": ("MODEL", ),
+            display_name="Differential Diffusion",
-            },
+            category="_for_testing",
-            "optional": {
+            inputs=[
-                "strength": ("FLOAT", {
+                io.Model.Input("model"),
-                    "default": 1.0,
+                io.Float.Input(
-                    "min": 0.0,
+                    "strength",
-                    "max": 1.0,
+                    default=1.0,
-                    "step": 0.01,
+                    min=0.0,
-                }),
+                    max=1.0,
-            }
+                    step=0.01,
-        }
+                    optional=True,
-    RETURN_TYPES = ("MODEL",)
+                ),
-    FUNCTION = "apply"
+            ],
-    CATEGORY = "_for_testing"
+            outputs=[io.Model.Output()],
-    INIT = False
+            is_experimental=True,
+        )
 
-    def apply(self, model, strength=1.0):
+    @classmethod
+    def execute(cls, model, strength=1.0) -> io.NodeOutput:
         model = model.clone()
-        model.set_model_denoise_mask_function(lambda *args, **kwargs: self.forward(*args, **kwargs, strength=strength))
+        model.set_model_denoise_mask_function(lambda *args, **kwargs: cls.forward(*args, **kwargs, strength=strength))
-        return (model, )
+        return io.NodeOutput(model)
 
-    def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict, strength: float):
+    @classmethod
+    def forward(cls, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict, strength: float):
         model = extra_options["model"]
         step_sigmas = extra_options["sigmas"]
         sigma_to = model.inner_model.model_sampling.sigma_min
@@ -53,9 +60,13 @@ class DifferentialDiffusion():
             return binary_mask
 
 
-NODE_CLASS_MAPPINGS = {
+class DifferentialDiffusionExtension(ComfyExtension):
-    "DifferentialDiffusion": DifferentialDiffusion,
+    @override
-}
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-NODE_DISPLAY_NAME_MAPPINGS = {
+        return [
-    "DifferentialDiffusion": "Differential Diffusion",
+            DifferentialDiffusion,
-}
+        ]
+
+
+async def comfy_entrypoint() -> DifferentialDiffusionExtension:
+    return DifferentialDiffusionExtension()