ComfyUI Node

DGX Spark VAE Loader

The small-file loader that still matters

By redstonewhite·Created 6 months ago·Updated 6 months ago· 0
DGX Spark VAE Loader
    • vae
    vae_name
    devicecuda:0

    A VAE is tiny next to a diffusion model, so the natural reaction to a "fast VAE loader" is: why bother? The answer is the same reason the rest of this pack exists - it's not about the file being small, it's about how the DGX Spark moves it. The stock loader reads safetensors through system RAM, which on this machine can transiently use up to 2× the file size in RAM and adds a delay every time you switch VAEs. DGX Spark VAE Loader does the transfer zero-copy with fastsafetensors GPUDirect, straight into unified memory.

    It's a drop-in replacement for ComfyUI's Load VAE. Same vae/ folder, same VAE output, same wiring into your VAE Decode node. If you're replacing all four loaders in a Spark workflow with this pack - checkpoint, CLIP, VAE, and the diffusion model loader - you get one consistent memory story where every piece is cached in the registry and released by the same unloader.

    How it loads

    The mechanism is the pack's standard one: SafeTensorsFileLoader copies the file directly to cuda:0 via GPUDirect, the state dict loads with assign=True so nothing gets duplicated into fresh tensors, and the VAE's patcher is pinned so its load and offload devices are both cuda:0 - ComfyUI never tries to shuffle a VAE off to CPU, which on the Spark would just be extra work for nothing.

    Then the result goes into the pack's global registry, keyed by filename. Re-run the workflow and the node hands back the already-loaded VAE from memory. The IS_CHANGED logic makes that work: it reports "always changed" until the VAE is in the registry, then switches to returning a stable load id so the cache engages. First load pays full price, everything after is instant.

    The two inputs

    • vae_name - the VAE to load, from your vae/ folder. Same dropdown as the stock Load VAE.
    • device - cuda:0, the only choice. On the Spark's unified-memory architecture there's one device and this is it. Leave it alone.

    It outputs a single vae socket that feeds your VAE Decode, exactly where the normal loader's output goes.

    Installing it

    Search ComfyUI-DGXSparkFastSafetensorsLoaders in ComfyUI Manager, or:

    cd ComfyUI/custom_nodes
    git clone https://github.com/redstonewhite/ComfyUI-DGXSparkFastSafetensorsLoaders.git
    pip install fastsafetensors   # activate your venv first if applicable
    

    Restart ComfyUI; the node lands in the loaders category. No model downloads - it reads your existing vae/ files.

    Common issues

    • The VAE never unloads. Same story as the other loaders here: memory isn't managed by ComfyUI, so your loaded VAE sits in unified memory until the pack's Model Unloader runs (or you restart). For a small VAE that's usually fine, but if you're swapping VAEs mid-session it adds up.
    • Wrong VAE, wrong decode. This node is a plain loader - it doesn't know or care about your checkpoint, it just hands back whatever file you picked. If decodes look off after switching, check vae_name, not the node.
    • You're not on a Spark. Then this is solving a problem you don't have. On a normal GPU the stock VAE loader is already fast enough that this is just another dependency in the graph.

    The honest take: this is the least glamorous node in the pack, but it's also the one that makes the others usable. The DGX Spark's whole advantage is keeping models resident in memory, and that only works if every piece of the pipeline - VAE included - loads the fast way and gets freed the controlled way.

    Categoryloaders

    Inputs (2)

    NameTypeDefaultDescription
    vae_nameCOMBOThe VAE model to load.
    deviceCOMBOcuda:0The device to load to.

    Outputs (1)

    NameTypeDescription
    vaeVAE