fix: apply CodeRabbit auto-fixes

Fixed 1 file(s) based on 1 unresolved review comment.

Co-authored-by: CodeRabbit <noreply@coderabbit.ai>

README.md

@@ -1,32 +1,41 @@
 # ComfyUI GPU Resident Loader
 
-A ComfyUI custom-node pack that targets **time-to-VRAM** and **sticky GPU residency**, not just lower peak host RAM.
+A ComfyUI custom-node pack for **faster time-to-VRAM**, **selective safetensors loading**, and **sticky GPU residency control**.
 
-It does two things:
+This repo does three related jobs:
 
-1. **Installs startup-time loader and residency patches** before any workflow nodes run.
-2. Ships **KJ-compatible loader nodes** for diffusion models and checkpoints, plus preload/pin/evict/report nodes for manual residency control.
+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.
+
+It is not just a “clean RAM” addon. The main target is the path from model file -> tensors -> live ComfyUI object -> VRAM retention.
 
 ## Why this exists
 
-Stock ComfyUI makes separate decisions for:
+ComfyUI’s default behavior mixes together two separate concerns:
 
-- where a model **lives after load**, and
-- where checkpoint tensors are **materialized first**.
+- **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 thing.
+Those are not the same problem.
 
-This repo targets the second problem directly for `.safetensors` by steering eligible loads toward direct GPU ingest, narrowing resident diffusion-model loads down to UNet tensors only, then targets the first problem by overriding offload policy and by teaching `free_memory()` to protect high-priority sticky entries until the VRAM budget says otherwise.
+This repo focuses on both:
 
-## What is included
+- 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.
 
-### Startup patcher
+## What changes at startup
 
-Installed automatically from `__init__.py` when the custom node loads.
+`__init__.py` calls `startup.install_patches()`, which applies the monkey patches exactly once when the custom node is imported.
+
+### 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`
@@ -35,77 +44,349 @@ Current patch surface:
 - `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`
 
-### Loader nodes
+### What those patches do
 
-- **Diffusion Model Loader Resident**
-- **Checkpoint Loader Resident**
-- **Checkpoint Model Loader Resident**
-- **Checkpoint Clip Loader Resident**
-- **Checkpoint VAE Loader Resident**
-- **Diffusion Model Selector Resident**
+#### 1) `load_torch_file` becomes residency-aware
 
-`Diffusion Model Loader Resident` mirrors the relevant KJ diffusion-loader feature surface:
+The patched loader:
 
-- weight dtype override
-- compute dtype override
-- cublas-ops toggle
-- SageAttention override
-- fp16 accumulation toggle
-- optional extra-state-dict merge
+- 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.
 
-On `.safetensors`, the resident diffusion-model path now reads only the detected UNet keys and merges only matching keys from any extra state dict. Repeated resident loads also reuse a live equivalent object when the source path and loader-relevant options still match.
+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.
 
-Before a GPU-bound resident load needs fresh VRAM, the loader now estimates the upcoming footprint from the model or checkpoint file itself and trims only enough lower-priority residency to cover that request plus a small internal headroom margin. The existing sticky ordering and partial-unload preference stay intact. That decision lives in the loader path, not in a separate user-configured VRAM target node.
+#### 2) Loader contexts are attached to stock ComfyUI load paths
 
-### Residency nodes
+These stock paths are wrapped with registry context and output binding:
 
-- **Set Global Residency Policy**
-- **Registry Snapshot**
-- **Pin Model/CLIP/VAE Residency**
-- **Preload Model/CLIP/VAE To GPU**
-- **Evict Model/CLIP/VAE From GPU**
-- **Report Model/CLIP/VAE Residency**
+- 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.
+
+#### 3) 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.
+
+#### 4) `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.
+
+#### 5) 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
+
+`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
+
+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.
+
+That behavior is not a general performance feature toggle. It exists to reduce avoidable VRAM spikes while sticky residency is active.
 
 ## Policies
 
-The startup patcher exposes four policies:
+The registry exposes four global policies:
 
-- `legacy` — leave ingest/offload behavior close to stock ComfyUI.
-- `balanced` — keep the registry and diagnostics, but do not aggressively steer ingest to GPU.
-- `prefer_gpu` — prefer GPU ingest, keep UNet/ControlNet/CLIP on the faster side of the device policy, but do not auto-pin tracked objects.
-- `sticky_gpu` — prefer GPU ingest, auto-pin the highest-value tracked outputs, and let lower-priority sticky entries yield first when VRAM pressure rises.
+### `legacy`
 
-Default selection order:
+Stay closest to stock ComfyUI behavior. Registry tracking still exists, but the patcher does not aggressively steer ingest/offload toward the GPU path.
 
-1. `COMFYUI_GPU_RESIDENT_POLICY` environment variable, if set.
-2. `sticky_gpu` when `--gpu_only` is active.
-3. `sticky_gpu` when `--highvram` is active.
+### `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`.
 
-## Important scope limits
+Supported values are:
 
-### Best path: `.safetensors`
+- `legacy`
+- `balanced`
+- `prefer_gpu`
+- `sticky_gpu`
 
-This repo is optimized around `.safetensors`.
+## Included nodes
 
-Direct GPU ingest is attempted for `.safetensors` loads. Resident diffusion-model loads take a header-first selective path and fetch only the detected UNet tensors instead of loading the whole file and filtering later. If the direct path fails, the patcher falls back to CPU read + GPU copy and records that fallback in the registry.
+All nodes live under the `GPU Resident Loader` category.
 
-### `.ckpt` / `.pt` remain CPU-first under PyTorch
+### Loader nodes
 
-Those formats still go through `torch.load()`. The repo tracks that path and can still keep the resulting model hot in VRAM, but it does **not** claim true direct-to-GPU checkpoint ingest for pickle-based formats. Resident checkpoint nodes now warn about this when a GPU-resident policy is active so the compatibility path is not mistaken for the fast path.
+#### Diffusion Model Selector Resident
 
-Use the included conversion helper to migrate hot models to `.safetensors`.
+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`.
+
+#### 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 the whole registry as formatted JSON.
+
+#### Pin Model Residency / Pin CLIP Residency / Pin VAE Residency
+
+Marks a tracked 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.
+
+#### Evict Model From GPU / Evict CLIP From GPU / Evict VAE From GPU
+
+Attempts to unload the selected 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.
+
+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, plus
+- 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.
+
+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.
+
+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,
+- or a recorded load failure.
+
+## What gets tracked
+
+Tracked/bound 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.
+
+## 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 contents alive after ComfyUI or WSL exits. CUDA memory lifetime is process/context scoped. Achieving persistence across process shutdown requires a long-lived keeper process or server that owns the CUDA context.
+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.
+
+### 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.
+
+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.
 
 ## Installation
 
@@ -115,67 +396,68 @@ Clone into `custom_nodes`:
 git clone https://github.com/xmarre/ComfyUI-GPU-Resident-Loader ComfyUI/custom_nodes/ComfyUI-GPU-Resident-Loader
 ```
 
-Install dependencies inside the same Python environment ComfyUI uses:
+Install dependencies into the same Python environment ComfyUI uses:
 
 ```bash
 pip install -r ComfyUI/custom_nodes/ComfyUI-GPU-Resident-Loader/requirements.txt
 ```
 
-Optional SageAttention dependencies are **not** installed by default. Install those separately if you plan to use the SageAttention loader modes.
+Requirements declared by the repo:
 
-## Basic usage
+- Python `>=3.10`
+- `safetensors>=0.4.3`
 
-### For direct diffusion-model loading
+Optional SageAttention dependencies are **not** installed by default. Install those separately if you plan to use a SageAttention mode in the resident loaders.
 
-Use **Diffusion Model Loader Resident**.
+## Basic usage patterns
 
-Recommended on a large VRAM machine:
+### 1) Large-VRAM, mostly resident workflow
 
-- policy: `sticky_gpu`
-- model format: `.safetensors`
-- preload with **Preload Model To GPU**
-- inspect with **Report Model Residency** or **Registry Snapshot**
+Recommended baseline:
 
-### For full checkpoints
+- 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**.
 
-Use **Checkpoint Loader Resident**.
+### 2) Full checkpoint workflow
 
-That tracks and binds the resulting diffusion model, CLIP, and VAE independently so they appear in the registry snapshot. If an equivalent live model, CLIP, or VAE already exists, the loader reuses it instead of rebuilding it.
+Use **Checkpoint Loader Resident** when you want `MODEL + CLIP + VAE` together.
 
-### For staged checkpoint loads
+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 residency control
 
 Use:
 
-- **Checkpoint Model Loader Resident** when the workflow only needs the diffusion model
-- **Checkpoint Clip Loader Resident** when the workflow only needs the text encoder
-- **Checkpoint VAE Loader Resident** when the workflow only needs the VAE
+- **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.
 
-The model-only checkpoint node takes the same selective safetensors UNet fast path as the diffusion-model loader. CLIP-only and VAE-only nodes still use ComfyUI's checkpoint construction logic, but they avoid materializing the other outputs and can reuse an already-live equivalent object.
+## Notes on compatibility and migration
 
-### For manual residency control
+### Legacy wiring: `extra_state_dict` used as a policy string
 
-- use **Pin ... Residency** to mark a tracked object sticky or evictable
-- use **Preload ... To GPU** to fully materialize it in VRAM immediately
-- use **Evict ... From GPU** to unload it from the current loaded-model set
+The resident diffusion-model loader contains a compatibility shim for older graphs:
 
-## Observability
+- 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.
 
-Every tracked load stores:
+New graphs should connect policy strings to **`policy_override`**, not to `extra_state_dict`.
 
-- source path
-- last load method
-- requested device
-- actual device
-- sticky flag
-- current loaded bytes
-- total bytes
-- current/offload/load device
+### Convert hot pickle checkpoints to safetensors
 
-That data is surfaced through the report nodes and the registry snapshot node.
-
-## Conversion helper
-
-`scripts/convert_checkpoint_to_safetensors.py` is included for one-time conversion of hot `.ckpt` / `.pt` files into `.safetensors`.
+`scripts/convert_checkpoint_to_safetensors.py` is included for one-time conversion of hot `.ckpt` / `.pt` / `.pth` style checkpoints.
 
 Example:
 
@@ -185,8 +467,13 @@ python ComfyUI/custom_nodes/ComfyUI-GPU-Resident-Loader/scripts/convert_checkpoi
   --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 intentionally stays GPL-compatible because it adapts behavior from GPL-licensed ComfyUI and mirrors feature behavior from the GPL-3.0-licensed KJNodes diffusion loader.
+This repo stays GPL-compatible because it adapts behavior from GPL-licensed ComfyUI and mirrors relevant loader behavior from GPL-3.0-licensed KJNodes.

cleanup.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import os
+from contextlib import contextmanager
 from typing import Any
 
 import comfy.model_management as model_management
 import torch
 
+from .external_residency import EXTERNAL_REGISTRY, ensure_external_integrations_installed
 from .residency import REGISTRY
 
 _ADAPTIVE_HEADROOM_RATIO = 0.125
@@ -33,6 +35,69 @@ def _safe_is_dead(loaded) -> bool:
         return False
 
 
+def _device_matches(device_a, device_b) -> bool:
+    normalized_a = _normalize_trim_device(device_a)
+    normalized_b = _normalize_trim_device(device_b)
+    if normalized_a is None or normalized_b is None:
+        return False
+    if isinstance(normalized_a, torch.device) and isinstance(normalized_b, torch.device):
+        if normalized_a.type != normalized_b.type:
+            return False
+        if normalized_a.type == "cuda":
+            index_a = 0 if normalized_a.index is None else normalized_a.index
+            index_b = 0 if normalized_b.index is None else normalized_b.index
+            return index_a == index_b
+        return True
+    return str(normalized_a) == str(normalized_b)
+
+
+def _should_force_cpu_offload(model: Any, *, active_device=None, force: bool = False) -> bool:
+    if force:
+        return True
+    active_device = _normalize_trim_device(active_device)
+    if not isinstance(active_device, torch.device) or active_device.type != "cuda":
+        return False
+    return _device_matches(active_device, getattr(model, "offload_device", None))
+
+
+@contextmanager
+def _temporary_offload_device(model: Any, device):
+    if model is None or device is None or not hasattr(model, "offload_device"):
+        yield
+        return
+
+    original_device = model.offload_device
+    model.offload_device = device
+    try:
+        yield
+    finally:
+        model.offload_device = original_device
+
+
+def unload_loaded_model(
+    loaded,
+    *,
+    active_device: str | torch.device | None = None,
+    force_offload_to_cpu: bool = False,
+    unpatch_weights: bool = True,
+) -> bool:
+    # This helper is intentionally scoped to full unloads owned by this plugin.
+    model = getattr(loaded, "model", None)
+    if model is None:
+        return False
+
+    active_device = _normalize_trim_device(active_device if active_device is not None else getattr(loaded, "device", None))
+    force_cpu_offload = _should_force_cpu_offload(
+        model,
+        active_device=active_device,
+        force=force_offload_to_cpu,
+    )
+    unload_target = torch.device("cpu") if force_cpu_offload else None
+
+    with _temporary_offload_device(model, unload_target):
+        return loaded.model_unload(None, unpatch_weights=unpatch_weights)
+
+
 def _sort_key_for_candidate(entry, *, sticky_respected: bool) -> tuple[int, int, float]:
     if not sticky_respected:
         return (0, 0, getattr(entry, "last_touched", 0.0) if entry is not None else 0.0)
@@ -65,8 +130,9 @@ def _trim_candidates(
     respect_sticky: bool,
     sticky_floor_priority: int,
     keep_models: tuple[Any, ...],
-) -> list[tuple[Any, Any, bool]]:
-    candidates: list[tuple[Any, Any, bool]] = []
+) -> list[tuple[Any, Any, bool, bool]]:
+    candidates: list[tuple[Any, Any, bool, bool]] = []
+    ensure_external_integrations_installed()
     for loaded in list(model_management.current_loaded_models):
         if device is not None and loaded.device != device:
             continue
@@ -86,7 +152,15 @@ def _trim_candidates(
             and bool(getattr(entry, "sticky", False))
             and int(getattr(entry, "priority", 0)) >= int(sticky_floor_priority)
         )
-        candidates.append((loaded, entry, sticky_respected))
+        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))
 
     candidates.sort(key=lambda item: _sort_key_for_candidate(item[1], sticky_respected=item[2]))
     return candidates
@@ -101,11 +175,13 @@ def trim_resident_vram(
     allow_partial_unload: bool,
     keep_models: tuple[Any, ...] = (),
 ) -> dict[str, Any]:
+    ensure_external_integrations_installed()
     cleanup_models_gc = getattr(model_management, "cleanup_models_gc", None)
     if callable(cleanup_models_gc):
         cleanup_models_gc()
 
     REGISTRY.refresh_runtime_state()
+    EXTERNAL_REGISTRY.refresh_runtime_state()
     device = _normalize_trim_device(device)
     if device is None:
         device = model_management.get_torch_device()
@@ -131,13 +207,14 @@ def trim_resident_vram(
             stopped_reason = "no_candidates"
             break
 
-        loaded, entry, sticky_respected = candidates[0]
-        model = loaded.model
-        loaded_before = int(loaded.model_loaded_memory())
+        candidate, entry, sticky_respected, is_external_candidate = candidates[0]
+        model = candidate if is_external_candidate else candidate.model
+        loaded_before = int(getattr(entry, "loaded_bytes", 0)) if is_external_candidate else int(candidate.model_loaded_memory())
         action = {
             "entry_id": getattr(entry, "entry_id", None),
             "basename": None if entry is None else os.path.basename(getattr(entry, "source_path", "") or ""),
             "tracked": entry is not None,
+            "external_candidate": bool(is_external_candidate),
             "sticky_respected": sticky_respected,
             "priority": None if entry is None else int(getattr(entry, "priority", 0)),
             "need_before_bytes": need,
@@ -145,7 +222,7 @@ def trim_resident_vram(
             "freed_pinned_ram_bytes": 0,
         }
 
-        if allow_partial_unload and hasattr(model, "pinned_memory_size") and hasattr(model, "partially_unload_ram"):
+        if (not is_external_candidate and allow_partial_unload and hasattr(model, "pinned_memory_size") and hasattr(model, "partially_unload_ram")):
             try:
                 pinned_memory = int(model.pinned_memory_size())
                 if pinned_memory > 0:
@@ -155,31 +232,42 @@ def trim_resident_vram(
             except Exception as exc:
                 action["pinned_ram_warning"] = str(exc)
 
-        try:
-            fully_unloaded = loaded.model_unload(need if allow_partial_unload else None)
-            action["mode"] = "full_unload" if fully_unloaded else "partial_unload"
-        except Exception as exc:
-            if allow_partial_unload:
-                try:
-                    fully_unloaded = loaded.model_unload(None)
-                    action["mode"] = "full_unload_fallback"
-                    action["partial_unload_warning"] = str(exc)
-                except Exception as fallback_exc:
-                    action["mode"] = "error"
-                    action["error"] = str(fallback_exc)
-                    actions.append(action)
-                    stopped_reason = "error"
-                    break
-            else:
+        if is_external_candidate:
+            try:
+                fully_unloaded = EXTERNAL_REGISTRY.evict(entry)
+                action["mode"] = "external_evict"
+            except Exception as exc:
                 action["mode"] = "error"
                 action["error"] = str(exc)
                 actions.append(action)
                 stopped_reason = "error"
                 break
-
-        if fully_unloaded:
+        else:
             try:
-                model_management.current_loaded_models.remove(loaded)
+                fully_unloaded = candidate.model_unload(need if allow_partial_unload else None)
+                action["mode"] = "full_unload" if fully_unloaded else "partial_unload"
+            except Exception as exc:
+                if allow_partial_unload:
+                    try:
+                        fully_unloaded = candidate.model_unload(None)
+                        action["mode"] = "full_unload_fallback"
+                        action["partial_unload_warning"] = str(exc)
+                    except Exception as fallback_exc:
+                        action["mode"] = "error"
+                        action["error"] = str(fallback_exc)
+                        actions.append(action)
+                        stopped_reason = "error"
+                        break
+                else:
+                    action["mode"] = "error"
+                    action["error"] = str(exc)
+                    actions.append(action)
+                    stopped_reason = "error"
+                    break
+
+        if fully_unloaded and not is_external_candidate:
+            try:
+                model_management.current_loaded_models.remove(candidate)
             except ValueError:
                 pass
 
@@ -187,8 +275,12 @@ def trim_resident_vram(
             soft_empty_cache()
 
         REGISTRY.refresh_runtime_state()
+        EXTERNAL_REGISTRY.refresh_runtime_state()
         free_after = _safe_free_memory(device)
-        action["loaded_after_bytes"] = 0 if fully_unloaded else int(loaded.model_loaded_memory())
+        if is_external_candidate:
+            action["loaded_after_bytes"] = 0 if fully_unloaded else int(getattr(entry, "loaded_bytes", 0))
+        else:
+            action["loaded_after_bytes"] = 0 if fully_unloaded else int(candidate.model_loaded_memory())
         action["freed_vram_bytes"] = max(0, free_after - free_now)
         action["free_after_bytes"] = free_after
         actions.append(action)

external_residency.py

@@ -0,0 +1,644 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+import contextlib
+import dataclasses
+import inspect
+import logging
+import os
+import sys
+import threading
+import time
+import weakref
+from typing import Any, Callable
+
+import torch
+
+from .residency import KIND_MODEL, KIND_VAE, REGISTRY
+
+_LOG = logging.getLogger(__name__)
+_SEEDVR2_PATCHED_CLASS_IDS: set[int] = set()
+_SEEDVR2_PATCHING_CLASS_IDS: set[int] = set()
+_SEEDVR2_PATCHING_THREAD_IDS: dict[int, int] = {}
+_SEEDVR2_PATCH_LOCK = threading.Lock()
+_SEEDVR2_PATCH_CONDITION = threading.Condition(_SEEDVR2_PATCH_LOCK)
+
+
+def _now() -> float:
+    return time.time()
+
+
+def _normalize_device(device: str | torch.device | None):
+    if device is None or isinstance(device, torch.device):
+        return device
+    try:
+        return torch.device(device)
+    except Exception:
+        return device
+
+
+def _device_matches(device_a, device_b) -> bool:
+    normalized_a = _normalize_device(device_a)
+    normalized_b = _normalize_device(device_b)
+    if normalized_a is None or normalized_b is None:
+        return False
+    if isinstance(normalized_a, torch.device) and isinstance(normalized_b, torch.device):
+        if normalized_a.type != normalized_b.type:
+            return False
+        if normalized_a.type == "cuda":
+            index_a = 0 if normalized_a.index is None else normalized_a.index
+            index_b = 0 if normalized_b.index is None else normalized_b.index
+            return index_a == index_b
+        return True
+    return str(normalized_a) == str(normalized_b)
+
+
+def _iter_seedvr2_wrapper_chain(model: Any):
+    if model is None:
+        return
+
+    stack = [model]
+    seen: set[int] = set()
+    while stack:
+        current = stack.pop()
+        if current is None:
+            continue
+        current_id = id(current)
+        if current_id in seen:
+            continue
+        seen.add(current_id)
+        yield current
+        for attr in ("_orig_mod", "dit_model"):
+            child = getattr(current, attr, None)
+            if child is not None:
+                stack.append(child)
+
+
+def _seedvr2_is_claimed(model: Any) -> bool:
+    return any(bool(getattr(current, "_seedvr2_cache_claimed", False)) for current in _iter_seedvr2_wrapper_chain(model))
+
+
+def _first_tensor_device(model: Any) -> str | None:
+    if model is None:
+        return None
+    try:
+        for param in model.parameters():
+            return str(param.device)
+    except Exception:
+        pass
+    try:
+        for buffer in model.buffers():
+            return str(buffer.device)
+    except Exception:
+        pass
+    return None
+
+
+def _unique_tensor_nbytes(model: Any) -> int:
+    if model is None:
+        return 0
+
+    total = 0
+    seen_storages: set[tuple[int, int, int]] = set()
+
+    def visit_tensor(tensor: torch.Tensor) -> None:
+        nonlocal total
+        if tensor is None:
+            return
+        try:
+            storage = tensor.untyped_storage()
+            key = (storage.data_ptr(), storage.nbytes(), int(tensor.device.index or 0) if tensor.device.type == "cuda" else -1)
+        except Exception:
+            key = (id(tensor), tensor.numel() * tensor.element_size(), -2)
+        if key in seen_storages:
+            return
+        seen_storages.add(key)
+        total += int(tensor.numel()) * int(tensor.element_size())
+
+    try:
+        for tensor in model.parameters():
+            visit_tensor(tensor)
+    except Exception:
+        pass
+    try:
+        for tensor in model.buffers():
+            visit_tensor(tensor)
+    except Exception:
+        pass
+    return int(total)
+
+
+def _seedvr2_entry_key(kind: str, node_id: Any) -> str:
+    return f"seedvr2:{kind}:{node_id}"
+
+
+def _seedvr2_source_path(kind: str, node_id: Any, config: dict[str, Any], model: Any) -> str:
+    model_name = config.get("model") or getattr(model, "_model_name", None) or f"node_{node_id}"
+    return f"seedvr2/{kind}/{node_id}/{model_name}"
+
+
+def _seedvr2_state_provider(
+    model_ref: Callable[[], Any | None],
+    config: dict[str, Any],
+) -> Callable[[], dict[str, Any]]:
+    def provider() -> dict[str, Any]:
+        model = model_ref()
+        if model is None:
+            return {}
+        total_bytes = _unique_tensor_nbytes(model)
+        return {
+            "current_device": _first_tensor_device(model),
+            "load_device": config.get("device"),
+            "offload_device": config.get("offload_device"),
+            "claimed": _seedvr2_is_claimed(model),
+            "loaded_bytes": total_bytes,
+            "total_bytes": total_bytes,
+        }
+
+    return provider
+
+
+def _coerce_external_bytes(value: Any, fallback: int, *, cache_key: str, field_name: str) -> int:
+    try:
+        return int(value)
+    except (TypeError, ValueError) as exc:
+        _LOG.debug(
+            "GPU Resident Loader: ignoring invalid external %s for %s: %r (%s)",
+            field_name,
+            cache_key,
+            value,
+            exc,
+        )
+        return int(fallback)
+
+
+@dataclasses.dataclass(slots=True)
+class ExternalResidencyEntry:
+    entry_id: str
+    cache_key: str
+    kind: str
+    source_path: str
+    sticky: bool
+    priority: int
+    created_at: float = dataclasses.field(default_factory=_now)
+    last_touched: float = dataclasses.field(default_factory=_now)
+    loaded_bytes: int = 0
+    total_bytes: int = 0
+    load_device: str | None = None
+    offload_device: str | None = None
+    current_device: str | None = None
+    claimed: bool = False
+    notes: list[str] = dataclasses.field(default_factory=list)
+    object_ref: weakref.ReferenceType[Any] | None = None
+    state_provider: Callable[[], dict[str, Any]] | None = None
+    evict_callback: Callable[[], bool] | None = None
+
+    def object(self) -> Any | None:
+        return None if self.object_ref is None else self.object_ref()
+
+    def is_alive(self) -> bool:
+        return self.object() is not None
+
+    def as_dict(self) -> dict[str, Any]:
+        return {
+            "entry_id": self.entry_id,
+            "cache_key": self.cache_key,
+            "kind": self.kind,
+            "source_path": self.source_path,
+            "basename": os.path.basename(self.source_path) if self.source_path else None,
+            "sticky": self.sticky,
+            "priority": self.priority,
+            "created_at": self.created_at,
+            "last_touched": self.last_touched,
+            "loaded_bytes": self.loaded_bytes,
+            "total_bytes": self.total_bytes,
+            "load_device": self.load_device,
+            "offload_device": self.offload_device,
+            "current_device": self.current_device,
+            "claimed": self.claimed,
+            "notes": list(self.notes),
+            "alive": self.is_alive(),
+            "external": True,
+        }
+
+
+class ExternalResidencyRegistry:
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self._entries: dict[str, ExternalResidencyEntry] = {}
+        self._cache_key_to_entry: dict[str, str] = {}
+        self._next_entry_seq = 1
+
+    def _make_entry_id(self, kind: str, source_path: str) -> str:
+        basename = os.path.basename(source_path) or "anonymous"
+        entry_id = f"external:{kind}:{basename}:{self._next_entry_seq}"
+        self._next_entry_seq += 1
+        return entry_id
+
+    def bind(
+        self,
+        *,
+        cache_key: str,
+        obj: Any,
+        kind: str,
+        source_path: str,
+        state_provider: Callable[[], dict[str, Any]],
+        evict_callback: Callable[[], bool],
+        sticky: bool = False,
+        priority: int | None = None,
+        note: str | None = None,
+    ) -> ExternalResidencyEntry:
+        if obj is None:
+            raise ValueError("Cannot bind None into external residency registry")
+
+        with self._lock:
+            entry_id = self._cache_key_to_entry.get(cache_key)
+            if entry_id is not None and entry_id in self._entries:
+                entry = self._entries[entry_id]
+            else:
+                entry_id = self._make_entry_id(kind, source_path)
+                entry = ExternalResidencyEntry(
+                    entry_id=entry_id,
+                    cache_key=cache_key,
+                    kind=kind,
+                    source_path=source_path,
+                    sticky=bool(sticky),
+                    priority=REGISTRY.default_priority(kind) if priority is None else int(priority),
+                )
+                self._entries[entry_id] = entry
+                self._cache_key_to_entry[cache_key] = entry_id
+
+            try:
+                entry.object_ref = weakref.ref(obj)
+            except TypeError:
+                entry.object_ref = None
+            entry.cache_key = cache_key
+            entry.kind = kind
+            entry.source_path = source_path
+            entry.sticky = bool(sticky)
+            entry.priority = REGISTRY.default_priority(kind) if priority is None else int(priority)
+            entry.state_provider = state_provider
+            entry.evict_callback = evict_callback
+            entry.last_touched = _now()
+            if note:
+                entry.notes.append(note)
+
+        self.refresh_runtime_state()
+        return entry
+
+    def remove(self, *, cache_key: str) -> bool:
+        with self._lock:
+            entry_id = self._cache_key_to_entry.pop(cache_key, None)
+            if entry_id is None:
+                return False
+            return self._entries.pop(entry_id, None) is not None
+
+    def refresh_runtime_state(self) -> None:
+        ensure_external_integrations_installed()
+        stale_keys: list[str] = []
+        with self._lock:
+            for cache_key, entry_id in list(self._cache_key_to_entry.items()):
+                entry = self._entries.get(entry_id)
+                if entry is None:
+                    stale_keys.append(cache_key)
+                    continue
+                obj = entry.object()
+                if obj is None:
+                    stale_keys.append(cache_key)
+                    continue
+                state_provider = entry.state_provider
+                if state_provider is None:
+                    continue
+                try:
+                    state = state_provider() or {}
+                except Exception as exc:
+                    _LOG.debug("GPU Resident Loader: failed to refresh external cache state for %s: %s", cache_key, exc)
+                    continue
+                entry.current_device = state.get("current_device")
+                entry.load_device = state.get("load_device")
+                entry.offload_device = state.get("offload_device")
+                entry.claimed = bool(state.get("claimed", False))
+                entry.loaded_bytes = _coerce_external_bytes(
+                    state.get("loaded_bytes", entry.loaded_bytes or 0),
+                    entry.loaded_bytes or 0,
+                    cache_key=cache_key,
+                    field_name="loaded_bytes",
+                )
+                entry.total_bytes = _coerce_external_bytes(
+                    state.get("total_bytes", entry.total_bytes or entry.loaded_bytes or 0),
+                    entry.total_bytes or entry.loaded_bytes or 0,
+                    cache_key=cache_key,
+                    field_name="total_bytes",
+                )
+            for cache_key in stale_keys:
+                entry_id = self._cache_key_to_entry.pop(cache_key, None)
+                if entry_id is not None:
+                    self._entries.pop(entry_id, None)
+
+    def candidates(
+        self,
+        *,
+        device: str | torch.device | None,
+        respect_sticky: bool,
+        sticky_floor_priority: int,
+        keep_models: tuple[Any, ...],
+    ) -> list[tuple[Any, ExternalResidencyEntry, bool]]:
+        self.refresh_runtime_state()
+        output: list[tuple[Any, ExternalResidencyEntry, bool]] = []
+        with self._lock:
+            for entry in self._entries.values():
+                obj = entry.object()
+                if obj is None:
+                    continue
+                if entry.claimed:
+                    continue
+                if device is not None and not _device_matches(device, entry.current_device):
+                    continue
+                if any(obj is keep for keep in keep_models if keep is not None):
+                    continue
+                sticky_respected = (
+                    respect_sticky
+                    and bool(entry.sticky)
+                    and int(entry.priority) >= int(sticky_floor_priority)
+                )
+                output.append((obj, entry, sticky_respected))
+        return output
+
+    def evict(self, entry: ExternalResidencyEntry) -> bool:
+        callback = entry.evict_callback
+        if callback is None:
+            return False
+        try:
+            result = bool(callback())
+        except Exception as exc:
+            _LOG.warning(
+                "GPU Resident Loader: external eviction callback failed for %s: %s",
+                entry.cache_key,
+                exc,
+            )
+            result = False
+        finally:
+            self.refresh_runtime_state()
+        return result
+
+    def snapshot(self) -> list[dict[str, Any]]:
+        self.refresh_runtime_state()
+        with self._lock:
+            items = [entry.as_dict() for entry in self._entries.values()]
+        items.sort(
+            key=lambda item: (
+                not item["sticky"],
+                item["kind"],
+                item["basename"] or "",
+            )
+        )
+        return items
+
+
+EXTERNAL_REGISTRY = ExternalResidencyRegistry()
+
+
+def _call_seedvr2_method_with_optional_expected_model(
+    method: Callable[..., Any],
+    *args: Any,
+    debug: Any = None,
+    expected_model: Any = None,
+) -> Any:
+    try:
+        parameters = inspect.signature(method).parameters
+    except (TypeError, ValueError):
+        parameters = {}
+
+    accepts_kwargs = any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values())
+    kwargs: dict[str, Any] = {}
+    if accepts_kwargs or "debug" in parameters:
+        kwargs["debug"] = debug
+    if expected_model is not None and (accepts_kwargs or "expected_model" in parameters):
+        kwargs["expected_model"] = expected_model
+    return method(*args, **kwargs)
+
+
+def _register_seedvr2_cached_model(global_cache: Any, *, kind: str, config: Any, model: Any) -> 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
+
+    node_id = config.get("node_id")
+    if node_id is None or model is None:
+        return
+
+    try:
+        model_ref = weakref.ref(model)
+    except TypeError:
+        def model_ref() -> Any | None:
+            return None
+
+    source_path = _seedvr2_source_path(kind, node_id, config, model)
+    cache_key = _seedvr2_entry_key(kind, node_id)
+    registry_kind = KIND_MODEL if kind == "dit" else KIND_VAE
+
+    if kind == "dit":
+        def evict_callback() -> bool:
+            return bool(
+                _call_seedvr2_method_with_optional_expected_model(
+                    global_cache.remove_dit,
+                    {"node_id": node_id},
+                    debug=None,
+                    expected_model=model_ref(),
+                )
+            )
+        note = f"SeedVR2 cached DiT node {node_id}"
+    else:
+        def evict_callback() -> bool:
+            return bool(
+                _call_seedvr2_method_with_optional_expected_model(
+                    global_cache.remove_vae,
+                    {"node_id": node_id},
+                    debug=None,
+                    expected_model=model_ref(),
+                )
+            )
+        note = f"SeedVR2 cached VAE node {node_id}"
+
+    EXTERNAL_REGISTRY.bind(
+        cache_key=cache_key,
+        obj=model,
+        kind=registry_kind,
+        source_path=source_path,
+        state_provider=_seedvr2_state_provider(model_ref, config),
+        evict_callback=evict_callback,
+        sticky=False,
+        priority=REGISTRY.default_priority(registry_kind),
+        note=note,
+    )
+
+
+def _install_seedvr2_integration_for_module(module: Any) -> bool:
+    model_cache_cls = getattr(module, "GlobalModelCache", None)
+    get_global_cache = getattr(module, "get_global_cache", None)
+    if model_cache_cls is None or not callable(get_global_cache):
+        return False
+
+    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
+
+    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
+
+        def set_dit_wrapper(self, dit_config, model, model_name, debug=None):
+            result = original_set_dit(self, dit_config, model, model_name, debug)
+            if result is not None:
+                _register_seedvr2_cached_model(self, kind="dit", config=dit_config, model=model)
+            return result
+
+        def set_vae_wrapper(self, vae_config, model, model_name, debug=None):
+            result = original_set_vae(self, vae_config, model, model_name, debug)
+            if result is not None:
+                _register_seedvr2_cached_model(self, kind="vae", config=vae_config, model=model)
+            return result
+
+        def replace_dit_wrapper(self, dit_config, model, debug=None, expected_model=None):
+            if original_replace_dit is None:
+                return False
+            result = _call_seedvr2_method_with_optional_expected_model(
+                original_replace_dit,
+                self,
+                dit_config,
+                model,
+                debug=debug,
+                expected_model=expected_model,
+            )
+            if result:
+                _register_seedvr2_cached_model(self, kind="dit", config=dit_config, model=model)
+            return result
+
+        def replace_vae_wrapper(self, vae_config, model, debug=None, expected_model=None):
+            if original_replace_vae is None:
+                return False
+            result = _call_seedvr2_method_with_optional_expected_model(
+                original_replace_vae,
+                self,
+                vae_config,
+                model,
+                debug=debug,
+                expected_model=expected_model,
+            )
+            if result:
+                _register_seedvr2_cached_model(self, kind="vae", config=vae_config, model=model)
+            return result
+
+        def remove_dit_wrapper(self, dit_config, debug=None, expected_model=None):
+            result = _call_seedvr2_method_with_optional_expected_model(
+                original_remove_dit,
+                self,
+                dit_config,
+                debug=debug,
+                expected_model=expected_model,
+            )
+            if result:
+                EXTERNAL_REGISTRY.remove(cache_key=_seedvr2_entry_key("dit", dit_config.get("node_id")))
+            return result
+
+        def remove_vae_wrapper(self, vae_config, debug=None, expected_model=None):
+            result = _call_seedvr2_method_with_optional_expected_model(
+                original_remove_vae,
+                self,
+                vae_config,
+                debug=debug,
+                expected_model=expected_model,
+            )
+            if result:
+                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:
+            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:
+        with _SEEDVR2_PATCH_CONDITION:
+            model_cache_cls.set_dit = original_set_dit
+            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
+            _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")
+    return True
+
+
+def ensure_external_integrations_installed() -> None:
+    for module in list(sys.modules.values()):
+        if module is None:
+            continue
+        module_file = getattr(module, "__file__", None)
+        if not module_file:
+            continue
+        normalized_file = os.path.abspath(module_file).replace("\\", "/")
+        if not normalized_file.endswith("/src/core/model_cache.py"):
+            continue
+        if "seedvr2" not in normalized_file.lower():
+            continue
+        if hasattr(module, "GlobalModelCache") and hasattr(module, "get_global_cache"):
+            try:
+                _install_seedvr2_integration_for_module(module)
+            except Exception as exc:
+                _LOG.warning(
+                    "GPU Resident Loader: failed to install SeedVR2 external integration from %s: %s",
+                    normalized_file,
+                    exc,
+                )

nodes.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import json
+import logging
 from typing import Any
 
 import comfy.model_management as model_management
 
+from .cleanup import unload_loaded_model
+from .external_residency import EXTERNAL_REGISTRY, ensure_external_integrations_installed
 from .kj_loader import (
     CheckpointClipLoaderResident,
     CheckpointLoaderResident,
@@ -15,6 +18,8 @@ from .kj_loader import (
 )
 from .residency import REGISTRY
 
+_LOG = logging.getLogger(__name__)
+
 
 def _entry_report_json(obj: Any) -> str:
     entry = REGISTRY.entry_for_object(obj)
@@ -59,9 +64,28 @@ def _evict_patcher(patcher, *, unpatch_weights: bool) -> bool:
         raise RuntimeError("Expected a patcher-capable object, but no patcher was found.")
     unloaded = False
     for loaded in list(model_management.current_loaded_models):
-        if loaded.model is patcher or loaded.model.is_clone(patcher):
-            loaded.model_unload(unpatch_weights=unpatch_weights)
-            unloaded = True
+        loaded_model = getattr(loaded, "model", None)
+        if loaded_model is None:
+            continue
+        safe_is_clone = False
+        if loaded_model is not patcher:
+            try:
+                safe_is_clone = loaded_model.is_clone(patcher)
+            except Exception as exc:
+                _LOG.warning("GPU Resident Loader: failed to evaluate clone state during eviction: %s", exc)
+        if loaded_model is patcher or safe_is_clone:
+            fully_unloaded = unload_loaded_model(
+                loaded,
+                active_device=getattr(loaded, "device", None),
+                force_offload_to_cpu=True,
+                unpatch_weights=unpatch_weights,
+            )
+            if fully_unloaded and unpatch_weights:
+                try:
+                    model_management.current_loaded_models.remove(loaded)
+                except ValueError:
+                    pass
+            unloaded = unloaded or fully_unloaded
     if unloaded and hasattr(model_management, "soft_empty_cache"):
         model_management.soft_empty_cache()
     REGISTRY.refresh_runtime_state()
@@ -100,7 +124,13 @@ class RegistrySnapshot:
     CATEGORY = "GPU Resident Loader/residency"
 
     def snapshot(self):
-        return (REGISTRY.snapshot_json(),)
+        ensure_external_integrations_installed()
+        payload = {
+            "policy": REGISTRY.get_policy(),
+            "entries": REGISTRY.snapshot(),
+            "external_entries": EXTERNAL_REGISTRY.snapshot(),
+        }
+        return (json.dumps(payload, indent=2, sort_keys=True),)
 
 
 class PinModelResidency:

patches.py

@@ -11,6 +11,13 @@ from typing import Any, Callable
 import torch
 from safetensors import safe_open
 
+from .cleanup import (
+    _device_matches as _shared_device_matches,
+    _should_force_cpu_offload,
+    adaptive_headroom_bytes,
+    trim_resident_vram,
+    unload_loaded_model,
+)
 from .residency import (
     KIND_CHECKPOINT,
     KIND_CLIP,
@@ -29,6 +36,8 @@ _METADATA_CPU_KEY_SUFFIXES = ("spiece_model", "tekken_model", "comfy_quant")
 _UNET_PREFIX_CANDIDATES = ("model.diffusion_model.", "model.model.", "net.")
 _WARNED_PICKLE_GPU_PATHS: set[str] = set()
 _SAFE_TENSORS_COMPONENT_CACHE_MAX = 32
+_STICKY_PROTECTION_VRAM_FLOOR_RATIO = 0.125
+_STICKY_PROTECTION_VRAM_FLOOR_CEIL_BYTES = 16 * 1024 ** 3
 _SAFETENSORS_DTYPE_MAP = {
     "BOOL": torch.bool,
     "U8": torch.uint8,
@@ -221,18 +230,10 @@ def _device_string(device: torch.device | None) -> str:
 
 
 def _devices_match(device_a: Any | None, device_b: Any | None) -> bool:
-    normalized_a = _normalize_device(device_a)
-    normalized_b = _normalize_device(device_b)
-    if normalized_a is not None and normalized_b is not None:
-        return normalized_a == normalized_b
-    return device_a == device_b
+    return _shared_device_matches(device_a, device_b)
 
 
 def _cpu_offload_required(model: Any, loaded_device: Any | None) -> bool:
-    offload_device = _normalize_device(getattr(model, "offload_device", None))
-    if offload_device is None or offload_device.type == "cpu":
-        return False
-
     current_device = None
     if hasattr(model, "current_loaded_device"):
         try:
@@ -241,10 +242,7 @@ def _cpu_offload_required(model: Any, loaded_device: Any | None) -> bool:
             current_device = None
     if current_device is None:
         current_device = _normalize_device(loaded_device)
-    if current_device is None or current_device.type == "cpu":
-        return False
-
-    return _devices_match(offload_device, current_device)
+    return _should_force_cpu_offload(model, active_device=current_device)
 
 
 @contextlib.contextmanager
@@ -408,6 +406,37 @@ def _warn_pickle_gpu_compatibility(path: str, requested_device: torch.device) ->
     )
 
 
+def _sticky_protection_target(memory_required: int, device: Any) -> int:
+    import comfy.model_management as model_management
+
+    required = max(0, int(memory_required))
+    target = required + adaptive_headroom_bytes(required)
+
+    minimum_inference_memory = getattr(model_management, "minimum_inference_memory", None)
+    if callable(minimum_inference_memory):
+        try:
+            target = max(target, int(minimum_inference_memory()))
+        except Exception:
+            pass
+
+    get_total_memory = getattr(model_management, "get_total_memory", None)
+    if callable(get_total_memory):
+        try:
+            total_memory = int(get_total_memory(device))
+        except Exception:
+            total_memory = 0
+        if total_memory > 0:
+            target = max(
+                target,
+                min(
+                    _STICKY_PROTECTION_VRAM_FLOOR_CEIL_BYTES,
+                    int(total_memory * _STICKY_PROTECTION_VRAM_FLOOR_RATIO),
+                ),
+            )
+
+    return target
+
+
 def _resolved_context(kind: str, source_path: str | None) -> tuple[torch.device | None, str, str | None]:
     ctx = REGISTRY.current_context()
     if ctx is not None and ctx.explicit_device is not None:
@@ -612,9 +641,224 @@ def _wrap_load_clip(func: Callable[..., Any]) -> Callable[..., Any]:
     return wrapper
 
 
+def _sticky_safe_batch_number(*, batch_count: int, free_memory: int, memory_used: int, device: Any) -> int:
+    batches = max(1, int(max(0, int(free_memory)) / max(1, int(memory_used))))
+    batches = min(max(1, int(batch_count)), batches)
+    if REGISTRY.get_policy() != "sticky_gpu" or device is None:
+        return batches
+
+    reserve = max(0, _sticky_protection_target(memory_used, device) - max(0, int(memory_used)))
+    safe_budget = max(0, int(free_memory) - reserve)
+    safe_batches = max(1, int(safe_budget / max(1, int(memory_used))))
+    capped = min(batches, safe_batches)
+    if capped < batches:
+        _LOG.debug(
+            "GPU Resident Loader: capped VAE batch from %s to %s to preserve %s bytes of transient headroom.",
+            batches,
+            capped,
+            reserve,
+        )
+    return max(1, capped)
+
+
+def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]:
+    @functools.wraps(func)
+    def wrapper(self, pixel_samples):
+        if REGISTRY.get_policy() != "sticky_gpu":
+            return func(self, pixel_samples)
+
+        import comfy.model_management as model_management
+
+        self.throw_exception_if_invalid()
+        pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
+        pixel_samples = pixel_samples.movedim(-1, 1)
+        do_tile = False
+        if self.latent_dim == 3 and pixel_samples.ndim < 5:
+            if not self.not_video:
+                pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
+            else:
+                pixel_samples = pixel_samples.unsqueeze(2)
+        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,
+                device=self.device,
+            )
+            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.")
+            do_tile = True
+
+        if do_tile:
+            model_management.soft_empty_cache()
+            if self.latent_dim == 3:
+                tile = 256
+                overlap = tile // 4
+                samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+            elif self.latent_dim == 1 or self.extra_1d_channel is not None:
+                samples = self.encode_tiled_1d(pixel_samples)
+            else:
+                samples = self.encode_tiled_(pixel_samples)
+
+        return samples
+
+    return wrapper
+
+
+def _wrap_vae_decode(func: Callable[..., Any]) -> Callable[..., Any]:
+    @functools.wraps(func)
+    def wrapper(self, samples_in, vae_options={}):
+        if REGISTRY.get_policy() != "sticky_gpu":
+            return func(self, samples_in, vae_options)
+
+        import comfy.model_management as model_management
+
+        self.throw_exception_if_invalid()
+        pixel_samples = None
+        do_tile = False
+        if self.latent_dim == 2 and samples_in.ndim == 5:
+            samples_in = samples_in[:, :, 0]
+        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,
+                device=self.device,
+            )
+
+            preallocated = False
+            if 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,
+                    dtype=self.vae_output_dtype(),
+                )
+                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:]),
+                            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.")
+            do_tile = True
+
+        if do_tile:
+            model_management.soft_empty_cache()
+            dims = samples_in.ndim - 2
+            if dims == 1 or self.extra_1d_channel is not None:
+                pixel_samples = self.decode_tiled_1d(samples_in)
+            elif dims == 2:
+                pixel_samples = self.decode_tiled_(samples_in)
+            elif dims == 3:
+                tile = 256 // self.spacial_compression_decode()
+                overlap = tile // 4
+                pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+
+        pixel_samples = pixel_samples.to(self.output_device).movedim(1, -1)
+        return pixel_samples
+
+    return wrapper
+
+
 def _wrap_load_models_gpu(func: Callable[..., Any]) -> Callable[..., Any]:
     @functools.wraps(func)
     def wrapper(models, *args, **kwargs):
+        import comfy.model_management as model_management
+
+        requested_models = set()
+        for model in list(models):
+            requested_models.add(model)
+            for additional in model.model_patches_models():
+                requested_models.add(additional)
+
+        clone_conflicts: list[Any] = []
+        seen_loaded_ids: set[int] = set()
+        for requested in requested_models:
+            is_clone = getattr(requested, "is_clone", None)
+            if not callable(is_clone):
+                continue
+            for loaded in list(model_management.current_loaded_models):
+                try:
+                    dead = loaded.is_dead()
+                except Exception:
+                    dead = False
+                if id(loaded) in seen_loaded_ids or dead:
+                    continue
+                loaded_model = getattr(loaded, "model", None)
+                if loaded_model is None or loaded_model is requested:
+                    continue
+                try:
+                    if not requested.is_clone(loaded_model):
+                        continue
+                except Exception as exc:
+                    raise RuntimeError(
+                        "GPU Resident Loader: failed to evaluate clone-conflict state before replacement"
+                    ) from exc
+                clone_conflicts.append(loaded)
+                seen_loaded_ids.add(id(loaded))
+
+        clone_conflicts_unloaded = 0
+        try:
+            for loaded in clone_conflicts:
+                # ComfyUI's built-in clone replacement pops the wrapper and only calls detach(False),
+                # which does not unpatch base weights. Fully unload before replacement or fail closed.
+                if not unload_loaded_model(
+                    loaded,
+                    active_device=getattr(loaded, "device", None),
+                    force_offload_to_cpu=True,
+                ):
+                    raise RuntimeError("GPU Resident Loader: failed to fully unload a clone-conflict wrapper before replacement")
+                try:
+                    model_management.current_loaded_models.remove(loaded)
+                except ValueError:
+                    pass
+                clone_conflicts_unloaded += 1
+        finally:
+            if clone_conflicts_unloaded > 0:
+                if hasattr(model_management, "soft_empty_cache"):
+                    model_management.soft_empty_cache()
+                REGISTRY.refresh_runtime_state()
+
         result = func(models, *args, **kwargs)
         for model in list(models):
             REGISTRY.touch(model)
@@ -668,12 +912,30 @@ def _wrap_free_memory(func: Callable[..., Any]) -> Callable[..., Any]:
         import comfy.model_management as model_management
 
         keep_loaded = list(keep_loaded or [])
-        sticky_wrappers = []
-        if REGISTRY.get_policy() == "sticky_gpu":
-            sticky_wrappers = [w for w in REGISTRY.sticky_loaded_wrappers(device) if w not in keep_loaded]
-
         protected_wrappers: list[Any] = []
+        sticky_wrappers: list[Any] = []
+        if REGISTRY.get_policy() == "sticky_gpu" and device is not None:
+            sticky_wrappers = [w for w in REGISTRY.sticky_loaded_wrappers(device) if w not in keep_loaded]
         if device is not None and sticky_wrappers:
+            protection_target = _sticky_protection_target(memory_required, device)
+            keep_models = tuple(
+                model
+                for model in (getattr(loaded_wrapper, "model", None) for loaded_wrapper in keep_loaded)
+                if model is not None
+            )
+            try:
+                trim_resident_vram(
+                    device=device,
+                    target_free_vram_bytes=protection_target,
+                    respect_sticky=True,
+                    sticky_floor_priority=0,
+                    allow_partial_unload=True,
+                    keep_models=keep_models,
+                )
+            except Exception as exc:
+                _LOG.debug("GPU Resident Loader: sticky pre-trim failed for free_memory(%s): %s", memory_required, exc)
+
+            sticky_wrappers = [w for w in REGISTRY.sticky_loaded_wrappers(device) if w not in keep_loaded]
             try:
                 free_now = model_management.get_free_memory(device)
             except Exception:
@@ -687,7 +949,7 @@ def _wrap_free_memory(func: Callable[..., Any]) -> Callable[..., Any]:
                         unloadable_wrappers.append(loaded)
                 available_for_protection = max(
                     0,
-                    free_now + sum(max(0, loaded.model_loaded_memory()) for loaded in unloadable_wrappers) - memory_required,
+                    free_now + sum(max(0, loaded.model_loaded_memory()) for loaded in unloadable_wrappers) - protection_target,
                 )
                 protected_memory = 0
                 for loaded in sticky_wrappers:
@@ -814,6 +1076,14 @@ def install_patches() -> None:
     if comfy_sd.load_clip is original_load_clip:
         comfy_sd.load_clip = _wrap_load_clip(original_load_clip)
 
+    original_vae_encode = _remember_original("sd.VAE.encode", comfy_sd.VAE.encode)
+    if comfy_sd.VAE.encode is original_vae_encode:
+        comfy_sd.VAE.encode = _wrap_vae_encode(original_vae_encode)
+
+    original_vae_decode = _remember_original("sd.VAE.decode", comfy_sd.VAE.decode)
+    if comfy_sd.VAE.decode is original_vae_decode:
+        comfy_sd.VAE.decode = _wrap_vae_decode(original_vae_decode)
+
     original_clip_vision_load = _remember_original("clip_vision.load", clip_vision.load)
     if clip_vision.load is original_clip_vision_load:
         clip_vision.load = _wrap_with_load_context(

residency.py

@@ -525,7 +525,6 @@ class ResidencyRegistry:
                 entry.loaded_bytes = int(loaded.model_loaded_memory())
                 entry.total_bytes = int(loaded.model_memory())
                 entry.current_device = str(loaded.device)
-                entry.last_touched = _now()
 
     def snapshot(self) -> list[dict[str, Any]]:
         self.refresh_runtime_state()