Checkpoint Loader (Unified Memory)

If you run ComfyUI on a DGX Spark or GB10, you've probably watched the same pointless dance: your first Flux run hangs for minutes while ComfyUI loads a checkpoint, and the model spends part of that trip in places it doesn't need to be. On those machines the CPU and GPU share one 128 GB memory pool, and ComfyUI's stock loader doesn't know what to do with that. Checkpoint Loader (Unified Memory) is a drop-in replacement for CheckpointLoaderSimple that loads the whole checkpoint straight into CUDA and keeps it there.

What it actually is

This is the flagship of the DGX Nodes pack by broken-gage - a set of loaders that swap the stock ComfyUI loading pipeline for a direct-to-CUDA path tuned for NVIDIA's Grace-Blackwell unified-memory systems. It takes a checkpoint file from your checkpoints/ folder and hands you the same three outputs as the stock node: MODEL, CLIP, and VAE. Wire MODEL into your sampler, CLIP into your text encoder, VAE into VAEDecode. Nothing about your workflow changes; you just swap the node and the first load stops being a coffee break.

Why the stock loader is bad at this

The pack's source has the root cause spelled out. ComfyUI decides where to stage weights with unet_inital_load_device(), which uses a strict "is GPU memory bigger than CPU memory?" check. On a GB10 both come from the same pool, so they're equal, the check fails, and the model gets loaded to CPU instead. Then the tensors get copied CPU→GPU at inference time anyway. The result on some workloads: disk → CUDA → copy to CPU → CUDA freed, and the GPU re-loads everything when sampling starts. One machine's benchmark showed a Flux.2 Dev first run taking ~400 s on the native path.

The DGX path cuts that by skipping the CPU staging entirely: safetensors tensors are read directly into CUDA, assign=True makes them the model weights with no copy, and ComfyUI's memory tracking is told the model is already GPU-resident so it doesn't evict your other models. The author's own tests dropped that ~400 s first run to ~73 s with the instanttensor backend - and, importantly, generation speed is identical to native. This only saves loading time, and it genuinely does.

The inputs that matter

You set three things besides the file:

• dgx_mode - the on/off toggle. ON uses the DGX direct path; OFF is plain stock loading. Defaults to ON.
• device - which CUDA device to load onto (cuda:0).
• storage_backend - auto tries instanttensor first (loads at ~1× memory, the one you want), then fastsafetensors (host-mmap + CUDA DMA, ~2× peak memory). auto is the sensible default.

Installing it

Install the pack through ComfyUI Manager by searching DGX Nodes, or clone it:

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

Restart ComfyUI and the nodes appear under the DGX Nodes category. This pack needs ComfyUI 0.24.0+.

Where people get burned

• Safetensors only in DGX mode. Point this at a .ckpt or .pt with dgx_mode ON and it raises and tells you to use CheckpointLoaderSimple. Flip dgx_mode off (stock path handles any format) or use a safetensors checkpoint.
• It's a DGX-hardware thing. On Windows or a normal x86 rig with no DGX backends installed, the node auto-falls back to stock loading - so it's harmless, but it's also not doing anything. Leave dgx_mode off there.
• Don't load weights bigger than available memory with dgx_mode on; the README warns it can cause instability and OOMs.
• The pack is honest with you: it's a "vibe-code" project, GPLv3, agentic-AI-generated, and not guaranteed fully working in every environment. For a loader that mostly delegates to ComfyUI's own machinery, that's a fair trade - but don't treat it as production infrastructure.