diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ed8ebf5 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..a448e41 --- /dev/null +++ b/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). \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6c04118 --- /dev/null +++ b/__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"] diff --git a/nodes/torch_compile_controlnet.py b/nodes/torch_compile_controlnet.py new file mode 100644 index 0000000..83dc9cd --- /dev/null +++ b/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,) diff --git a/nodes/torch_compile_vae.py b/nodes/torch_compile_vae.py new file mode 100644 index 0000000..66dccdc --- /dev/null +++ b/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,)