Nodes/KJNodes for ComfyUI/TorchCompileVAE
ComfyUI Node Runs on cloud

TorchCompileVAE

Torch.compile for just the VAE, with real backend and mode controls

By kijai·Created 3 years ago·Updated about 12 hours ago· 2,930
TorchCompileVAE
  • vae
  • VAE
backend
fullgraphfalse
modedefault
compile_encodertrue
compile_decodertrue

Unlike its three deprecated siblings in this pack (the LTX, Wan Video, and Qwen-Image TorchCompileModel* nodes, all of which now just say "use TorchCompileModelAdvanced instead"), this one is current - and it's the node that actually shows you what a properly configured torch.compile wrapper looks like in KJNodes. Instead of compiling the whole diffusion model, it compiles the VAE specifically: the part of the pipeline that turns latents into pixels (and back). For video especially, VAE decode is often a real chunk of your total run time, since you're decoding dozens of frames instead of one image - which is exactly the repeat-call, fixed-shape workload torch.compile is built to speed up.

How it works

torch.compile traces a model's forward pass once and produces a faster compiled version, at the cost of a slower first call while that trace happens. Everything past that first run, at the same shape, is where you get the speedup. This node gives you actual control over how that compile happens instead of a single fixed behavior.

The inputs that matter

  • vae (VAE, required) - the loaded VAE going in.
  • backend - inductor or cudagraphs. inductor is PyTorch's general-purpose compiler backend and does the most aggressive optimization, but it needs Triton - the same dependency that makes SageAttention such a headache to install on Windows. cudagraphs is a lighter alternative: it captures your GPU calls as a CUDA graph to cut launch overhead, and doesn't need Triton at all. If inductor won't install or won't run, cudagraphs is your fallback, not a downgrade you should feel bad about.
  • mode (default default) - one of default, max-autotune, max-autotune-no-cudagraphs, reduce-overhead. max-autotune spends real extra time at compile stage benchmarking multiple kernel variants to find the fastest one - worth it if you'll decode the same shape many times in a session, wasted if you won't. reduce-overhead leans on CUDA graphs specifically to cut down Python-side dispatch overhead between calls. max-autotune-no-cudagraphs gives you the autotuning without the CUDA-graph piece, which is the one to try if reduce-overhead or the cudagraphs backend misbehaves with variable or tiled shapes.
  • fullgraph (default false) - forces the compiler to trace your VAE as a single unbroken graph, and throws instead of silently falling back to eager mode wherever it can't. Leave it off unless you specifically need to catch graph breaks; turning it on is how you find out your VAE's forward pass has something in it torch.compile can't fully trace.
  • compile_encoder / compile_decoder (both default true) - compile either half independently. For most video workflows the decoder is doing the real work (many frames out), while the encoder often only runs once or twice per graph. If compiling the encoder causes trouble for no real speed benefit, turn compile_encoder off and leave compile_decoder on.

Output: VAE - wire it exactly where your original VAE would go, into VAE Decode or VAE Encode.

How to install it

Comes with the pack, no separate download.

Via ComfyUI Manager: search KJNodes for ComfyUI, install, restart.

Manually:

cd ComfyUI/custom_nodes
git clone https://github.com/kijai/ComfyUI-KJNodes
pip install -r ComfyUI-KJNodes/requirements.txt

Then restart. Nothing extra to fetch - it compiles a VAE you already have.

Common issues & troubleshooting

Your first decode after adding this is much slower. Expected - that's the compile trace happening once. If your resolution or batch size changes every run, you're paying that cost repeatedly instead of amortizing it, and you may end up net slower overall.

inductor fails to compile, especially on Windows. That backend needs Triton, which is a genuinely rough install on Windows - bad enough that Windows Triton builds are largely a one-person community effort outside official PyTorch releases. Switch backend to cudagraphs and you sidestep the dependency entirely.

Tiled VAE decode misbehaves or errors with fullgraph on, or with an aggressive mode. Tiled decoding at higher resolutions produces a variable number of tiles per call, and strict single-graph tracing or the cudagraphs optimization path doesn't always tolerate that well. Turn fullgraph off first; if it's still unstable, try max-autotune-no-cudagraphs before giving up on compiling the decoder entirely.

CategoryKJNodes/torchcompile

Inputs (6)

NameTypeDefaultDescription
vaeVAE
backendCOMBO2 options: inductor, cudagraphs
fullgraphBOOLEANfalseEnable full graph mode
modeCOMBOdefault4 options: default, max-autotune, max-autotune-no-cudagraphs, reduce-overhead
compile_encoderBOOLEANtrueCompile encoder
compile_decoderBOOLEANtrueCompile decoder

Outputs (1)

NameTypeDescription
VAEVAE