TorchCompileVAE
Torch.compile for just the VAE, with real backend and mode controls
- vae
- VAE
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-inductororcudagraphs.inductoris 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.cudagraphsis a lighter alternative: it captures your GPU calls as a CUDA graph to cut launch overhead, and doesn't need Triton at all. Ifinductorwon't install or won't run,cudagraphsis your fallback, not a downgrade you should feel bad about.mode(defaultdefault) - one ofdefault,max-autotune,max-autotune-no-cudagraphs,reduce-overhead.max-autotunespends 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-overheadleans on CUDA graphs specifically to cut down Python-side dispatch overhead between calls.max-autotune-no-cudagraphsgives you the autotuning without the CUDA-graph piece, which is the one to try ifreduce-overheador thecudagraphsbackend misbehaves with variable or tiled shapes.fullgraph(defaultfalse) - 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 defaulttrue) - 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, turncompile_encoderoff and leavecompile_decoderon.
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.
Inputs (6)
| Name | Type | Default | Description |
|---|---|---|---|
| vae | VAE | — | |
| backend | COMBO | 2 options: inductor, cudagraphs | |
| fullgraph | BOOLEAN | false | Enable full graph mode |
| mode | COMBO | default | 4 options: default, max-autotune, max-autotune-no-cudagraphs, reduce-overhead |
| compile_encoder | BOOLEAN | true | Compile encoder |
| compile_decoder | BOOLEAN | true | Compile decoder |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| VAE | VAE | — |