VAE Loader (Unified Memory)

The VAE is the last mile of every image you generate: it's what turns the sampler's latent noise into actual pixels at decode time. It's also small enough that nobody benchmarks it, which is exactly why VAE Loader (Unified Memory) deserves a sentence of honesty before the details: like the pack's upscaler node, this isn't about shaving minutes off a load. It's about running your entire model set - UNET, text encoder, and decoder - through one unified-memory-aware loading path on a DGX Spark or GB10, so the VAE stays warm on the GPU instead of getting re-staged through CPU between runs.

What it does

It loads a VAE from vae/ and outputs a VAE, which you wire into VAEDecode (or VAEEncode for img2img). With dgx_mode ON, regular safetensors VAE files take the direct-to-CUDA path; everything else falls back to stock loading. The nice part is the file list itself: the vae_name dropdown mirrors native ComfyUI 0.24, which means beyond your actual .safetensors files you get the built-in pseudo entries - pixel_space (the "just pass pixels through" VAE used in some new workflows) and the tiny approximate VAEs (taesd, taesdxl, taesd3, taef1, taef2 for images; taehv, lighttaew2_2, lighttaew2_1, lighttaehy1_5, taeltx_2 for video). Those are constructed from your vae_approx/ folder, exactly as the stock loader builds them, and they always use the stock path because there's nothing to accelerate.

The inputs

• vae_name - the VAE file or pseudo entry, as above.
• dgx_mode - ON (default) for the direct path, OFF for stock.
• device - cuda:0.
• storage_backend - auto (instanttensor, then fastsafetensors) or pinned.

How the fallbacks work

The source is explicit about what takes the stock route: pixel_space, the image and video TAEs, non-safetensors files, and any system without the DGX backends (Windows, non-DGX Linux). None of these error out - they just quietly load the standard way. The direct path is reserved for regular vae/ safetensors, which get read into CUDA, built with the pack's assign-style construction, and registered as GPU-resident so ComfyUI doesn't decide to evict the decoder right before you need it.

Installing it

It ships in the DGX Nodes pack with six siblings - install once:

cd ComfyUI/custom_nodes
git clone https://github.com/broken-gage/ComfyUI-DGX-Nodes
pip install -r requirements.txt
pip install instanttensor fastsafetensors   # optional backends for the DGX path

Or via ComfyUI Manager, searching DGX Nodes. Restart ComfyUI and find the nodes under DGX Nodes. Requires ComfyUI 0.24.0+.

Gotchas

• The pseudo entries are always stock. Don't read "pixel_space" in the dropdown and assume it's accelerated - it's a native entry that ComfyUI handles itself. The DGX path only touches regular safetensors VAEs.
• Silent fallbacks are the pack's theme. On a normal x86 or Windows box this node is just VAELoader with extra steps. Harmless to keep in shared workflows, but it's only earning its keep on a Spark/GB10 with the backends installed.
• One VAE-specific reminder that predates this pack: for most models the bundled VAE is fine, but if you see washed-out or muddy colors, the first suspect is a mismatched VAE - and a loader that quietly picks the wrong file is worse than one that tells you. This node at least shows you exactly which entry you selected.
• Same pack caveats, said once: experimental, GPLv3, AI-assisted vibe-code with no guaranteed maintenance. Low stakes here - but the pattern is worth knowing before you trust the big UNET loader on the same pack.