VAE (Improved)

AugmentWaitVAE, shown in the node menu as "VAE (Improved)", is the trigger-aware passthrough for the VAE type in Augment's Flow Control family. It exists for one reason: to pin when and where a VAE moves through your graph, and to stop ComfyUI's output cache from quietly skipping stages that share a value. If you've built anything multi-stage - encode, edit, decode, save - you know the pain it's aimed at.

The model itself is where most VAE trouble lives: mismatched VAEs are a classic source of washed-out or "broken" outputs, and swapping one mid-pipeline is exactly the kind of edit that ComfyUI's caching likes to pretend never happened. Putting the VAE through a Wait node makes the swap and its downstream decode stages re-run in a deterministic order every time.

How it works

Same engine as every node in this family. The node doesn't literally hold the VAE until a trigger arrives - reading the source, execute passes your value straight through and spits out a trigger token. The sequencing magic lives in the pack's executor patch: on import, Augment rewrites parts of ComfyUI's execution internals (execution.execute and the cache's change-detection) so any node that touches a TRIGGER type always re-runs and is never pruned by the "inputs unchanged, skip it" logic. Chain the trigger outputs into the next stage's trigger input and you've forced a real order onto a graph that otherwise wants to run everything at once.

In plain terms: it's a "make sure this runs here, in this order" sticker, not a gatekeeper.

The inputs and outputs

• value (VAE, required) - the VAE to carry. Wire it from your checkpoint loader or a dedicated VAE loader.
• trigger (optional) - the sequencing input. Connect the previous stage's trigger output here so this node runs after it.
• Outputs value (VAE), json (AUGMENT_JSON), trigger (TRIGGER) - VAE passes through untouched, the trigger passes on, and json carries a small structured dump (node type and value as a string) that you can feed to Augment's JSON Extract or JSON Viewer.

There's no widget on this one - the VAE always comes in on the wire, so there's nothing to set on the node itself. Its whole job is ordering and re-execution.

How to install

Augment isn't on the ComfyUI registry yet, so clone it by hand - and mind the folder name:

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

Restart ComfyUI and the nodes appear under Augment. The clone target must be augment: the pack imports everything as custom_nodes.augment.py.*, so a folder named augment-ComfyUI breaks all node loading. That's the single most common "I installed it and nothing shows up" cause with this pack.

Dependencies are just opencv-python, numpy, torch, Pillow and requests - all of which you already have from running ComfyUI. No model downloads, and no API key: only the pack's vector (PNG-to-SVG / SVG-to-PNG) nodes talk to the Augment Studio API.

Common issues

• VAE passthrough "lost" or decode stage skipped - that's the cache doing its thing; inserting a Wait node in the VAE chain is the fix, since trigger nodes always re-run.
• Node missing after install - folder name, almost certainly. Rename to augment and restart.
• Odd execution errors after a ComfyUI update - the executor override patches internal APIs that shift between versions. If the console shows an [augment] warning about the executor patch, that's your culprit; update the pack or pin your ComfyUI version.