From 3e68bc342cd60b909b4117c1b68a3afc62ef875c Mon Sep 17 00:00:00 2001
From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com>
Date: Fri, 3 Oct 2025 21:43:54 +0300
Subject: [PATCH] convert nodes_torch_compile.py to V3 schema (#10173)

---
 comfy_extras/nodes_torch_compile.py | 46 +++++++++++++++++++----------
 1 file changed, 31 insertions(+), 15 deletions(-)

diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py
index 605536678..adbeece2f 100644
--- a/comfy_extras/nodes_torch_compile.py
+++ b/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)
 
-NODE_CLASS_MAPPINGS = {
-    "TorchCompileModel": TorchCompileModel,
-}
+
+class TorchCompileExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            TorchCompileModel,
+        ]
+
+
+async def comfy_entrypoint() -> TorchCompileExtension:
+    return TorchCompileExtension()