convert nodes_eps.py to V3 schema (comfyanonymous#10172)

Author: bigcat88

comfy_extras/nodes_eps.py

@@ -1,4 +1,9 @@
-class EpsilonScaling:
+from typing_extensions import override
+
+from comfy_api.latest import ComfyExtension, io
+
+
+class EpsilonScaling(io.ComfyNode):
     """
     Implements the Epsilon Scaling method from 'Elucidating the Exposure Bias in Diffusion Models'
     (https://arxiv.org/abs/2308.15321v6).
@@ -8,26 +13,28 @@ class EpsilonScaling:
     recommended by the paper for its practicality and effectiveness.
     """
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "scaling_factor": ("FLOAT", {
-                    "default": 1.005,
-                    "min": 0.5,
-                    "max": 1.5,
-                    "step": 0.001,
-                    "display": "number"
-                }),
-            }
-        }
-
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "model_patches/unet"
-
-    def patch(self, model, scaling_factor):
+    def define_schema(cls):
+        return io.Schema(
+            node_id="Epsilon Scaling",
+            category="model_patches/unet",
+            inputs=[
+                io.Model.Input("model"),
+                io.Float.Input(
+                    "scaling_factor",
+                    default=1.005,
+                    min=0.5,
+                    max=1.5,
+                    step=0.001,
+                    display_mode=io.NumberDisplay.number,
+                ),
+            ],
+            outputs=[
+                io.Model.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model, scaling_factor) -> io.NodeOutput:
         # Prevent division by zero, though the UI's min value should prevent this.
         if scaling_factor == 0:
             scaling_factor = 1e-9
@@ -53,8 +60,15 @@ def epsilon_scaling_function(args):
 
         model_clone.set_model_sampler_post_cfg_function(epsilon_scaling_function)
 
-        return (model_clone,)
+        return io.NodeOutput(model_clone)
+
+
+class EpsilonScalingExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            EpsilonScaling,
+        ]
 
-NODE_CLASS_MAPPINGS = {
-    "Epsilon Scaling": EpsilonScaling
-}
+async def comfy_entrypoint() -> EpsilonScalingExtension:
+    return EpsilonScalingExtension()