Add sticky VAE preflight tiling

[codex] restore SeedVR2 cache reuse by default

README.md

@@ -1,14 +1,14 @@
 # ComfyUI GPU Resident Loader
 
-A ComfyUI custom-node pack for **faster time-to-VRAM**, **selective safetensors loading**, and **sticky GPU residency control**.
+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 registry** with preload / pin / evict / report controls for tracked objects.
+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.
+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
 
@@ -24,11 +24,14 @@ 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 **manual control**, it exposes nodes that let you preload, pin, evict, and inspect tracked models, CLIPs, and VAEs.
+- 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 monkey patches exactly once when the custom node is imported.
+`__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
 
@@ -55,68 +58,100 @@ Current patch surface:
 - `comfy.controlnet.load_controlnet`
 - `comfy.diffusers_load.load_diffusers`
 
-### What those patches do
+### Lazily installed external integrations
 
-#### 1) `load_torch_file` becomes residency-aware
+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.
+- 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
+### 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.
+- checkpoint loads
+- diffusion-model loads
+- CLIP loads
+- CLIP Vision loads
+- ControlNet loads
+- diffusers loads
 
-That means the registry is not limited to the custom resident nodes. Stock ComfyUI loaders that pass through these paths are also tracked.
+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) Device/offload policy is overridden
+### 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.
+- 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.
 
-#### 4) `free_memory()` becomes sticky-aware
+### 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.
+- 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
 
-#### 5) Clone replacement is hardened
+### 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.
 
-#### 6) Unload / detach is redirected to CPU when needed
+### 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.
 
-#### 7) VAE encode/decode gets a sticky-safe path
+### 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, and
-- retry with tiled VAE encode/decode on OOM.
+- 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.
 
@@ -136,9 +171,9 @@ Keep registry tracking and diagnostics without aggressive GPU residency behavior
 
 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.
+- diffusion models
+- CLIP / text encoders
+- ControlNets
 
 This policy does **not** auto-pin tracked objects.
 
@@ -146,19 +181,19 @@ This policy does **not** auto-pin tracked objects.
 
 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.
+- 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`.
+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:
 
@@ -179,8 +214,8 @@ Returns an absolute path string for a selected diffusion model.
 
 Notes:
 
-- resolves from `diffusion_models`, and
-- also exposes `text_encoders` entries whose filename contains `connector`.
+- resolves from `diffusion_models`
+- also exposes `text_encoders` entries whose filename contains `connector`
 
 #### Diffusion Model Loader Resident
 
@@ -196,10 +231,10 @@ KJ-style diffusion-model loader with these controls:
 
 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.
+- 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
 
@@ -211,9 +246,9 @@ Full checkpoint loader that returns:
 
 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.
+- 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
 
@@ -221,9 +256,9 @@ 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.
+- 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
 
@@ -231,8 +266,8 @@ 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.
+- can reuse a live equivalent CLIP object
+- avoids rebuilding the diffusion model and VAE outputs when only CLIP is needed
 
 #### Checkpoint VAE Loader Resident
 
@@ -240,8 +275,8 @@ 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.
+- can reuse a live equivalent VAE object
+- avoids rebuilding the diffusion model and CLIP outputs when only VAE is needed
 
 ### Residency nodes
 
@@ -253,25 +288,29 @@ The loader nodes also expose an optional `policy_override` string input for one-
 
 #### Registry Snapshot
 
-Returns the whole registry as formatted JSON.
+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 object as sticky or non-sticky and optionally changes its priority.
+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 object, then updates sticky state / priority in the registry.
+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 object from the current loaded-model set.
+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 object.
+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`.
 
@@ -281,33 +320,41 @@ 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.
+- 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, plus
-- adaptive headroom.
+- the estimated load size
+- adaptive headroom
 
 Current adaptive headroom policy:
 
-- ratio: `12.5%` of the estimated load,
-- floor: `256 MiB`,
-- ceiling: `1 GiB`.
+- 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.
+- 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 registry tracks residency metadata for bound objects.
+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:
 
@@ -332,25 +379,61 @@ Typical per-entry fields include:
 
 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,
-- or a recorded load failure.
+- 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
 
-Tracked/bound paths include:
+### Native tracked/bound paths
 
-- resident node loads from this repo,
-- stock checkpoint loads,
-- stock diffusion-model loads,
-- stock CLIP loads,
-- stock CLIP Vision loads,
-- stock diffusers loads.
+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`
@@ -359,10 +442,10 @@ 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.
+- 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
 
@@ -370,9 +453,9 @@ 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.
+- 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.
 
@@ -382,11 +465,19 @@ This repo does **not** keep VRAM allocations alive after ComfyUI, Python, or WSL
 
 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 registry only sees objects that pass through the patched ComfyUI load paths or through this repo’s resident nodes.
+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 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.
+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
 
@@ -415,11 +506,11 @@ Optional SageAttention dependencies are **not** installed by default. Install th
 
 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**.
+- 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
 
@@ -431,17 +522,27 @@ That path can reuse already-live components instead of always rebuilding all thr
 
 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.
+- **Checkpoint Model Loader Resident** for diffusion model only
+- **Checkpoint Clip Loader Resident** for CLIP only
+- **Checkpoint VAE Loader Resident** for VAE only
 
-### 4) Manual residency control
+### 4) Manual native residency control
 
 Use:
 
-- **Pin ... Residency** to mark a tracked entry sticky / non-sticky,
-- **Preload ... To GPU** to force a full live load now,
-- **Evict ... From GPU** to unload it from the current loaded-model set.
+- **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
 
@@ -449,9 +550,9 @@ Use:
 
 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.
+- 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`.
 

cleanup.py

@@ -7,7 +7,7 @@ from typing import Any
 import comfy.model_management as model_management
 import torch
 
-from .external_residency import EXTERNAL_REGISTRY, ensure_external_integrations_installed
+from .external_residency import EXTERNAL_REGISTRY, ensure_external_integrations_installed, external_trim_enabled
 from .residency import REGISTRY
 
 _ADAPTIVE_HEADROOM_RATIO = 0.125
@@ -154,13 +154,14 @@ def _trim_candidates(
         )
         candidates.append((loaded, entry, sticky_respected, False))
 
-    for external_obj, entry, sticky_respected in EXTERNAL_REGISTRY.candidates(
-        device=device,
-        respect_sticky=respect_sticky,
-        sticky_floor_priority=sticky_floor_priority,
-        keep_models=keep_models,
-    ):
-        candidates.append((external_obj, entry, sticky_respected, True))
+    if external_trim_enabled():
+        for external_obj, entry, sticky_respected in EXTERNAL_REGISTRY.candidates(
+            device=device,
+            respect_sticky=respect_sticky,
+            sticky_floor_priority=sticky_floor_priority,
+            keep_models=keep_models,
+        ):
+            candidates.append((external_obj, entry, sticky_respected, True))
 
     candidates.sort(key=lambda item: _sort_key_for_candidate(item[1], sticky_respected=item[2]))
     return candidates
@@ -305,6 +306,7 @@ def trim_resident_vram(
         "respect_sticky": bool(respect_sticky),
         "sticky_floor_priority": int(sticky_floor_priority),
         "allow_partial_unload": bool(allow_partial_unload),
+        "external_trim_enabled": bool(external_trim_enabled()),
         "actions": actions,
     }
 

external_residency.py

@@ -280,7 +280,7 @@ class ExternalResidencyRegistry:
             entry.state_provider = state_provider
             entry.evict_callback = evict_callback
             entry.last_touched = _now()
-            if note:
+            if note and note not in entry.notes:
                 entry.notes.append(note)
 
         self.refresh_runtime_state()
@@ -398,6 +398,11 @@ class ExternalResidencyRegistry:
 EXTERNAL_REGISTRY = ExternalResidencyRegistry()
 
 
+def external_trim_enabled() -> bool:
+    value = os.environ.get("COMFYUI_GPU_RESIDENT_TRIM_EXTERNAL", "").strip().lower()
+    return value in {"1", "true", "yes", "on"}
+
+
 def _call_seedvr2_method_with_optional_expected_model(
     method: Callable[..., Any],
     *args: Any,
@@ -418,18 +423,18 @@ def _call_seedvr2_method_with_optional_expected_model(
     return method(*args, **kwargs)
 
 
-def _register_seedvr2_cached_model(global_cache: Any, *, kind: str, config: Any, model: Any) -> None:
+def _register_seedvr2_cached_model(global_cache: Any, *, kind: str, config: Any, model: Any) -> ExternalResidencyEntry | None:
     if not isinstance(config, Mapping):
         _LOG.debug(
             "GPU Resident Loader: skipping SeedVR2 %s cache entry with unexpected config type: %s",
             kind,
             type(config).__name__,
         )
-        return
+        return None
 
     node_id = config.get("node_id")
     if node_id is None or model is None:
-        return
+        return None
 
     try:
         model_ref = weakref.ref(model)
@@ -464,7 +469,7 @@ def _register_seedvr2_cached_model(global_cache: Any, *, kind: str, config: Any,
             )
         note = f"SeedVR2 cached VAE node {node_id}"
 
-    EXTERNAL_REGISTRY.bind(
+    return EXTERNAL_REGISTRY.bind(
         cache_key=cache_key,
         obj=model,
         kind=registry_kind,
@@ -485,24 +490,33 @@ def _install_seedvr2_integration_for_module(module: Any) -> bool:
 
     class_id = id(model_cache_cls)
     thread_id = threading.get_ident()
-    with _SEEDVR2_PATCH_CONDITION:
-        while (
-            class_id in _SEEDVR2_PATCHING_CLASS_IDS
-            and _SEEDVR2_PATCHING_THREAD_IDS.get(class_id) != thread_id
-        ):
-            _SEEDVR2_PATCH_CONDITION.wait()
-        if class_id in _SEEDVR2_PATCHED_CLASS_IDS or _SEEDVR2_PATCHING_THREAD_IDS.get(class_id) == thread_id:
-            return True
-        _SEEDVR2_PATCHING_CLASS_IDS.add(class_id)
-        _SEEDVR2_PATCHING_THREAD_IDS[class_id] = thread_id
+    original_set_dit = None
+    original_set_vae = None
+    original_replace_dit = None
+    original_replace_vae = None
+    original_remove_dit = None
+    original_remove_vae = None
+    provisional_cache_keys: set[str] = set()
+    provisional_cache_ownership: dict[str, str] = {}
 
     try:
-        original_set_dit = model_cache_cls.set_dit
-        original_set_vae = model_cache_cls.set_vae
-        original_replace_dit = getattr(model_cache_cls, "replace_dit", None)
-        original_replace_vae = getattr(model_cache_cls, "replace_vae", None)
-        original_remove_dit = model_cache_cls.remove_dit
-        original_remove_vae = model_cache_cls.remove_vae
+        with _SEEDVR2_PATCH_CONDITION:
+            while (
+                class_id in _SEEDVR2_PATCHING_CLASS_IDS
+                and _SEEDVR2_PATCHING_THREAD_IDS.get(class_id) != thread_id
+            ):
+                _SEEDVR2_PATCH_CONDITION.wait()
+            if class_id in _SEEDVR2_PATCHED_CLASS_IDS or _SEEDVR2_PATCHING_THREAD_IDS.get(class_id) == thread_id:
+                return True
+            _SEEDVR2_PATCHING_CLASS_IDS.add(class_id)
+            _SEEDVR2_PATCHING_THREAD_IDS[class_id] = thread_id
+
+            original_set_dit = model_cache_cls.set_dit
+            original_set_vae = model_cache_cls.set_vae
+            original_replace_dit = getattr(model_cache_cls, "replace_dit", None)
+            original_replace_vae = getattr(model_cache_cls, "replace_vae", None)
+            original_remove_dit = model_cache_cls.remove_dit
+            original_remove_vae = model_cache_cls.remove_vae
 
         def set_dit_wrapper(self, dit_config, model, model_name, debug=None):
             result = original_set_dit(self, dit_config, model, model_name, debug)
@@ -570,54 +584,92 @@ def _install_seedvr2_integration_for_module(module: Any) -> bool:
                 EXTERNAL_REGISTRY.remove(cache_key=_seedvr2_entry_key("vae", vae_config.get("node_id")))
             return result
 
-        model_cache_cls.set_dit = set_dit_wrapper
-        model_cache_cls.set_vae = set_vae_wrapper
-        if original_replace_dit is not None:
-            model_cache_cls.replace_dit = replace_dit_wrapper
-        if original_replace_vae is not None:
-            model_cache_cls.replace_vae = replace_vae_wrapper
-        model_cache_cls.remove_dit = remove_dit_wrapper
-        model_cache_cls.remove_vae = remove_vae_wrapper
         global_cache = get_global_cache()
         model_cache_lock = getattr(global_cache, "_model_cache_lock", None)
         lock_context = model_cache_lock if model_cache_lock is not None else contextlib.nullcontext()
         with lock_context:
+            model_cache_cls.set_dit = set_dit_wrapper
+            model_cache_cls.set_vae = set_vae_wrapper
+            if original_replace_dit is not None:
+                model_cache_cls.replace_dit = replace_dit_wrapper
+            if original_replace_vae is not None:
+                model_cache_cls.replace_vae = replace_vae_wrapper
+            model_cache_cls.remove_dit = remove_dit_wrapper
+            model_cache_cls.remove_vae = remove_vae_wrapper
             dit_items = list(getattr(global_cache, "_dit_models", {}).items())
             vae_items = list(getattr(global_cache, "_vae_models", {}).items())
-        for node_id, entry in dit_items:
-            if not isinstance(entry, tuple) or len(entry) != 2:
-                continue
-            model, config = entry
-            if model is not None:
-                _register_seedvr2_cached_model(global_cache, kind="dit", config=config, model=model)
-        for node_id, entry in vae_items:
-            if not isinstance(entry, tuple) or len(entry) != 2:
-                continue
-            model, config = entry
-            if model is not None:
-                _register_seedvr2_cached_model(global_cache, kind="vae", config=config, model=model)
-    except Exception:
+            for _node_id, entry in dit_items:
+                if not isinstance(entry, tuple) or len(entry) != 2:
+                    continue
+                model, config = entry
+                if model is not None:
+                    if isinstance(config, Mapping) and config.get("node_id") is not None:
+                        cache_key = _seedvr2_entry_key("dit", config.get("node_id"))
+                        provisional_cache_keys.add(cache_key)
+                        prior_entry_id = EXTERNAL_REGISTRY._cache_key_to_entry.get(cache_key)
+                    else:
+                        cache_key = None
+                        prior_entry_id = None
+                    registered_entry = _register_seedvr2_cached_model(global_cache, kind="dit", config=config, model=model)
+                    if cache_key is not None and registered_entry is not None:
+                        if prior_entry_id != registered_entry.entry_id:
+                            provisional_cache_ownership[cache_key] = registered_entry.entry_id
+            for _node_id, entry in vae_items:
+                if not isinstance(entry, tuple) or len(entry) != 2:
+                    continue
+                model, config = entry
+                if model is not None:
+                    if isinstance(config, Mapping) and config.get("node_id") is not None:
+                        cache_key = _seedvr2_entry_key("vae", config.get("node_id"))
+                        provisional_cache_keys.add(cache_key)
+                        prior_entry_id = EXTERNAL_REGISTRY._cache_key_to_entry.get(cache_key)
+                    else:
+                        cache_key = None
+                        prior_entry_id = None
+                    registered_entry = _register_seedvr2_cached_model(global_cache, kind="vae", config=config, model=model)
+                    if cache_key is not None and registered_entry is not None:
+                        if prior_entry_id != registered_entry.entry_id:
+                            provisional_cache_ownership[cache_key] = registered_entry.entry_id
         with _SEEDVR2_PATCH_CONDITION:
-            model_cache_cls.set_dit = original_set_dit
-            model_cache_cls.set_vae = original_set_vae
+            _SEEDVR2_PATCHING_CLASS_IDS.discard(class_id)
+            _SEEDVR2_PATCHED_CLASS_IDS.add(class_id)
+            _SEEDVR2_PATCHING_THREAD_IDS.pop(class_id, None)
+            _SEEDVR2_PATCH_CONDITION.notify_all()
+    except Exception:
+        _LOG.debug(
+            "GPU Resident Loader: rolling back SeedVR2 integration for class_id=%s provisional_cache_keys=%s",
+            class_id,
+            sorted(provisional_cache_keys),
+            exc_info=True,
+        )
+        for cache_key in provisional_cache_keys:
+            created_entry_id = provisional_cache_ownership.get(cache_key)
+            if created_entry_id is not None:
+                with EXTERNAL_REGISTRY._lock:
+                    current_entry_id = EXTERNAL_REGISTRY._cache_key_to_entry.get(cache_key)
+                    if current_entry_id == created_entry_id:
+                        EXTERNAL_REGISTRY.remove(cache_key=cache_key)
+        with lock_context:
+            if original_set_dit is not None:
+                model_cache_cls.set_dit = original_set_dit
+            if original_set_vae is not None:
+                model_cache_cls.set_vae = original_set_vae
             if original_replace_dit is not None:
                 model_cache_cls.replace_dit = original_replace_dit
             if original_replace_vae is not None:
                 model_cache_cls.replace_vae = original_replace_vae
-            model_cache_cls.remove_dit = original_remove_dit
-            model_cache_cls.remove_vae = original_remove_vae
+            if original_remove_dit is not None:
+                model_cache_cls.remove_dit = original_remove_dit
+            if original_remove_vae is not None:
+                model_cache_cls.remove_vae = original_remove_vae
+        with _SEEDVR2_PATCH_CONDITION:
             _SEEDVR2_PATCHING_CLASS_IDS.discard(class_id)
             _SEEDVR2_PATCHED_CLASS_IDS.discard(class_id)
             _SEEDVR2_PATCHING_THREAD_IDS.pop(class_id, None)
             _SEEDVR2_PATCH_CONDITION.notify_all()
         raise
 
-    with _SEEDVR2_PATCH_CONDITION:
-        _SEEDVR2_PATCHING_CLASS_IDS.discard(class_id)
-        _SEEDVR2_PATCHED_CLASS_IDS.add(class_id)
-        _SEEDVR2_PATCHING_THREAD_IDS.pop(class_id, None)
-        _SEEDVR2_PATCH_CONDITION.notify_all()
-    _LOG.info("GPU Resident Loader: integrated external SeedVR2 cache eviction hooks")
+    _LOG.info("GPU Resident Loader: integrated external SeedVR2 cache visibility hooks")
     return True
 
 

patches.py

@@ -661,6 +661,67 @@ def _sticky_safe_batch_number(*, batch_count: int, free_memory: int, memory_used
     return max(1, capped)
 
 
+def _scaled_batch_memory(total_memory_used: int, total_batch_count: int, batch_number: int) -> int:
+    total_memory = max(1, int(total_memory_used))
+    total_batches = max(1, int(total_batch_count))
+    current_batch = max(1, min(int(batch_number), total_batches))
+    return max(1, (total_memory * current_batch + total_batches - 1) // total_batches)
+
+
+def _prepare_sticky_vae_batch(
+    *,
+    device: Any,
+    patcher: Any,
+    total_memory_used: int,
+    total_batch_count: int,
+) -> tuple[int, bool]:
+    free_memory = int(patcher.get_free_memory(device))
+    batch_number = _sticky_safe_batch_number(
+        batch_count=total_batch_count,
+        free_memory=free_memory,
+        memory_used=total_memory_used,
+        device=device,
+    )
+    batch_memory_used = _scaled_batch_memory(total_memory_used, total_batch_count, batch_number)
+
+    if REGISTRY.get_policy() != "sticky_gpu" or device is None:
+        return batch_number, False
+
+    target_free = _sticky_protection_target(batch_memory_used, device)
+    if free_memory < target_free:
+        try:
+            trim_resident_vram(
+                device=device,
+                target_free_vram_bytes=target_free,
+                respect_sticky=True,
+                sticky_floor_priority=0,
+                allow_partial_unload=True,
+                keep_models=(patcher,),
+            )
+        except Exception as exc:
+            _LOG.debug("GPU Resident Loader: proactive VAE trim failed for %s bytes: %s", batch_memory_used, exc)
+
+        free_memory = int(patcher.get_free_memory(device))
+        batch_number = _sticky_safe_batch_number(
+            batch_count=total_batch_count,
+            free_memory=free_memory,
+            memory_used=total_memory_used,
+            device=device,
+        )
+        batch_memory_used = _scaled_batch_memory(total_memory_used, total_batch_count, batch_number)
+        target_free = _sticky_protection_target(batch_memory_used, device)
+
+    should_tile = free_memory < target_free and batch_number <= 1
+    if should_tile:
+        _LOG.info(
+            "GPU Resident Loader: skipping regular VAE pass and switching directly to tiled mode; free=%s target=%s batch_memory=%s",
+            free_memory,
+            target_free,
+            batch_memory_used,
+        )
+    return batch_number, should_tile
+
+
 def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]:
     @functools.wraps(func)
     def wrapper(self, pixel_samples):
@@ -681,29 +742,31 @@ def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]:
         try:
             memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
             model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
-            free_memory = self.patcher.get_free_memory(self.device)
-            batch_number = _sticky_safe_batch_number(
-                batch_count=pixel_samples.shape[0],
-                free_memory=free_memory,
-                memory_used=memory_used,
+            batch_number, should_tile = _prepare_sticky_vae_batch(
                 device=self.device,
+                patcher=self.patcher,
+                total_memory_used=memory_used,
+                total_batch_count=pixel_samples.shape[0],
             )
-            samples = None
-            for x in range(0, pixel_samples.shape[0], batch_number):
-                pixels_in = self.process_input(pixel_samples[x:x + batch_number]).to(self.vae_dtype)
-                if getattr(self.first_stage_model, "comfy_has_chunked_io", False):
-                    out = self.first_stage_model.encode(pixels_in, device=self.device)
-                else:
-                    pixels_in = pixels_in.to(self.device)
-                    out = self.first_stage_model.encode(pixels_in)
-                out = out.to(self.output_device).to(dtype=self.vae_output_dtype())
-                if samples is None:
-                    samples = torch.empty(
-                        (pixel_samples.shape[0],) + tuple(out.shape[1:]),
-                        device=self.output_device,
-                        dtype=self.vae_output_dtype(),
-                    )
-                samples[x:x + batch_number] = out
+            if should_tile:
+                do_tile = True
+            else:
+                samples = None
+                for x in range(0, pixel_samples.shape[0], batch_number):
+                    pixels_in = self.process_input(pixel_samples[x:x + batch_number]).to(self.vae_dtype)
+                    if getattr(self.first_stage_model, "comfy_has_chunked_io", False):
+                        out = self.first_stage_model.encode(pixels_in, device=self.device)
+                    else:
+                        pixels_in = pixels_in.to(self.device)
+                        out = self.first_stage_model.encode(pixels_in)
+                    out = out.to(self.output_device).to(dtype=self.vae_output_dtype())
+                    if samples is None:
+                        samples = torch.empty(
+                            (pixel_samples.shape[0],) + tuple(out.shape[1:]),
+                            device=self.output_device,
+                            dtype=self.vae_output_dtype(),
+                        )
+                    samples[x:x + batch_number] = out
         except Exception as e:
             model_management.raise_non_oom(e)
             _LOG.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
@@ -741,16 +804,15 @@ def _wrap_vae_decode(func: Callable[..., Any]) -> Callable[..., Any]:
         try:
             memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
             model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
-            free_memory = self.patcher.get_free_memory(self.device)
-            batch_number = _sticky_safe_batch_number(
-                batch_count=samples_in.shape[0],
-                free_memory=free_memory,
-                memory_used=memory_used,
+            batch_number, should_tile = _prepare_sticky_vae_batch(
                 device=self.device,
+                patcher=self.patcher,
+                total_memory_used=memory_used,
+                total_batch_count=samples_in.shape[0],
             )
 
             preallocated = False
-            if getattr(self.first_stage_model, "comfy_has_chunked_io", False):
+            if not should_tile and getattr(self.first_stage_model, "comfy_has_chunked_io", False):
                 pixel_samples = torch.empty(
                     self.first_stage_model.decode_output_shape(samples_in.shape),
                     device=self.output_device,
@@ -758,25 +820,28 @@ def _wrap_vae_decode(func: Callable[..., Any]) -> Callable[..., Any]:
                 )
                 preallocated = True
 
-            for x in range(0, samples_in.shape[0], batch_number):
-                samples = samples_in[x:x + batch_number].to(device=self.device, dtype=self.vae_dtype)
-                if preallocated:
-                    self.first_stage_model.decode(samples, output_buffer=pixel_samples[x:x + batch_number], **vae_options)
-                else:
-                    out = self.first_stage_model.decode(samples, **vae_options).to(
-                        device=self.output_device,
-                        dtype=self.vae_output_dtype(),
-                        copy=True,
-                    )
-                    if pixel_samples is None:
-                        pixel_samples = torch.empty(
-                            (samples_in.shape[0],) + tuple(out.shape[1:]),
+            if should_tile:
+                do_tile = True
+            else:
+                for x in range(0, samples_in.shape[0], batch_number):
+                    samples = samples_in[x:x + batch_number].to(device=self.device, dtype=self.vae_dtype)
+                    if preallocated:
+                        self.first_stage_model.decode(samples, output_buffer=pixel_samples[x:x + batch_number], **vae_options)
+                    else:
+                        out = self.first_stage_model.decode(samples, **vae_options).to(
                             device=self.output_device,
                             dtype=self.vae_output_dtype(),
+                            copy=True,
                         )
-                    pixel_samples[x:x + batch_number].copy_(out)
-                    del out
-                self.process_output(pixel_samples[x:x + batch_number])
+                        if pixel_samples is None:
+                            pixel_samples = torch.empty(
+                                (samples_in.shape[0],) + tuple(out.shape[1:]),
+                                device=self.output_device,
+                                dtype=self.vae_output_dtype(),
+                            )
+                        pixel_samples[x:x + batch_number].copy_(out)
+                        del out
+                    self.process_output(pixel_samples[x:x + batch_number])
         except Exception as e:
             model_management.raise_non_oom(e)
             _LOG.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")