Nodes/ComfyUI-MultiGPU/VAELoaderDisTorch2MultiGPU
ComfyUI Node

VAELoaderDisTorch2MultiGPU

Keep the VAE decode off your busy GPU

By pollockjj·Created 2 years ago·Updated 3 months ago· 925
VAELoaderDisTorch2MultiGPU
    • VAE
    vae_name
    compute_devicecpu
    virtual_vram_gb4.0
    donor_devicecpu
    expert_mode_allocations
    eject_modelstrue

    The VAE is the smallest of the three big model pieces (checkpoint/UNet, text encoder, VAE) - usually a few hundred MB - so on its own it's rarely what's blowing your VRAM budget. VAELoaderDisTorch2MultiGPU exists anyway because it's part of the same family as the pack's UNet and CLIP DisTorch2 loaders: it lets you pick which device does the encode/decode and, if you want, push part of it off your compute card. On a single-GPU box that mostly means "run the VAE decode on CPU so it doesn't compete with the sampler." On a multi-GPU box it means the VAE can live on your second card entirely, out of the way.

    What it actually does

    Same underlying idea as the rest of the DisTorch2 family covered in this pack: it doesn't parallelize anything, it moves where the model's weights sit so your primary compute device has more headroom for the thing that actually needs it - latent space during sampling. The VAE only runs at the encode (start of img2img) and decode (end of every run) steps, so unlike the UNet it isn't touched every sampling step - which is exactly why offloading it costs you almost nothing in speed.

    Inputs and output

    • vae_name - pick your VAE file from ComfyUI/models/vae. One entry worth knowing about: pixel_space. Some current models (HiDream-O1 being the clearest example) don't use a VAE at all - they generate and edit raw pixels directly. pixel_space is the placeholder you select for those; picking a real VAE file for a model that has none (or the reverse) is a documented source of "why is this broken" confusion in ComfyUI workflows.
    • compute_device (default cpu) - which device runs the actual encode/decode math.
    • virtual_vram_gb (default 4, 0–128) - how much of the VAE's footprint to offload to the donor.
    • donor_device (default cpu) - where the offloaded part lives.
    • expert_mode_allocations - the advanced byte/ratio/fraction override string, same syntax as the pack's other DisTorch2 nodes. Leave it blank unless you're deliberately splitting across three-plus devices.
    • eject_models - defaults to true here, unlike the CLIP loader's default-off. Worth noticing: the VAE decode is usually the last big memory-hungry step of a run, so clearing other models off the target device first before it runs is the more sensible default.

    Output is a single VAE - the same type the stock VAELoader produces, so it drops straight into whatever VAEDecode / VAEEncode nodes you already have.

    Installing it

    ComfyUI Manager is the easy way - search "ComfyUI-MultiGPU," install, restart. Manual install:

    cd ComfyUI/custom_nodes
    git clone https://github.com/pollockjj/ComfyUI-MultiGPU
    

    No extra dependencies. This is a wrapper around ComfyUI's own VAE loading, not a separate model or library - it just changes device placement. It shows up under the "multigpu" category once ComfyUI restarts.

    Common issues & troubleshooting

    Cannot handle this data type or garbage output after picking a VAE. Almost always a mismatch - a VAE from one architecture doesn't decode another architecture's latents cleanly (different channel counts under the hood). Use the VAE the checkpoint actually shipped with unless you know the two are latent-compatible.

    You picked pixel_space but your workflow still has a real decode step failing. Pixel-space models genuinely have no VAE and no separate CLIP - if your graph still routes through this loader expecting a normal decode, that's the mismatch. Check the model card for whether it's a latent-space or pixel-space architecture before wiring the workflow.

    Offloading the VAE didn't free as much VRAM as you hoped. That's expected - it's the smallest of the three components. If you're still hitting OOM, the UNet or the text encoder loader is almost always the bigger lever; this node is the easy, low-risk one to offload first, not the one that solves a tight budget by itself.

    Second GPU as donor_device isn't helping. Same caveat as the rest of the pack: on a typical consumer board, GPU-to-GPU transfer over a split PCIe bus can be slower than CPU offload through pinned memory. For something this small it usually doesn't matter either way, but if you're chasing every last bit of speed, test CPU as the donor first.

    Categorymultigpu/distorch_2

    Inputs (6)

    NameTypeDefaultDescription
    vae_nameCOMBO1 options: pixel_space
    compute_deviceoptCOMBOcpu1 options: cpu
    virtual_vram_gboptFLOAT4.00–128
    donor_deviceoptCOMBOcpu1 options: cpu
    expert_mode_allocationsoptSTRING
    eject_modelsoptBOOLEANtrue

    Outputs (1)

    NameTypeDescription
    VAEVAE