Nodes/ComfyUI-Torch-Compile/TorchCompileLoadVAE
ComfyUI Node

TorchCompileLoadVAE

Your VAE is a hidden speed limit — TorchCompileLoadVAE compiles it away

By yondonfu·Created 2 years ago·Updated about a year ago· 11
TorchCompileLoadVAE
  • vae
  • VAE
backend
fullgraphfalse
modedefault
compile_encodertrue
compile_decodertrue

Everyone obsesses over the sampler, but the VAE is where a lot of your wall-clock actually goes. Decode a 1024px image and it's one pass; decode a 30-second video frame by frame, or a big batch, and the VAE quietly becomes the bottleneck. That's the niche TorchCompileLoadVAE exists to fix. It takes a VAE you already loaded, wraps its encoder and decoder in torch.compile, and hands you back the same VAE - faster on the second call onward.

It's a tiny node, and honestly, that's the point. If you generate one image at a time and don't care about a few hundred milliseconds, skip it. If you batch, loop, upscale a lot, or push video through ComfyUI, this is one of the few genuinely free speedups left.

How it works

TorchCompileLoadVAE is a pass-through that mutates the loaded VAE in place. Looking at the source, it grabs vae.first_stage_model, takes the modules named encoder and decoder, and replaces them with torch.compile()-wrapped versions. The compile_encoder and compile_decoder toggles let you opt each half in or out, and it's smart about naming: if your VAE is a TAESD, it compiles taesd_encoder/taesd_decoder instead. That TAESD detection is the actual reason this pack exists - the author forked these nodes out of Kijai's KJNodes because torch.compile silently failed on TAESD VAEs there, and fixed it.

The compile is JIT, so the first decode after wiring this in is slower - torch builds and optimizes the graph, which can take tens of seconds depending on your mode. Every decode after that uses the cached compiled graph. That's why this node shines for repeated work: a single one-shot generation pays the compile tax and gets almost nothing back; batch and video work amortize it instantly.

The inputs that matter

Only one of these is a model: vae accepts the output of your normal Load VAE (or a VAE baked into a checkpoint). The rest are tuning knobs you can mostly ignore on day one:

  • backend - inductor (default, the safe general choice) or cudagraphs (CUDA graph capture, faster but pickier about your GPU and dynamic shapes).
  • mode - default is fine to start. max-autotune hunts for the fastest kernel config at the cost of a brutal first-compile time; reduce-overhead is the middle ground; max-autotune-no-cudagraphs is what the community tends to land on when cudagraphs misbehave.
  • fullgraph - off by default. Turning it on compiles the whole module as one graph, which can help, but it will error out if the model has any Python control flow torch can't fuse.

The output is one VAE, which you wire straight into your existing VAE Decode / VAE Encode nodes exactly as you would the uncompiled one.

Installing it

Install the pack once and both nodes come with it. Easiest path is ComfyUI Manager - search "ComfyUI-Torch-Compile" and hit install, then restart. Or do it by hand:

cd ComfyUI/custom_nodes
git clone https://github.com/yondonfu/ComfyUI-Torch-Compile
cd ComfyUI-Torch-Compile
pip install -r requirements.txt

There's nothing heavy here - the requirements file is literally just torch, and there are no model downloads. The only real dependency is having a working torch.compile in your PyTorch build.

Common issues

  • Nothing gets faster. Check your logs - if you see compile errors, you probably lack triton, which inductor needs on NVIDIA/Linux. It's not in the requirements file, so install it yourself (pip install triton), and note it can be fiddly on Windows.
  • cudagraphs backend crashes or hangs. Very common on some cards (30-series especially). Drop to inductor, or use max-autotune-no-cudagraphs.
  • max-autotune won't even start on your GPU. It requires a card with 80+ SMs - that's roughly RTX 30/40-class high end. A 4060-class card won't qualify; default still works fine.
  • First decode seems to freeze. That's the JIT compile, not a crash. Leave it running; it's one-time per graph.

The author is Yondon Fu, the engineer behind ComfyStream (Livepeer's realtime ComfyUI streaming project) - so the "compile the expensive parts, amortize over many calls" philosophy here comes from someone whose whole job is shaving milliseconds off repeated inference. The nodes are MIT-licensed, two files, no bloat. Worth having in your toolbox even if you only reach for it on video days.

Categorytorch-compile

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