Skip to content

Commit

Permalink
Add backend option to TorchCompileModel.
Browse files Browse the repository at this point in the history
If you want to use the cudagraphs backend you need to: --disable-cuda-malloc

If you get other backends working feel free to make a PR to add them.
  • Loading branch information
comfyanonymous committed Sep 27, 2024
1 parent d72e871 commit 83b01f9
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions comfy_extras/nodes_torch_compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,17 @@ class TorchCompileModel:
@classmethod
def INPUT_TYPES(s):
return {"required": { "model": ("MODEL",),
"backend": (["inductor", "cudagraphs"],),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"

CATEGORY = "_for_testing"
EXPERIMENTAL = True

def patch(self, model):
def patch(self, model, backend):
m = model.clone()
m.add_object_patch("diffusion_model", torch.compile(model=m.get_model_object("diffusion_model")))
m.add_object_patch("diffusion_model", torch.compile(model=m.get_model_object("diffusion_model"), backend=backend))
return (m, )

NODE_CLASS_MAPPINGS = {
Expand Down

0 comments on commit 83b01f9

Please sign in to comment.