Initial commit

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -0,0 +1,5 @@
+# ComfyUI-Torch-Compile
+
+ComfyUI-Torch-Compile is a set of [ComfyUI](https://www.comfy.org/) nodes for running torch.compile.
+
+These nodes are based off of the ones found in [ComfyUI-KJNodes](https://github.com/kijai/ComfyUI-KJNodes/blob/8c590fd5a023ee14b5617347567752bf62ea4cd6/nodes/model_optimization_nodes.py).

__init__.py

@@ -0,0 +1,9 @@
+from .nodes.torch_compile_vae import TorchCompileVAE
+from .nodes.torch_compile_controlnet import TorchCompileControlNet
+
+NODE_CLASS_MAPPINGS = {
+    "TorchCompileVAE": TorchCompileVAE,
+    "TorchCompileControlNet": TorchCompileControlNet,
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes/torch_compile_controlnet.py

@@ -0,0 +1,44 @@
+import torch
+
+# Based on https://github.com/kijai/ComfyUI-KJNodes/blob/8c590fd5a023ee14b5617347567752bf62ea4cd6/nodes/model_optimization_nodes.py#L359
+
+
+class TorchCompileControlNet:
+    CATEGORY = "torch-compile"
+    RETURN_TYPES = ("CONTROL_NET",)
+    FUNCTION = "compile"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "controlnet": ("CONTROL_NET",),
+                "backend": (["inductor", "cudagraphs"],),
+                "fullgraph": (
+                    "BOOLEAN",
+                    {"default": False, "tooltip": "Enable full graph mode"},
+                ),
+                "mode": (
+                    [
+                        "default",
+                        "max-autotune",
+                        "max-autotune-no-cudagraphs",
+                        "reduce-overhead",
+                    ],
+                    {"default": "default"},
+                ),
+            }
+        }
+
+    def compile(self, controlnet, backend: str, mode: str, fullgraph: bool):
+        try:
+            controlnet.control_model = torch.compile(
+                controlnet.control_model,
+                mode=mode,
+                fullgraph=fullgraph,
+                backend=backend,
+            )
+        except:
+            raise RuntimeError("Failed to compile model")
+
+        return (controlnet,)

nodes/torch_compile_vae.py

@@ -0,0 +1,86 @@
+import torch
+
+# Based on https://github.com/kijai/ComfyUI-KJNodes/blob/8c590fd5a023ee14b5617347567752bf62ea4cd6/nodes/model_optimization_nodes.py#L321
+
+
+class TorchCompileVAE:
+    CATEGORY = "torch-compile-utils"
+    RETURN_TYPES = ("VAE",)
+    FUNCTION = "compile"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "vae": ("VAE",),
+                "backend": (["inductor", "cudagraphs"],),
+                "fullgraph": (
+                    "BOOLEAN",
+                    {"default": False, "tooltip": "Enable full graph mode"},
+                ),
+                "mode": (
+                    [
+                        "default",
+                        "max-autotune",
+                        "max-autotune-no-cudagraphs",
+                        "reduce-overhead",
+                    ],
+                    {"default": "default"},
+                ),
+                "compile_encoder": (
+                    "BOOLEAN",
+                    {"default": True, "tooltip": "Compile encoder"},
+                ),
+                "compile_decoder": (
+                    "BOOLEAN",
+                    {"default": True, "tooltip": "Compile decoder"},
+                ),
+            }
+        }
+
+    def compile(
+        self,
+        vae,
+        backend: str,
+        mode: str,
+        fullgraph: bool,
+        compile_encoder: bool,
+        compile_decoder: bool,
+    ):
+        if compile_encoder:
+            encoder_name = "encoder"
+            if hasattr(vae.first_stage_model, "taesd_encoder"):
+                encoder_name = "taesd_encoder"
+
+            try:
+                setattr(
+                    vae.first_stage_model,
+                    encoder_name,
+                    torch.compile(
+                        getattr(vae.first_stage_model, encoder_name),
+                        mode=mode,
+                        fullgraph=fullgraph,
+                        backend=backend,
+                    ),
+                )
+            except:
+                raise RuntimeError("Failed to compile model")
+        if compile_decoder:
+            decoder_name = "decoder"
+            if hasattr(vae.first_stage_model, "taesd_decoder"):
+                decoder_name = "taesd_decoder"
+
+            try:
+                setattr(
+                    vae.first_stage_model,
+                    decoder_name,
+                    torch.compile(
+                        getattr(vae.first_stage_model, decoder_name),
+                        mode=mode,
+                        fullgraph=fullgraph,
+                        backend=backend,
+                    ),
+                )
+            except:
+                raise RuntimeError("Failed to compile model")
+        return (vae,)