Apply Torch Compile

If you've been on r/comfyui long enough you've seen the Wan speed posts: "torch.compile + SageAttention, 5:30 instead of 40 minutes." Torch.compile is a genuinely real speedup for video models - the community ran on it all through the Wan 2.1/2.2 era - but wiring it up is fiddly. This node is the "apply" half of that. It takes your model, wraps its forward pass in torch.compile, and hands you back a MODEL you can wire straight back into your graph. The companion TorchCompileSpeedSettings node tells it how to compile; this one does the deed.

The hook is that it's non-intrusive. It doesn't patch Kijai's WanVideo wrapper or fight ComfyUI's own compile support - it just clones your model, compiles its forward, and returns the clone. That means it works on anything with a MODEL input, not just Wan, though Wan is clearly the target audience.

How it works

Peek at nodes.py and it's refreshingly honest. ApplyTorchCompile clones the model, then does a few things in order:

1. If the args came from "speed" mode or have experimental_ptx on, it flips torch inductor config flags directly: CUDA graphs off, max autotune on for GEMMs and pointwise kernels.
2. It sets torch._dynamo.config.cache_size_limit and recompile_limit from your args.
3. If experimental_ptx is on, it runs a small warmup - a 256×256 triton.ops.matmul, falling back to a compiled 512×512 matmul if triton.ops isn't in your build - to seed the PTX/kernel cache. That's why your next generation starts faster, not just your next step.
4. It wraps model.forward with torch.compile(backend=..., mode=..., dynamic=...) and returns the clone.

There's also a weakref cache keyed on the model object plus the compile config, so repeated calls for the same model + settings skip a recompile. Two honest caveats from reading the source: compile_transformer_blocks_only is carried in the args dict but never actually read here, so don't trust that knob to shrink compile scope yet. And because every graph run clones a fresh model, the node's own cache rarely hits across separate workflow executions - the repeat-run speedup mostly comes from torch's own compile cache, which is fine, because that's the one that actually works.

Inputs and outputs

Only three wires, all of them required:

• model (MODEL) - whatever you want compiled, from any loader.
• compile_args (WANCOMPILEARGS) - the dict from TorchCompileSpeedSettings. This custom type is what lets you also plug straight into the WanVideo Cython Model Loader's compile_args socket instead of this node.
• model (MODEL) out - the compiled clone. Wire it where the original went.

Installing

No models, no pip deps, no Triton requirement in the repo - it's pure Python glue around torch that's already in your ComfyUI env. Install via ComfyUI Manager (search "ComfyUI-TorchCompileSpeed") or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/eddyhhlure1Eddy/ComfyUI-TorchCompileSpeed

Restart ComfyUI. That's the whole install.

Troubleshooting

• First run is slow, second run is fast. That's compilation + autotune, not a bug. The README's benchmarks on an RTX 5090 show the gap: ~10s first run, ~0ms second.
• Console says "Compiled and cached forward" then "Reused compiled forward from cache." The first is the compile happening; if you see "ERROR: Compilation failed" the node safely returns your original uncompiled model - you keep generating, just slower.
• OOM or VRAM pressure: the README suggests lowering dynamo_cache_size_limit in the Settings node. Compiling the whole model eats memory while it autotunes.
• No triton.ops: the PTX warmup falls back to a plain compiled matmul and prints a [TorchCompileSpeed] line so you know. On Windows, Triton may be missing entirely; the fallback still works, just less aggressively.

If you're only doing Wan, honestly, the cleanest setup is Settings → WanVideo Cython Model Loader and skip this node entirely. Reach for ApplyTorchCompile when you want the same settings on a model the loader won't compile for you.