convert nodes_torch_compile.py to V3 schema (#10173)

Author: bigcat88

comfy_extras/nodes_torch_compile.py

@@ -1,23 +1,39 @@
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, io
 from comfy_api.torch_helpers import set_torch_compile_wrapper
 
 
-class TorchCompileModel:
+class TorchCompileModel(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "backend": (["inductor", "cudagraphs"],),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="TorchCompileModel",
+            category="_for_testing",
+            inputs=[
+                io.Model.Input("model"),
+                io.Combo.Input(
+                    "backend",
+                    options=["inductor", "cudagraphs"],
+                ),
+            ],
+            outputs=[io.Model.Output()],
+            is_experimental=True,
+        )
 
-    CATEGORY = "_for_testing"
-    EXPERIMENTAL = True
-
-    def patch(self, model, backend):
+    @classmethod
+    def execute(cls, model, backend) -> io.NodeOutput:
         m = model.clone()
         set_torch_compile_wrapper(model=m, backend=backend)
-        return (m, )
+        return io.NodeOutput(m)
+
+
+class TorchCompileExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            TorchCompileModel,
+        ]
+
 
-NODE_CLASS_MAPPINGS = {
-    "TorchCompileModel": TorchCompileModel,
-}
+async def comfy_entrypoint() -> TorchCompileExtension:
+    return TorchCompileExtension()