ComfyUI GPU Resident Loader

A ComfyUI custom-node pack for faster time-to-VRAM, selective safetensors loading, sticky GPU residency control, and visibility into compatible external GPU model caches.

This repo does three related jobs:

1. Installs startup-time monkey patches before any workflow nodes run.
2. Ships KJ-style resident loader nodes for diffusion models and checkpoints.
3. Maintains a live residency system for native ComfyUI objects and compatible external caches, with preload / pin / evict / report controls for native tracked objects and automatic snapshot support plus provider-specific eviction for compatible external entries.

It is not just a “clean RAM” addon. The main target is the path from model file -> tensors -> live ComfyUI object -> VRAM retention, including GPU-resident caches that live outside ComfyUI’s normal loaded-model list.

Why this exists

ComfyUI’s default behavior mixes together two separate concerns:

• ingest path — where tensors are first materialized while a model is being read, and
• residency policy — where the finished model tends to live afterwards.

Those are not the same problem.

This repo focuses on both:

• For .safetensors, it tries to keep eligible loads on the narrowest, most GPU-friendly path it can.
• For resident diffusion and checkpoint-model loads, it avoids broad checkpoint materialization by selecting only the detected UNet keys where possible.
• For runtime VRAM pressure, it adds a sticky-priority registry and teaches ComfyUI’s unload path to protect higher-value resident entries until enough VRAM must be reclaimed.
• For compatible external GPU caches that bypass comfy.model_management.current_loaded_models, it can discover supported providers at runtime and include their entries in snapshot output and provider-specific eviction decisions, with automatic trim kept opt-in.
• For manual control, it exposes nodes that let you preload, pin, evict, and inspect tracked native models, CLIPs, and VAEs.

What changes at startup

__init__.py calls startup.install_patches(), which applies the core monkey patches exactly once when the custom node is imported.

Those startup patches cover the built-in ComfyUI load and residency paths below. Compatible external-cache integrations are installed lazily later, on demand, when a supported module is actually present in the running process.

Patched functions / methods

Current patch surface:

• comfy.utils.load_torch_file
• comfy.clip_vision.load_torch_file (redirected to the patched comfy.utils.load_torch_file when present)
• comfy.model_management.free_memory
• comfy.model_management.load_models_gpu
• comfy.model_management.unet_offload_device
• comfy.model_management.text_encoder_offload_device
• comfy.model_management.vae_offload_device
• comfy.model_management.text_encoder_device
• comfy.model_management.vae_device
• comfy.model_management.unet_inital_load_device
• comfy.model_management.LoadedModel.model_unload
• comfy.model_patcher.ModelPatcher.detach
• comfy.sd.load_checkpoint_guess_config
• comfy.sd.load_diffusion_model
• comfy.sd.load_clip
• comfy.sd.VAE.encode
• comfy.sd.VAE.decode
• comfy.clip_vision.load
• comfy.controlnet.load_controlnet
• comfy.diffusers_load.load_diffusers

Lazily installed external integrations

Current external integration surface:

• compatible SeedVR2 src/core/model_cache.py modules discovered at runtime

When a compatible SeedVR2 cache module is present, the repo wraps:

• GlobalModelCache.set_dit
• GlobalModelCache.set_vae
• GlobalModelCache.replace_dit (when that method exists in the installed SeedVR2 build)
• GlobalModelCache.replace_vae (when that method exists in the installed SeedVR2 build)
• GlobalModelCache.remove_dit
• GlobalModelCache.remove_vae

That lazy integration lets the loader:

• mirror SeedVR2-owned cached DiT and VAE objects into a separate external residency registry
• refresh byte / device / claimed-state metadata from the live cached object
• evict those entries through SeedVR2’s own removal path instead of assuming they live in comfy.model_management.current_loaded_models

What those patches do

1) load_torch_file becomes residency-aware

The patched loader:

• detects the active load context (model, clip, vae, checkpoint, etc.)
• picks an explicit GPU target device when the active policy wants GPU ingest
• attempts direct safetensors reads on the requested device
• falls back to CPU read + tensor-by-tensor copy if direct GPU safetensors loading fails
• still uses CPU-first torch.load() for pickle formats (.ckpt, .pt, .pth, .bin)
• records the actual load method in the residency registry

For safetensors loads happening inside a model / clip / vae context, it can select only the detected component keys from the file header instead of pulling the full file into memory first.

2) Loader contexts are attached to stock ComfyUI load paths

These stock paths are wrapped with registry context and output binding:

• checkpoint loads
• diffusion-model loads
• CLIP loads
• CLIP Vision loads
• ControlNet loads
• diffusers loads

That means the native registry is not limited to the custom resident nodes. Stock ComfyUI loaders that pass through these paths are also tracked.

3) Compatible external caches can join the residency system lazily

When a compatible SeedVR2 cache module is present, the loader installs cache-level hooks that register SeedVR2-owned cached DiT / VAE objects into a separate external registry.

Those entries:

• are refreshed from the live cached object at runtime
• appear in Registry Snapshot output under external_entries
• are evicted through SeedVR2’s own cache-removal methods rather than the normal Comfy unload path
• remain outside automatic trim unless COMFYUI_GPU_RESIDENT_TRIM_EXTERNAL=1 is set

4) Device/offload policy is overridden

Depending on the active policy, the patcher can steer:

• initial UNet load device
• CLIP/Text Encoder device
• VAE device
• offload devices for UNet / CLIP / VAE

This is how prefer_gpu and sticky_gpu keep more of the hot path on the GPU side than stock ComfyUI would.

5) free_memory() becomes sticky-aware

Under sticky_gpu, comfy.model_management.free_memory() is patched so that:

• sticky tracked wrappers are considered first
• higher-priority sticky entries are protected first
• lower-priority or older sticky entries yield first when VRAM must be reclaimed
• a transient protection floor is applied so ComfyUI does not immediately tear down high-value resident entries for small requests

6) Clone replacement is hardened

load_models_gpu() is patched to fully unload clone-conflict wrappers before replacement instead of relying on a shallow detach path that can leave base weights patched.

7) Unload / detach is redirected to CPU when needed

LoadedModel.model_unload() and ModelPatcher.detach() are patched so that unloads which would otherwise not reclaim VRAM are redirected through a CPU offload target first.

8) VAE encode/decode gets a sticky-safe path

Under sticky_gpu, patched VAE.encode() and VAE.decode():

• cap the working batch count when necessary to preserve transient VRAM headroom
• retry with tiled VAE encode/decode on OOM

That behavior is not a general performance feature toggle. It exists to reduce avoidable VRAM spikes while sticky residency is active.

Policies

The registry exposes four global policies:

legacy

Stay closest to stock ComfyUI behavior. Registry tracking still exists, but the patcher does not aggressively steer ingest/offload toward the GPU path.

balanced

Keep registry tracking and diagnostics without aggressive GPU residency behavior.

prefer_gpu

Prefer GPU ingest for tracked model-like loads and keep the faster side of the device/offload policy for:

• diffusion models
• CLIP / text encoders
• ControlNets

This policy does not auto-pin tracked objects.

sticky_gpu

Builds on prefer_gpu and additionally:

• auto-pins newly bound models and CLIPs
• keeps VAE offload on the GPU side as well
• patches free_memory() to protect sticky tracked wrappers by priority
• uses the sticky-safe VAE encode/decode behavior

Default policy selection

Selection order is:

1. COMFYUI_GPU_RESIDENT_POLICY, if set to a supported value
2. sticky_gpu when ComfyUI is started with --gpu_only
3. sticky_gpu when ComfyUI is started with --highvram
4. otherwise prefer_gpu

Supported values are:

• legacy
• balanced
• prefer_gpu
• sticky_gpu

Included nodes

All nodes live under the GPU Resident Loader category.

Loader nodes

Diffusion Model Selector Resident

Returns an absolute path string for a selected diffusion model.

Notes:

• resolves from diffusion_models
• also exposes text_encoders entries whose filename contains connector

Diffusion Model Loader Resident

KJ-style diffusion-model loader with these controls:

• weight_dtype
• compute_dtype
• patch_cublaslinear
• sage_attention
• enable_fp16_accumulation
• optional extra_state_dict
• optional policy_override

Behavior:

• for .safetensors, it loads only the detected UNet portion of the file
• if extra_state_dict is provided, only matching UNet keys are merged
• repeated loads reuse a live equivalent model when the source path and loader-relevant options still match
• before GPU-bound loads, it estimates the upcoming footprint and trims only enough lower-priority residency to cover the request plus adaptive headroom

Checkpoint Loader Resident

Full checkpoint loader that returns:

• MODEL
• CLIP
• VAE

Behavior:

• shares the same tuning knobs as the resident diffusion-model loader for the model component
• reuses already-live equivalent components when possible
• composes the final output from model / clip / vae component loaders instead of always rebuilding the whole checkpoint path from scratch

Checkpoint Model Loader Resident

Model-only checkpoint loader.

Behavior:

• takes the same selective safetensors UNet fast path as the diffusion-model loader
• reuses a live equivalent model when available
• uses the same dtype / attention / cublas / fp16-accumulation knobs as the full checkpoint loader

Checkpoint Clip Loader Resident

CLIP-only checkpoint loader.

Behavior:

• can reuse a live equivalent CLIP object
• avoids rebuilding the diffusion model and VAE outputs when only CLIP is needed

Checkpoint VAE Loader Resident

VAE-only checkpoint loader.

Behavior:

• can reuse a live equivalent VAE object
• avoids rebuilding the diffusion model and CLIP outputs when only VAE is needed

Residency nodes

Set Global Residency Policy

Sets the active global policy and returns it as a STRING.

The loader nodes also expose an optional policy_override string input for one-off loads.

Registry Snapshot

Returns a composite formatted JSON snapshot with:

• policy for the active global policy
• entries for native Comfy-managed registry entries
• external_entries for compatible external cache entries discovered at runtime

Pin Model Residency / Pin CLIP Residency / Pin VAE Residency

Marks a tracked native object as sticky or non-sticky and optionally changes its priority.

Preload Model To GPU / Preload CLIP To GPU / Preload VAE To GPU

Calls load_models_gpu(..., force_full_load=True) for the selected native object, then updates sticky state / priority in the registry.

Evict Model From GPU / Evict CLIP From GPU / Evict VAE From GPU

Attempts to unload the selected native object from the current loaded-model set.

unpatch_weights=True performs a full unload path. When eviction succeeds, the node returns evicted; otherwise not_loaded.

Report Model Residency / Report CLIP Residency / Report VAE Residency

Returns a JSON report for a single tracked native object.

If the object is not currently bound in the registry, the node returns a JSON payload with tracked: false.

Adaptive trimming before resident loads

The resident loaders now do load-scoped VRAM trimming themselves.

Before a GPU-bound resident load, the loader estimates required bytes from:

• the safetensors header when possible
• the detected checkpoint component subset when possible
• otherwise the source file size as a fallback

It then requests enough free VRAM for:

• the estimated load size
• adaptive headroom

Current adaptive headroom policy:

• ratio: 12.5% of the estimated load
• floor: 256 MiB
• ceiling: 1 GiB

The trim path prefers to:

• unload non-sticky entries first
• then lower-priority sticky entries
• preserve explicitly kept models
• use partial unload where available
• include compatible external cache entries in the same candidate search only when COMFYUI_GPU_RESIDENT_TRIM_EXTERNAL=1, they are visible, on the same device, and not currently claimed/in use

This logic lives in the resident loader path. You do not need a separate “target free VRAM” node for it.

Registry and observability

The snapshot now exposes two collections:

• entries for native Comfy-managed tracked objects
• external_entries for compatible external cache objects

Native registry entries

The native registry tracks residency metadata for bound objects.

Typical per-entry fields include:

• entry_id
• kind
• source_path
• basename
• sticky
• priority
• created_at
• last_touched
• loaded_bytes
• total_bytes
• load_device
• offload_device
• current_device
• last_method
• last_report
• loader_key
• notes
• alive

The last_method / last_report fields let you see whether a load actually used:

• direct safetensors GPU ingest
• safetensors CPU -> CUDA fallback
• safetensors component-only load
• CPU-first torch.load() compatibility path
• a recorded load failure

External registry entries

Typical external-entry fields include:

• entry_id
• cache_key
• kind
• source_path
• basename
• sticky
• priority
• created_at
• last_touched
• loaded_bytes
• total_bytes
• load_device
• offload_device
• current_device
• claimed
• notes
• alive
• external

For SeedVR2-backed entries, claimed: true means the cache object is currently marked in use and is skipped by the external trim candidate search when that opt-in path is enabled.

What gets tracked

Native tracked/bound paths

Native tracked paths include:

• resident node loads from this repo
• stock checkpoint loads
• stock diffusion-model loads
• stock CLIP loads
• stock CLIP Vision loads
• stock diffusers loads

ControlNet loads also participate in the patched load context and device-policy path, but this repo does not currently expose dedicated ControlNet residency nodes.

External tracked paths

Current external integration coverage is:

• SeedVR2 global cached DiT entries
• SeedVR2 global cached VAE entries

Those entries are discovered lazily from compatible SeedVR2 cache modules at runtime. They are tracked separately from the native registry and participate in snapshot plus provider-specific eviction decisions, with load-scoped trim available only through the explicit external-trim opt-in.

Important limits and non-goals

Best path is still .safetensors

The narrow fast path is built around .safetensors.

That is where this repo can:

• inspect headers cheaply
• select only model / clip / vae subsets
• estimate component bytes more accurately
• attempt direct device-targeted reads

.ckpt / .pt / pickle formats are still CPU-first

For pickle-based formats, PyTorch still goes through torch.load() on CPU first.

The repo can still:

• track those loads
• keep the resulting live objects resident
• reuse equivalent live objects later

It does not claim direct-to-GPU ingest for those formats.

Cross-process persistence is out of scope

This repo does not keep VRAM allocations alive after ComfyUI, Python, or WSL exits.

CUDA memory lifetime is process/context scoped. True persistence across process shutdown would need a separate long-lived keeper process or service that owns the CUDA context.

External integrations are compatibility-based, not universal

The external registry does not automatically manage every third-party cache.

At the moment, the documented external integration target is SeedVR2. Other custom nodes with private caches remain invisible until this repo grows a provider-specific integration for them.

It does not automatically capture arbitrary custom loader implementations

The native registry only sees objects that pass through the patched ComfyUI load paths or through this repo’s resident nodes.

If another custom node loads models through its own private code path and bypasses those patched entry points, that object may never become a tracked native registry entry. In that case, the preload / pin / evict / report nodes from this repo cannot manage it until that external loader is integrated or patched.

Likewise, even for supported external providers such as SeedVR2, the current external integration is about observation + provider-specific eviction, with automatic trim kept opt-in. This repo does not yet expose dedicated external preload / pin / report / evict nodes for provider-owned cache entries.

Installation

Clone into custom_nodes:

git clone https://github.com/xmarre/ComfyUI-GPU-Resident-Loader ComfyUI/custom_nodes/ComfyUI-GPU-Resident-Loader

Install dependencies into the same Python environment ComfyUI uses:

pip install -r ComfyUI/custom_nodes/ComfyUI-GPU-Resident-Loader/requirements.txt

Requirements declared by the repo:

• Python >=3.10
• safetensors>=0.4.3

Optional SageAttention dependencies are not installed by default. Install those separately if you plan to use a SageAttention mode in the resident loaders.

Basic usage patterns

1) Large-VRAM, mostly resident workflow

Recommended baseline:

• start ComfyUI with --highvram or set policy manually to sticky_gpu
• prefer .safetensors for hot models
• load diffusion models through Diffusion Model Loader Resident
• use Preload ... To GPU for models you know you will reuse
• inspect with Report ... Residency or Registry Snapshot

2) Full checkpoint workflow

Use Checkpoint Loader Resident when you want MODEL + CLIP + VAE together.

That path can reuse already-live components instead of always rebuilding all three outputs.

3) Staged checkpoint workflow

Use component loaders when the graph does not need the whole checkpoint at once:

• Checkpoint Model Loader Resident for diffusion model only
• Checkpoint Clip Loader Resident for CLIP only
• Checkpoint VAE Loader Resident for VAE only

4) Manual native residency control

Use:

• Pin ... Residency to mark a tracked native entry sticky / non-sticky
• Preload ... To GPU to force a full live native load now
• Evict ... From GPU to unload it from the current native loaded-model set

5) Mixed workflows with SeedVR2 external caching

If SeedVR2 keeps DiT or VAE models in its own global cache, those objects can now show up in Registry Snapshot under external_entries.

That means:

• you can see that those bytes exist even though they are outside current_loaded_models
• resident loader trim only reclaims them when COMFYUI_GPU_RESIDENT_TRIM_EXTERNAL=1, they are visible, on the same device, and not currently claimed/in use
• eviction goes through SeedVR2’s own cache-removal path instead of a normal Comfy wrapper unload

Notes on compatibility and migration

Legacy wiring: extra_state_dict used as a policy string

The resident diffusion-model loader contains a compatibility shim for older graphs:

• if extra_state_dict receives one of the known policy names
• and that value is not an existing file path
• it is interpreted as policy_override instead

New graphs should connect policy strings to policy_override, not to extra_state_dict.

Convert hot pickle checkpoints to safetensors

scripts/convert_checkpoint_to_safetensors.py is included for one-time conversion of hot .ckpt / .pt / .pth style checkpoints.

Example:

python ComfyUI/custom_nodes/ComfyUI-GPU-Resident-Loader/scripts/convert_checkpoint_to_safetensors.py \
  --input /path/to/model.ckpt \
  --output /path/to/model.safetensors

Optional flags:

• --state-dict-key <key> to extract a different top-level dict key
• --allow-non-tensor-values to skip non-tensor entries instead of failing

License

GPL-3.0-or-later.

This repo stays GPL-compatible because it adapts behavior from GPL-licensed ComfyUI and mirrors relevant loader behavior from GPL-3.0-licensed KJNodes.