Preserve sticky VAE headroom

[codex] Fix sticky reclaim before Face Detailer VAE encode

README.md

@@ -16,7 +16,7 @@ Stock ComfyUI makes separate decisions for:
 
 Those are not the same thing.
 
-This repo targets the second problem directly for `.safetensors` by steering eligible loads toward direct GPU ingest, then targets the first problem by overriding offload policy and by teaching `free_memory()` to respect sticky entries until the VRAM budget is actually exceeded.
+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.
 
 ## What is included
 
@@ -46,6 +46,9 @@ Current patch surface:
 
 - **Diffusion Model Loader Resident**
 - **Checkpoint Loader Resident**
+- **Checkpoint Model Loader Resident**
+- **Checkpoint Clip Loader Resident**
+- **Checkpoint VAE Loader Resident**
 - **Diffusion Model Selector Resident**
 
 `Diffusion Model Loader Resident` mirrors the relevant KJ diffusion-loader feature surface:
@@ -57,6 +60,10 @@ Current patch surface:
 - fp16 accumulation toggle
 - optional extra-state-dict merge
 
+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.
+
+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.
+
 ### Residency nodes
 
 - **Set Global Residency Policy**
@@ -72,8 +79,8 @@ The startup patcher exposes four 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 and GPU offload devices, but do not auto-pin tracked objects.
-- `sticky_gpu` — prefer GPU ingest, prefer GPU offload devices, and auto-mark tracked loader outputs sticky.
+- `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.
 
 Default selection order:
 
@@ -88,11 +95,11 @@ Default selection order:
 
 This repo is optimized around `.safetensors`.
 
-Direct GPU ingest is attempted for `.safetensors` loads. If the direct path fails, the patcher falls back to CPU read + GPU copy and records that fallback in the registry.
+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.
 
 ### `.ckpt` / `.pt` remain CPU-first under PyTorch
 
-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.
+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.
 
 Use the included conversion helper to migrate hot models to `.safetensors`.
 
@@ -133,7 +140,17 @@ Recommended on a large VRAM machine:
 
 Use **Checkpoint Loader Resident**.
 
-That tracks and binds the resulting diffusion model, CLIP, and VAE independently so they appear in the registry snapshot.
+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.
+
+### For staged checkpoint loads
+
+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
+
+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.
 
 ### For manual residency control
 

cleanup.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from typing import Any
+
+import comfy.model_management as model_management
+import torch
+
+from .residency import REGISTRY
+
+_ADAPTIVE_HEADROOM_RATIO = 0.125
+_ADAPTIVE_HEADROOM_FLOOR_BYTES = 256 * 1024 * 1024
+_ADAPTIVE_HEADROOM_CEIL_BYTES = 1024 * 1024 * 1024
+
+
+def _safe_free_memory(device) -> int:
+    return int(model_management.get_free_memory(device))
+
+
+def _normalize_trim_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 _safe_is_dead(loaded) -> bool:
+    try:
+        return loaded.is_dead()
+    except Exception:
+        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)
+    priority = int(getattr(entry, "priority", 0)) if entry is not None else 0
+    last_touched = float(getattr(entry, "last_touched", 0.0)) if entry is not None else 0.0
+    return (1, priority, last_touched)
+
+
+def _should_keep_loaded_model(model: Any, keep_models: tuple[Any, ...]) -> bool:
+    if not keep_models:
+        return False
+    for keep in keep_models:
+        if keep is None:
+            continue
+        if model is keep:
+            return True
+        is_clone = getattr(model, "is_clone", None)
+        if callable(is_clone):
+            try:
+                if is_clone(keep):
+                    return True
+            except Exception:
+                pass
+    return False
+
+
+def _trim_candidates(
+    *,
+    device,
+    respect_sticky: bool,
+    sticky_floor_priority: int,
+    keep_models: tuple[Any, ...],
+) -> list[tuple[Any, Any, bool]]:
+    candidates: list[tuple[Any, Any, bool]] = []
+    for loaded in list(model_management.current_loaded_models):
+        if device is not None and loaded.device != device:
+            continue
+        if _safe_is_dead(loaded):
+            continue
+
+        model = getattr(loaded, "model", None)
+        if model is None:
+            continue
+        if _should_keep_loaded_model(model, keep_models):
+            continue
+
+        entry = REGISTRY.entry_for_object(model)
+        sticky_respected = (
+            respect_sticky
+            and entry is not None
+            and bool(getattr(entry, "sticky", False))
+            and int(getattr(entry, "priority", 0)) >= int(sticky_floor_priority)
+        )
+        candidates.append((loaded, entry, sticky_respected))
+
+    candidates.sort(key=lambda item: _sort_key_for_candidate(item[1], sticky_respected=item[2]))
+    return candidates
+
+
+def trim_resident_vram(
+    *,
+    device: str | torch.device | None = None,
+    target_free_vram_bytes: int,
+    respect_sticky: bool,
+    sticky_floor_priority: int,
+    allow_partial_unload: bool,
+    keep_models: tuple[Any, ...] = (),
+) -> dict[str, Any]:
+    cleanup_models_gc = getattr(model_management, "cleanup_models_gc", None)
+    if callable(cleanup_models_gc):
+        cleanup_models_gc()
+
+    REGISTRY.refresh_runtime_state()
+    device = _normalize_trim_device(device)
+    if device is None:
+        device = model_management.get_torch_device()
+    free_before = _safe_free_memory(device)
+    actions: list[dict[str, Any]] = []
+    soft_empty_cache = getattr(model_management, "soft_empty_cache", None)
+    stopped_reason = "target_met"
+
+    while True:
+        free_now = _safe_free_memory(device)
+        need = int(target_free_vram_bytes) - free_now
+        if need <= 0:
+            stopped_reason = "target_met"
+            break
+
+        candidates = _trim_candidates(
+            device=device,
+            respect_sticky=respect_sticky,
+            sticky_floor_priority=sticky_floor_priority,
+            keep_models=keep_models,
+        )
+        if not candidates:
+            stopped_reason = "no_candidates"
+            break
+
+        loaded, entry, sticky_respected = candidates[0]
+        model = loaded.model
+        loaded_before = int(loaded.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,
+            "sticky_respected": sticky_respected,
+            "priority": None if entry is None else int(getattr(entry, "priority", 0)),
+            "need_before_bytes": need,
+            "loaded_before_bytes": loaded_before,
+            "freed_pinned_ram_bytes": 0,
+        }
+
+        if 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:
+                    pinned_budget = min(pinned_memory, max(need, 0))
+                    model.partially_unload_ram(pinned_budget)
+                    action["freed_pinned_ram_bytes"] = pinned_budget
+            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:
+                action["mode"] = "error"
+                action["error"] = str(exc)
+                actions.append(action)
+                stopped_reason = "error"
+                break
+
+        if fully_unloaded:
+            try:
+                model_management.current_loaded_models.remove(loaded)
+            except ValueError:
+                pass
+
+        if callable(soft_empty_cache):
+            soft_empty_cache()
+
+        REGISTRY.refresh_runtime_state()
+        free_after = _safe_free_memory(device)
+        action["loaded_after_bytes"] = 0 if fully_unloaded else int(loaded.model_loaded_memory())
+        action["freed_vram_bytes"] = max(0, free_after - free_now)
+        action["free_after_bytes"] = free_after
+        actions.append(action)
+
+        if action["freed_vram_bytes"] <= 0 and not fully_unloaded:
+            stopped_reason = "no_progress"
+            break
+
+    free_after = _safe_free_memory(device)
+    target_met = free_after >= int(target_free_vram_bytes)
+    if target_met:
+        stopped_reason = "target_met"
+    return {
+        "status": "met_target" if target_met else ("error" if stopped_reason == "error" else "partial"),
+        "stopped_reason": stopped_reason,
+        "target_met": target_met,
+        "device": str(device),
+        "target_free_vram_bytes": int(target_free_vram_bytes),
+        "free_before_bytes": free_before,
+        "free_after_bytes": free_after,
+        "freed_vram_bytes": max(0, free_after - free_before),
+        "respect_sticky": bool(respect_sticky),
+        "sticky_floor_priority": int(sticky_floor_priority),
+        "allow_partial_unload": bool(allow_partial_unload),
+        "actions": actions,
+    }
+
+
+def adaptive_headroom_bytes(required_bytes: int) -> int:
+    required = max(0, int(required_bytes))
+    if required == 0:
+        return 0
+    return min(
+        _ADAPTIVE_HEADROOM_CEIL_BYTES,
+        max(_ADAPTIVE_HEADROOM_FLOOR_BYTES, int(required * _ADAPTIVE_HEADROOM_RATIO)),
+    )
+
+
+def trim_resident_vram_for_load(
+    *,
+    required_bytes: int,
+    reason: str,
+    device: str | torch.device | None = None,
+    respect_sticky: bool = True,
+    sticky_floor_priority: int = 0,
+    allow_partial_unload: bool = True,
+    keep_models: tuple[Any, ...] = (),
+) -> dict[str, Any]:
+    estimated_load_bytes = max(0, int(required_bytes))
+    headroom_bytes = adaptive_headroom_bytes(estimated_load_bytes)
+    target_free_vram_bytes = estimated_load_bytes + headroom_bytes
+
+    report = trim_resident_vram(
+        device=device,
+        target_free_vram_bytes=target_free_vram_bytes,
+        respect_sticky=respect_sticky,
+        sticky_floor_priority=sticky_floor_priority,
+        allow_partial_unload=allow_partial_unload,
+        keep_models=keep_models,
+    )
+    report["trim_strategy"] = "adaptive_load_request"
+    report["trim_reason"] = str(reason)
+    report["estimated_load_bytes"] = estimated_load_bytes
+    report["adaptive_headroom_bytes"] = headroom_bytes
+    report["kept_loaded_models"] = len([model for model in keep_models if model is not None])
+    return report

kj_loader.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import json
 import logging
 import os
 from typing import Any
@@ -13,6 +14,14 @@ import comfy.utils
 from comfy.cli_args import PerformanceFeature, args
 from comfy.ldm.modules.attention import attention_pytorch, wrap_attn
 
+from .cleanup import trim_resident_vram_for_load
+from .patches import (
+    checkpoint_component_info_from_header,
+    estimate_checkpoint_component_bytes,
+    estimate_safetensors_tensor_bytes,
+    infer_unet_prefix_from_keys,
+    load_safetensors_state_dict,
+)
 from .residency import KIND_CHECKPOINT, KIND_CLIP, KIND_MODEL, KIND_VAE, POLICIES, REGISTRY
 
 
@@ -37,6 +46,13 @@ DTYPE_MAP = {
     "fp32": torch.float32,
 }
 
+UNET_PREFIX_CANDIDATES = (
+    "model.diffusion_model.",
+    "model.model.",
+    "net.",
+    "model.",
+)
+
 
 def _set_cublas_linear(enabled: bool) -> None:
     if enabled:
@@ -252,6 +268,16 @@ def _normalize_optional_string(value: str | None) -> str | None:
     return normalized or None
 
 
+def _effective_policy_name(policy_override: str | None) -> str:
+    return REGISTRY.get_policy() if policy_override is None else policy_override
+
+
+def _make_loader_key(loader_name: str, **payload: Any) -> str:
+    normalized_payload = {"loader": loader_name}
+    normalized_payload.update(payload)
+    return json.dumps(normalized_payload, sort_keys=True, separators=(",", ":"))
+
+
 def _normalize_optional_policy(value: str | None) -> str | None:
     normalized = _normalize_optional_string(value)
     return None if normalized is None else normalized.lower()
@@ -291,6 +317,229 @@ def _resolve_loader_policy_and_extra_state_dict(
     return normalized_policy, normalized_extra
 
 
+def _is_safetensors_path(path: str) -> bool:
+    lowered = path.lower()
+    return lowered.endswith(".safetensors") or lowered.endswith(".sft")
+
+
+def _warn_pickle_checkpoint_gpu_compatibility(loader_name: str, path: str) -> None:
+    if _is_safetensors_path(path):
+        return
+    if not REGISTRY.wants_gpu_ingest(KIND_MODEL):
+        return
+    _LOG.warning(
+        "%s: %s is using the compatibility path through torch.load() on CPU before tensor-by-tensor copies to GPU. "
+        "Convert hot checkpoints to safetensors with scripts/convert_checkpoint_to_safetensors.py for the real fast path.",
+        loader_name,
+        path,
+    )
+
+
+def _fallback_file_size_bytes(path: str) -> int:
+    try:
+        return int(os.path.getsize(path))
+    except OSError:
+        return 0
+
+
+def _weight_dtype_override(weight_dtype: str) -> torch.dtype | None:
+    if weight_dtype == "fp8_e4m3fn_fast":
+        return torch.float8_e4m3fn
+    return DTYPE_MAP.get(weight_dtype)
+
+
+def _normalize_keep_models(*objects: Any) -> tuple[Any, ...]:
+    keep_models: list[Any] = []
+    for obj in objects:
+        if obj is None:
+            continue
+        patcher = getattr(obj, "patcher", None)
+        keep = patcher if patcher is not None else obj
+        if not any(existing is keep for existing in keep_models):
+            keep_models.append(keep)
+    return tuple(keep_models)
+
+
+def _estimate_extra_state_dict_bytes(extra_state_dict: str | None, *, weight_dtype: str) -> int:
+    if not extra_state_dict:
+        return 0
+    dtype_override = _weight_dtype_override(weight_dtype)
+    if _is_safetensors_path(extra_state_dict):
+        estimated = estimate_safetensors_tensor_bytes(extra_state_dict, dtype_override=dtype_override)
+        if estimated is not None:
+            return int(estimated)
+    return _fallback_file_size_bytes(extra_state_dict)
+
+
+def _estimate_model_load_bytes(
+    source_path: str,
+    *,
+    cache_scope: str,
+    weight_dtype: str,
+    extra_state_dict: str | None,
+) -> int:
+    dtype_override = _weight_dtype_override(weight_dtype)
+    estimated = None
+    if _is_safetensors_path(source_path):
+        if cache_scope == "checkpoint_model":
+            estimated = estimate_checkpoint_component_bytes(
+                source_path,
+                KIND_MODEL,
+                dtype_override=dtype_override,
+            )
+        else:
+            estimated = estimate_safetensors_tensor_bytes(source_path, dtype_override=dtype_override)
+    if estimated is None:
+        estimated = _fallback_file_size_bytes(source_path)
+    return int(estimated) + _estimate_extra_state_dict_bytes(extra_state_dict, weight_dtype=weight_dtype)
+
+
+def _estimate_checkpoint_aux_component_bytes(ckpt_path: str, *, kind: str) -> int:
+    estimated = None
+    if _is_safetensors_path(ckpt_path):
+        estimated = estimate_checkpoint_component_bytes(ckpt_path, kind)
+    if estimated is None:
+        estimated = _fallback_file_size_bytes(ckpt_path)
+    return int(estimated)
+
+
+def _maybe_trim_before_load(
+    *,
+    loader_name: str,
+    reason: str,
+    explicit_device: torch.device | None,
+    required_bytes: int,
+    keep_models: tuple[Any, ...] = (),
+) -> None:
+    if explicit_device is None or explicit_device.type != "cuda":
+        return
+    if required_bytes <= 0:
+        return
+
+    report = trim_resident_vram_for_load(
+        required_bytes=required_bytes,
+        reason=reason,
+        device=explicit_device,
+        keep_models=keep_models,
+    )
+    if report["freed_vram_bytes"] > 0:
+        _LOG.info(
+            "%s: adaptively freed %.2f GiB before load (%s, estimated %.2f GiB + %.2f GiB headroom)",
+            loader_name,
+            report["freed_vram_bytes"] / (1024 ** 3),
+            report["stopped_reason"],
+            report["estimated_load_bytes"] / (1024 ** 3),
+            report["adaptive_headroom_bytes"] / (1024 ** 3),
+        )
+    if not report["target_met"]:
+        _LOG.warning(
+            "%s: adaptive trim could not reach the estimated headroom for %s (%s, free %.2f GiB / target %.2f GiB)",
+            loader_name,
+            reason,
+            report["stopped_reason"],
+            report["free_after_bytes"] / (1024 ** 3),
+            report["target_free_vram_bytes"] / (1024 ** 3),
+        )
+
+
+def _selected_unet_key_map_from_header(
+    path: str,
+    *,
+    known_unet_keys: set[str] | None = None,
+) -> tuple[dict[str, str], str | None]:
+    from safetensors import safe_open
+
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        all_keys = list(handle.keys())
+
+    prefix = infer_unet_prefix_from_keys(all_keys)
+    if known_unet_keys is None:
+        selected = {key: key[len(prefix):] for key in all_keys if key.startswith(prefix)}
+        return (selected or {key: key for key in all_keys}), prefix
+
+    prefixes_to_try: list[str] = []
+    for candidate in (prefix, *UNET_PREFIX_CANDIDATES):
+        if candidate and candidate not in prefixes_to_try:
+            prefixes_to_try.append(candidate)
+
+    selected: dict[str, str] = {}
+    for key in all_keys:
+        if key in known_unet_keys:
+            selected[key] = key
+            continue
+        for candidate in prefixes_to_try:
+            if key.startswith(candidate):
+                stripped = key[len(candidate):]
+                if stripped in known_unet_keys:
+                    selected[key] = stripped
+                    break
+    return selected, prefix
+
+
+def _load_matching_extra_unet_state_dict(
+    extra_state_dict_path: str,
+    *,
+    requested_device: torch.device,
+    known_unet_keys: set[str],
+) -> dict[str, Any]:
+    if _is_safetensors_path(extra_state_dict_path):
+        selected_keys, _ = _selected_unet_key_map_from_header(
+            extra_state_dict_path,
+            known_unet_keys=known_unet_keys,
+        )
+        extra_sd, _, _, _ = load_safetensors_state_dict(
+            extra_state_dict_path,
+            requested_device,
+            selected_keys=selected_keys,
+        )
+        return extra_sd
+
+    extra_sd = comfy.utils.load_torch_file(extra_state_dict_path)
+    return _extract_unet_state_dict(extra_sd, known_unet_keys=known_unet_keys)
+
+
+def _extract_unet_state_dict(
+    sd: dict[str, Any],
+    *,
+    diffusion_model_prefix: str | None = None,
+    known_unet_keys: set[str] | None = None,
+) -> dict[str, Any]:
+    if known_unet_keys is None:
+        if diffusion_model_prefix is None:
+            diffusion_model_prefix = comfy.sd.model_detection.unet_prefix_from_state_dict(sd)
+        if diffusion_model_prefix:
+            prefix_len = len(diffusion_model_prefix)
+            extracted = {key[prefix_len:]: value for key, value in sd.items() if key.startswith(diffusion_model_prefix)}
+            if extracted:
+                return extracted
+        return dict(sd)
+
+    prefixes_to_try: list[str] = []
+
+    def add_prefix(prefix: str | None) -> None:
+        if prefix and prefix not in prefixes_to_try:
+            prefixes_to_try.append(prefix)
+
+    add_prefix(diffusion_model_prefix)
+    add_prefix(comfy.sd.model_detection.unet_prefix_from_state_dict(sd))
+    for prefix in UNET_PREFIX_CANDIDATES:
+        add_prefix(prefix)
+
+    extracted: dict[str, Any] = {}
+    for key, value in sd.items():
+        if key in known_unet_keys:
+            extracted[key] = value
+            continue
+        for prefix in prefixes_to_try:
+            if not key.startswith(prefix):
+                continue
+            stripped = key[len(prefix):]
+            if stripped in known_unet_keys:
+                extracted[stripped] = value
+                break
+    return extracted
+
+
 @contextlib.contextmanager
 def _apply_policy_override(policy_override: str | None):
     if policy_override is None:
@@ -309,6 +558,349 @@ def _apply_policy_override(policy_override: str | None):
         REGISTRY.set_policy(previous_policy)
 
 
+def _bind_model_for_reuse(model, *, source_path: str, note: str, loader_key: str) -> None:
+    REGISTRY.bind_object(model, source_path=source_path, kind=KIND_MODEL, note=note, loader_key=loader_key)
+
+
+def _bind_clip_for_reuse(clip, *, source_path: str, note: str, loader_key: str) -> None:
+    if clip is not None and getattr(clip, "patcher", None) is not None:
+        REGISTRY.bind_object(
+            clip.patcher,
+            source_path=source_path,
+            kind=KIND_CLIP,
+            note=note,
+            loader_key=loader_key,
+            reusable_obj=clip,
+        )
+
+
+def _bind_vae_for_reuse(vae, *, source_path: str, note: str, loader_key: str) -> None:
+    if vae is not None and getattr(vae, "patcher", None) is not None:
+        REGISTRY.bind_object(
+            vae.patcher,
+            source_path=source_path,
+            kind=KIND_VAE,
+            note=note,
+            loader_key=loader_key,
+            reusable_obj=vae,
+        )
+
+
+def _load_resident_diffusion_model(
+    *,
+    loader_name: str,
+    cache_scope: str,
+    source_path: str,
+    note: str,
+    weight_dtype: str,
+    compute_dtype: str,
+    patch_cublaslinear: bool,
+    sage_attention: str,
+    enable_fp16_accumulation: bool,
+    extra_state_dict: str | None = None,
+    policy_override: str | None = None,
+    keep_models: tuple[Any, ...] = (),
+) -> Any:
+    model_options = _build_model_options(weight_dtype)
+    effective_policy = _effective_policy_name(policy_override)
+    loader_key = _make_loader_key(
+        cache_scope,
+        component="model",
+        weight_dtype=weight_dtype,
+        compute_dtype=compute_dtype,
+        patch_cublaslinear=patch_cublaslinear,
+        sage_attention=sage_attention,
+        enable_fp16_accumulation=enable_fp16_accumulation,
+        extra_state_dict=extra_state_dict,
+        policy=effective_policy,
+    )
+    reused_model = REGISTRY.lookup_live_object(kind=KIND_MODEL, source_path=source_path, loader_key=loader_key)
+    if reused_model is not None:
+        _LOG.info("%s: reusing live model for %s", loader_name, source_path)
+        return reused_model
+
+    _warn_pickle_checkpoint_gpu_compatibility(loader_name, source_path)
+    explicit_device = REGISTRY.explicit_load_device(kind=KIND_MODEL, source_path=source_path)
+    _maybe_trim_before_load(
+        loader_name=loader_name,
+        reason=f"model load for {os.path.basename(source_path)}",
+        explicit_device=explicit_device,
+        required_bytes=_estimate_model_load_bytes(
+            source_path,
+            cache_scope=cache_scope,
+            weight_dtype=weight_dtype,
+            extra_state_dict=extra_state_dict,
+        ),
+        keep_models=keep_models,
+    )
+
+    with _temporary_backend_flags(
+        cublas=patch_cublaslinear,
+        fp16_accumulation=enable_fp16_accumulation,
+    ):
+        with REGISTRY.load_context(
+            kind=KIND_MODEL,
+            source_path=source_path,
+            explicit_device=explicit_device,
+            cache_key=loader_key,
+        ):
+            sd, metadata = comfy.utils.load_torch_file(source_path, return_metadata=True)
+            if not _is_safetensors_path(source_path):
+                sd = _extract_unet_state_dict(sd)
+            if extra_state_dict:
+                requested_device = explicit_device if explicit_device is not None else torch.device("cpu")
+                sd.update(
+                    _load_matching_extra_unet_state_dict(
+                        extra_state_dict,
+                        requested_device=requested_device,
+                        known_unet_keys=set(sd),
+                    )
+                )
+
+        model = comfy.sd.load_diffusion_model_state_dict(sd, model_options=model_options, metadata=metadata)
+        _apply_model_postload_options(model, compute_dtype=compute_dtype, sage_attention=sage_attention)
+
+    _bind_model_for_reuse(model, source_path=source_path, note=note, loader_key=loader_key)
+    return model
+
+
+def _checkpoint_model_loader_key(
+    *,
+    weight_dtype: str,
+    compute_dtype: str,
+    patch_cublaslinear: bool,
+    sage_attention: str,
+    enable_fp16_accumulation: bool,
+    policy_override: str | None,
+) -> str:
+    return _make_loader_key(
+        "checkpoint_model",
+        component="model",
+        weight_dtype=weight_dtype,
+        compute_dtype=compute_dtype,
+        patch_cublaslinear=patch_cublaslinear,
+        sage_attention=sage_attention,
+        enable_fp16_accumulation=enable_fp16_accumulation,
+        extra_state_dict=None,
+        policy=_effective_policy_name(policy_override),
+    )
+
+
+def _checkpoint_component_loader_key(component: str, policy_override: str | None) -> str:
+    return _make_loader_key(f"checkpoint_{component}", component=component, policy=_effective_policy_name(policy_override))
+
+
+def _load_checkpoint_clip_only(
+    *,
+    ckpt_path: str,
+    policy_override: str | None,
+    loader_name: str,
+    keep_models: tuple[Any, ...] = (),
+):
+    loader_key = _checkpoint_component_loader_key("clip", policy_override)
+    reused_clip = REGISTRY.lookup_live_object(kind=KIND_CLIP, source_path=ckpt_path, loader_key=loader_key)
+    if reused_clip is not None:
+        _LOG.info("%s: reusing live CLIP for %s", loader_name, ckpt_path)
+        return reused_clip
+
+    _warn_pickle_checkpoint_gpu_compatibility(loader_name, ckpt_path)
+    is_safetensors = _is_safetensors_path(ckpt_path)
+    header_info = checkpoint_component_info_from_header(ckpt_path) if is_safetensors else None
+    model_config = None if header_info is None else header_info.get("model_config")
+    explicit_device = REGISTRY.explicit_load_device(kind=KIND_CLIP, source_path=ckpt_path)
+    _maybe_trim_before_load(
+        loader_name=loader_name,
+        reason=f"CLIP load for {os.path.basename(ckpt_path)}",
+        explicit_device=explicit_device,
+        required_bytes=_estimate_checkpoint_aux_component_bytes(ckpt_path, kind=KIND_CLIP),
+        keep_models=keep_models,
+    )
+    with REGISTRY.load_context(kind=KIND_CLIP, source_path=ckpt_path, explicit_device=explicit_device, cache_key=loader_key):
+        if is_safetensors and model_config is None:
+            requested_device = explicit_device if explicit_device is not None else torch.device("cpu")
+            sd, metadata, _, _ = load_safetensors_state_dict(
+                ckpt_path,
+                requested_device,
+                return_metadata=True,
+                selected_keys=None,
+            )
+        else:
+            sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
+    if model_config is None:
+        _, clip, _, _ = comfy.sd.load_state_dict_guess_config(
+            sd,
+            output_vae=False,
+            output_clip=True,
+            output_model=False,
+            embedding_directory=folder_paths.get_folder_paths("embeddings"),
+            metadata=metadata,
+        )
+    else:
+        scaled_fp8_list = []
+        for key in list(sd.keys()):
+            if key.endswith(".scaled_fp8"):
+                scaled_fp8_list.append(key[:-len(".scaled_fp8")])
+
+        if scaled_fp8_list:
+            clip_source_sd: dict[str, Any] = {}
+            for key, value in sd.items():
+                if any(key.startswith(prefix) for prefix in scaled_fp8_list):
+                    continue
+                clip_source_sd[key] = value
+            for prefix in scaled_fp8_list:
+                quant_sd, _ = comfy.utils.convert_old_quants(sd, prefix, metadata=metadata or {})
+                clip_source_sd.update(quant_sd)
+        else:
+            clip_source_sd = sd
+
+        clip_target = model_config.clip_target(state_dict=clip_source_sd)
+        if clip_target is None:
+            clip = None
+        else:
+            clip_sd = model_config.process_clip_state_dict(clip_source_sd)
+            if len(clip_sd) == 0:
+                _LOG.warning("%s: no CLIP/text encoder weights found in %s after selective checkpoint load", loader_name, ckpt_path)
+                clip = None
+            else:
+                parameters = comfy.utils.calculate_parameters(clip_sd)
+                clip = comfy.sd.CLIP(
+                    clip_target,
+                    embedding_directory=folder_paths.get_folder_paths("embeddings"),
+                    tokenizer_data=clip_sd,
+                    parameters=parameters,
+                    state_dict=clip_sd,
+                    model_options={},
+                    disable_dynamic=False,
+                )
+    _bind_clip_for_reuse(clip, source_path=ckpt_path, note="checkpoint clip", loader_key=loader_key)
+    return clip
+
+
+def _load_checkpoint_vae_only(
+    *,
+    ckpt_path: str,
+    policy_override: str | None,
+    loader_name: str,
+    keep_models: tuple[Any, ...] = (),
+):
+    loader_key = _checkpoint_component_loader_key("vae", policy_override)
+    reused_vae = REGISTRY.lookup_live_object(kind=KIND_VAE, source_path=ckpt_path, loader_key=loader_key)
+    if reused_vae is not None:
+        _LOG.info("%s: reusing live VAE for %s", loader_name, ckpt_path)
+        return reused_vae
+
+    _warn_pickle_checkpoint_gpu_compatibility(loader_name, ckpt_path)
+    is_safetensors = _is_safetensors_path(ckpt_path)
+    header_info = checkpoint_component_info_from_header(ckpt_path) if is_safetensors else None
+    model_config = None if header_info is None else header_info.get("model_config")
+    explicit_device = REGISTRY.explicit_load_device(kind=KIND_VAE, source_path=ckpt_path)
+    _maybe_trim_before_load(
+        loader_name=loader_name,
+        reason=f"VAE load for {os.path.basename(ckpt_path)}",
+        explicit_device=explicit_device,
+        required_bytes=_estimate_checkpoint_aux_component_bytes(ckpt_path, kind=KIND_VAE),
+        keep_models=keep_models,
+    )
+    with REGISTRY.load_context(kind=KIND_VAE, source_path=ckpt_path, explicit_device=explicit_device, cache_key=loader_key):
+        if is_safetensors and model_config is None:
+            requested_device = explicit_device if explicit_device is not None else torch.device("cpu")
+            sd, metadata, _, _ = load_safetensors_state_dict(
+                ckpt_path,
+                requested_device,
+                return_metadata=True,
+                selected_keys=None,
+            )
+        else:
+            sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
+    if model_config is None:
+        _, _, vae, _ = comfy.sd.load_state_dict_guess_config(
+            sd,
+            output_vae=True,
+            output_clip=False,
+            output_model=False,
+            embedding_directory=folder_paths.get_folder_paths("embeddings"),
+            metadata=metadata,
+        )
+    else:
+        vae_sd = comfy.utils.state_dict_prefix_replace(
+            sd,
+            {prefix: "" for prefix in model_config.vae_key_prefix},
+            filter_keys=True,
+        )
+        vae_sd = model_config.process_vae_state_dict(vae_sd)
+        if len(vae_sd) == 0:
+            _LOG.warning("%s: no VAE weights found in %s after selective checkpoint load", loader_name, ckpt_path)
+            vae = None
+        else:
+            vae = comfy.sd.VAE(sd=vae_sd, metadata=metadata)
+    _bind_vae_for_reuse(vae, source_path=ckpt_path, note="checkpoint vae", loader_key=loader_key)
+    return vae
+
+
+def _load_full_checkpoint(
+    *,
+    ckpt_path: str,
+    weight_dtype: str,
+    compute_dtype: str,
+    patch_cublaslinear: bool,
+    sage_attention: str,
+    enable_fp16_accumulation: bool,
+    policy_override: str | None,
+    loader_name: str,
+):
+    model_key = _checkpoint_model_loader_key(
+        weight_dtype=weight_dtype,
+        compute_dtype=compute_dtype,
+        patch_cublaslinear=patch_cublaslinear,
+        sage_attention=sage_attention,
+        enable_fp16_accumulation=enable_fp16_accumulation,
+        policy_override=policy_override,
+    )
+    clip_key = _checkpoint_component_loader_key("clip", policy_override)
+    vae_key = _checkpoint_component_loader_key("vae", policy_override)
+
+    model = REGISTRY.lookup_live_object(kind=KIND_MODEL, source_path=ckpt_path, loader_key=model_key)
+    clip = REGISTRY.lookup_live_object(kind=KIND_CLIP, source_path=ckpt_path, loader_key=clip_key)
+    vae = REGISTRY.lookup_live_object(kind=KIND_VAE, source_path=ckpt_path, loader_key=vae_key)
+    missing = [name for name, value in (("model", model), ("clip", clip), ("vae", vae)) if value is None]
+    if not missing:
+        _LOG.info("%s: reusing live checkpoint outputs for %s", loader_name, ckpt_path)
+        return model, clip, vae
+
+    keep_models = _normalize_keep_models(model, clip, vae)
+    if model is None:
+        model = _load_resident_diffusion_model(
+            loader_name=loader_name,
+            cache_scope="checkpoint_model",
+            source_path=ckpt_path,
+            note="checkpoint model",
+            weight_dtype=weight_dtype,
+            compute_dtype=compute_dtype,
+            patch_cublaslinear=patch_cublaslinear,
+            sage_attention=sage_attention,
+            enable_fp16_accumulation=enable_fp16_accumulation,
+            policy_override=policy_override,
+            keep_models=keep_models,
+        )
+        keep_models = keep_models + _normalize_keep_models(model)
+    if clip is None:
+        clip = _load_checkpoint_clip_only(
+            ckpt_path=ckpt_path,
+            policy_override=policy_override,
+            loader_name=loader_name,
+            keep_models=keep_models,
+        )
+        keep_models = keep_models + _normalize_keep_models(clip)
+    if vae is None:
+        vae = _load_checkpoint_vae_only(
+            ckpt_path=ckpt_path,
+            policy_override=policy_override,
+            loader_name=loader_name,
+            keep_models=keep_models,
+        )
+    return model, clip, vae
+
+
 class DiffusionModelSelectorResident:
     @classmethod
     def INPUT_TYPES(cls):
@@ -409,26 +1001,20 @@ class DiffusionModelLoaderResident:
         )
 
         with _apply_policy_override(policy_override):
-            with _temporary_backend_flags(
-                cublas=patch_cublaslinear,
-                fp16_accumulation=enable_fp16_accumulation,
-            ):
-                model_options = _build_model_options(weight_dtype)
-                unet_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
-                explicit_device = REGISTRY.explicit_load_device(kind=KIND_MODEL, source_path=unet_path)
-
-                with REGISTRY.load_context(kind=KIND_MODEL, source_path=unet_path, explicit_device=explicit_device):
-                    sd, metadata = comfy.utils.load_torch_file(unet_path, return_metadata=True)
-                    if extra_state_dict:
-                        extra_sd = comfy.utils.load_torch_file(extra_state_dict)
-                        sd.update(extra_sd)
-                        del extra_sd
-
-                diffusion_model_prefix = comfy.sd.model_detection.unet_prefix_from_state_dict(sd)
-                sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=False)
-                model = comfy.sd.load_diffusion_model_state_dict(sd, model_options=model_options, metadata=metadata)
-                _apply_model_postload_options(model, compute_dtype=compute_dtype, sage_attention=sage_attention)
-            REGISTRY.bind_object(model, source_path=unet_path, kind=KIND_MODEL)
+            unet_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
+            model = _load_resident_diffusion_model(
+                loader_name="Diffusion Model Loader Resident",
+                cache_scope="diffusion_model",
+                source_path=unet_path,
+                note="diffusion model",
+                weight_dtype=weight_dtype,
+                compute_dtype=compute_dtype,
+                patch_cublaslinear=patch_cublaslinear,
+                sage_attention=sage_attention,
+                enable_fp16_accumulation=enable_fp16_accumulation,
+                extra_state_dict=extra_state_dict,
+                policy_override=policy_override,
+            )
             return (model,)
 
 
@@ -475,7 +1061,7 @@ class CheckpointLoaderResident:
     CATEGORY = "GPU Resident Loader/loaders"
     DESCRIPTION = (
         "Checkpoint loader with the KJ DiffusionModelLoader-style tuning knobs plus GPU-resident ingest. "
-        "It loads the whole checkpoint, then binds the model, CLIP, and VAE into the residency registry."
+        "It reuses live checkpoint components when possible and composes missing outputs from the model, CLIP, and VAE loaders instead of materializing a broad checkpoint state dict."
     )
 
     def load(
@@ -495,29 +1081,130 @@ class CheckpointLoaderResident:
         )
 
         with _apply_policy_override(policy_override):
-            with _temporary_backend_flags(
-                cublas=patch_cublaslinear,
-                fp16_accumulation=enable_fp16_accumulation,
-            ):
-                model_options = _build_model_options(weight_dtype)
-                ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
-                explicit_device = REGISTRY.explicit_load_device(kind=KIND_CHECKPOINT, source_path=ckpt_path)
-
-                with REGISTRY.load_context(kind=KIND_CHECKPOINT, source_path=ckpt_path, explicit_device=explicit_device):
-                    sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
-
-                model, clip, vae, _ = comfy.sd.load_state_dict_guess_config(
-                    sd,
-                    output_vae=True,
-                    output_clip=True,
-                    embedding_directory=folder_paths.get_folder_paths("embeddings"),
-                    metadata=metadata,
-                    model_options=model_options,
-                )
-                _apply_model_postload_options(model, compute_dtype=compute_dtype, sage_attention=sage_attention)
-            REGISTRY.bind_object(model, source_path=ckpt_path, kind=KIND_MODEL, note="checkpoint model")
-            if clip is not None and getattr(clip, "patcher", None) is not None:
-                REGISTRY.bind_object(clip.patcher, source_path=ckpt_path, kind=KIND_CLIP, note="checkpoint clip")
-            if vae is not None and getattr(vae, "patcher", None) is not None:
-                REGISTRY.bind_object(vae.patcher, source_path=ckpt_path, kind=KIND_VAE, note="checkpoint vae")
+            ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
+            model, clip, vae = _load_full_checkpoint(
+                ckpt_path=ckpt_path,
+                weight_dtype=weight_dtype,
+                compute_dtype=compute_dtype,
+                patch_cublaslinear=patch_cublaslinear,
+                sage_attention=sage_attention,
+                enable_fp16_accumulation=enable_fp16_accumulation,
+                policy_override=policy_override,
+                loader_name="Checkpoint Loader Resident",
+            )
             return model, clip, vae
+
+
+class CheckpointModelLoaderResident:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return CheckpointLoaderResident.INPUT_TYPES()
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load"
+    CATEGORY = "GPU Resident Loader/loaders"
+    DESCRIPTION = (
+        "Checkpoint model-only loader. Safetensors checkpoints take the selective UNet fast path, and live equivalent models are reused when available."
+    )
+
+    def load(
+        self,
+        ckpt_name: str,
+        weight_dtype: str,
+        compute_dtype: str,
+        patch_cublaslinear: bool,
+        sage_attention: str,
+        enable_fp16_accumulation: bool,
+        policy_override: str | None = None,
+    ):
+        policy_override, _ = _resolve_loader_policy_and_extra_state_dict(
+            loader_name="Checkpoint Model Loader Resident",
+            policy_override=policy_override,
+            extra_state_dict=None,
+        )
+
+        with _apply_policy_override(policy_override):
+            ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
+            model = _load_resident_diffusion_model(
+                loader_name="Checkpoint Model Loader Resident",
+                cache_scope="checkpoint_model",
+                source_path=ckpt_path,
+                note="checkpoint model",
+                weight_dtype=weight_dtype,
+                compute_dtype=compute_dtype,
+                patch_cublaslinear=patch_cublaslinear,
+                sage_attention=sage_attention,
+                enable_fp16_accumulation=enable_fp16_accumulation,
+                policy_override=policy_override,
+            )
+            return (model,)
+
+
+class CheckpointClipLoaderResident:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "ckpt_name": (
+                    folder_paths.get_filename_list("checkpoints"),
+                    {"tooltip": "Checkpoint file to load."},
+                ),
+            },
+            "optional": {
+                "policy_override": (
+                    "STRING",
+                    {
+                        "forceInput": True,
+                        "tooltip": "Optional residency policy override. Connect Set Global Residency Policy here.",
+                    },
+                )
+            },
+        }
+
+    RETURN_TYPES = ("CLIP",)
+    FUNCTION = "load"
+    CATEGORY = "GPU Resident Loader/loaders"
+    DESCRIPTION = "Checkpoint CLIP-only loader with live-object reuse. It avoids rebuilding the text encoder when an equivalent CLIP object is already alive."
+
+    def load(self, ckpt_name: str, policy_override: str | None = None):
+        policy_override, _ = _resolve_loader_policy_and_extra_state_dict(
+            loader_name="Checkpoint Clip Loader Resident",
+            policy_override=policy_override,
+            extra_state_dict=None,
+        )
+
+        with _apply_policy_override(policy_override):
+            ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
+            clip = _load_checkpoint_clip_only(
+                ckpt_path=ckpt_path,
+                policy_override=policy_override,
+                loader_name="Checkpoint Clip Loader Resident",
+            )
+            return (clip,)
+
+
+class CheckpointVAELoaderResident:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return CheckpointClipLoaderResident.INPUT_TYPES()
+
+    RETURN_TYPES = ("VAE",)
+    FUNCTION = "load"
+    CATEGORY = "GPU Resident Loader/loaders"
+    DESCRIPTION = "Checkpoint VAE-only loader with live-object reuse. It avoids rebuilding the VAE when an equivalent object is still alive."
+
+    def load(self, ckpt_name: str, policy_override: str | None = None):
+        policy_override, _ = _resolve_loader_policy_and_extra_state_dict(
+            loader_name="Checkpoint VAE Loader Resident",
+            policy_override=policy_override,
+            extra_state_dict=None,
+        )
+
+        with _apply_policy_override(policy_override):
+            ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
+            vae = _load_checkpoint_vae_only(
+                ckpt_path=ckpt_path,
+                policy_override=policy_override,
+                loader_name="Checkpoint VAE Loader Resident",
+            )
+            return (vae,)

nodes.py

@@ -1,13 +1,24 @@
 from __future__ import annotations
 
 import json
+import logging
 from typing import Any
 
 import comfy.model_management as model_management
 
-from .kj_loader import CheckpointLoaderResident, DiffusionModelLoaderResident, DiffusionModelSelectorResident
+from .cleanup import unload_loaded_model
+from .kj_loader import (
+    CheckpointClipLoaderResident,
+    CheckpointLoaderResident,
+    CheckpointModelLoaderResident,
+    CheckpointVAELoaderResident,
+    DiffusionModelLoaderResident,
+    DiffusionModelSelectorResident,
+)
 from .residency import REGISTRY
 
+_LOG = logging.getLogger(__name__)
+
 
 def _entry_report_json(obj: Any) -> str:
     entry = REGISTRY.entry_for_object(obj)
@@ -52,9 +63,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()
@@ -328,6 +358,9 @@ NODE_CLASS_MAPPINGS = {
     "DiffusionModelSelectorResident": DiffusionModelSelectorResident,
     "DiffusionModelLoaderResident": DiffusionModelLoaderResident,
     "CheckpointLoaderResident": CheckpointLoaderResident,
+    "CheckpointModelLoaderResident": CheckpointModelLoaderResident,
+    "CheckpointClipLoaderResident": CheckpointClipLoaderResident,
+    "CheckpointVAELoaderResident": CheckpointVAELoaderResident,
     "SetGlobalResidencyPolicy": SetGlobalResidencyPolicy,
     "RegistrySnapshot": RegistrySnapshot,
     "PinModelResidency": PinModelResidency,
@@ -348,6 +381,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DiffusionModelSelectorResident": "Diffusion Model Selector Resident",
     "DiffusionModelLoaderResident": "Diffusion Model Loader Resident",
     "CheckpointLoaderResident": "Checkpoint Loader Resident",
+    "CheckpointModelLoaderResident": "Checkpoint Model Loader Resident",
+    "CheckpointClipLoaderResident": "Checkpoint Clip Loader Resident",
+    "CheckpointVAELoaderResident": "Checkpoint VAE Loader Resident",
     "SetGlobalResidencyPolicy": "Set Global Residency Policy",
     "RegistrySnapshot": "Registry Snapshot",
     "PinModelResidency": "Pin Model Residency",

patches.py

@@ -1,13 +1,23 @@
 from __future__ import annotations
 
+import contextlib
 import functools
+import json
 import logging
 import os
+import struct
 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,
@@ -23,6 +33,183 @@ _LOG = logging.getLogger(__name__)
 _PATCHED = False
 _ORIGINALS: dict[str, Callable[..., Any]] = {}
 _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,
+    "I8": torch.int8,
+    "I16": torch.int16,
+    "U16": getattr(torch, "uint16", torch.int32),
+    "I32": torch.int32,
+    "U32": getattr(torch, "uint32", torch.int64),
+    "I64": torch.int64,
+    "U64": getattr(torch, "uint64", torch.int64),
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "F32": torch.float32,
+    "F64": torch.float64,
+    "F8_E4M3FN": getattr(torch, "float8_e4m3fn", torch.float16),
+    "F8_E5M2": getattr(torch, "float8_e5m2", torch.float16),
+}
+
+
+def _safetensors_header_cache_key(path: str) -> tuple[str, int, int]:
+    stat = os.stat(path)
+    return (os.path.abspath(path), stat.st_mtime_ns, stat.st_size)
+
+
+def _read_safetensors_header(path: str) -> tuple[dict[str, dict[str, Any]], dict[str, str] | None]:
+    with open(path, "rb") as handle:
+        header_size = struct.unpack("<Q", handle.read(8))[0]
+        header = json.loads(handle.read(header_size))
+
+    metadata = header.get("__metadata__")
+    tensor_headers = {
+        key: value
+        for key, value in header.items()
+        if key != "__metadata__" and isinstance(value, dict)
+    }
+    return tensor_headers, metadata if isinstance(metadata, dict) else None
+
+
+def _torch_dtype_from_safetensors_code(code: str | None) -> torch.dtype:
+    if code is None:
+        return torch.float32
+    return _SAFETENSORS_DTYPE_MAP.get(code, torch.float32)
+
+
+def _tensor_nbytes_from_header(
+    tensor_info: dict[str, Any],
+    *,
+    dtype_override: torch.dtype | None = None,
+) -> int:
+    dtype = dtype_override if dtype_override is not None else _torch_dtype_from_safetensors_code(tensor_info.get("dtype"))
+    numel = 1
+    for dim in tensor_info.get("shape", ()):
+        numel *= int(dim)
+    return int(numel) * int(torch.empty((), dtype=dtype).element_size())
+
+
+def _build_meta_state_dict_from_header(tensor_headers: dict[str, dict[str, Any]]) -> dict[str, torch.Tensor]:
+    meta_state_dict: dict[str, torch.Tensor] = {}
+    for key, tensor_info in tensor_headers.items():
+        shape = tuple(int(dim) for dim in tensor_info.get("shape", ()))
+        meta_state_dict[key] = torch.empty(
+            shape,
+            dtype=_torch_dtype_from_safetensors_code(tensor_info.get("dtype")),
+            device="meta",
+        )
+    return meta_state_dict
+
+
+@functools.lru_cache(maxsize=_SAFE_TENSORS_COMPONENT_CACHE_MAX)
+def _cached_component_key_maps(cache_key: tuple[str, int, int]) -> dict[str, Any]:
+    import comfy.sd as comfy_sd
+
+    path = cache_key[0]
+    tensor_headers, metadata = _read_safetensors_header(path)
+    all_keys = tuple(tensor_headers)
+    unet_prefix = infer_unet_prefix_from_keys(all_keys)
+    meta_state_dict = _build_meta_state_dict_from_header(tensor_headers)
+    model_config = comfy_sd.model_detection.model_config_from_unet(meta_state_dict, unet_prefix, metadata=metadata)
+
+    def select_prefixed(prefixes: tuple[str, ...] | list[str] | None) -> tuple[tuple[str, str], ...]:
+        if not prefixes:
+            return ()
+        return tuple(
+            (key, key)
+            for key in all_keys
+            if any(key.startswith(prefix) for prefix in prefixes)
+        )
+
+    return {
+        "metadata": metadata,
+        "unet_prefix": unet_prefix,
+        "model_config": model_config,
+        "model": tuple((key, key[len(unet_prefix):]) for key in all_keys if key.startswith(unet_prefix)),
+        "clip": select_prefixed(getattr(model_config, "text_encoder_key_prefix", None) or ()),
+        "vae": select_prefixed(getattr(model_config, "vae_key_prefix", None) or ()),
+    }
+
+
+def checkpoint_component_info_from_header(path: str) -> dict[str, Any] | None:
+    try:
+        return _cached_component_key_maps(_safetensors_header_cache_key(path))
+    except Exception as exc:
+        _LOG.warning("GPU Resident Loader: failed to build selective safetensors header map for %s: %s", path, exc)
+        return None
+
+
+def estimate_safetensors_tensor_bytes(
+    path: str,
+    *,
+    selected_keys: list[str] | tuple[str, ...] | set[str] | None = None,
+    dtype_override: torch.dtype | None = None,
+) -> int | None:
+    try:
+        tensor_headers, _ = _read_safetensors_header(path)
+    except Exception as exc:
+        _LOG.warning("GPU Resident Loader: failed to estimate safetensors tensor bytes for %s: %s", path, exc)
+        return None
+
+    selected = None if selected_keys is None else set(selected_keys)
+    total = 0
+    matched = 0
+    for key, tensor_info in tensor_headers.items():
+        if selected is not None and key not in selected:
+            continue
+        total += _tensor_nbytes_from_header(tensor_info, dtype_override=dtype_override)
+        matched += 1
+    if selected is not None and selected and matched == 0:
+        _LOG.warning(
+            "GPU Resident Loader: selected tensor keys were provided but none matched %s; "
+            "treating size as unknown for fallback",
+            path,
+        )
+        return None
+    return int(total)
+
+
+def estimate_checkpoint_component_bytes(
+    path: str,
+    kind: str,
+    *,
+    dtype_override: torch.dtype | None = None,
+) -> int | None:
+    component_maps = checkpoint_component_info_from_header(path)
+    if component_maps is None:
+        return None
+    pairs = component_maps.get(kind, ())
+    if not pairs:
+        return None
+    return estimate_safetensors_tensor_bytes(
+        path,
+        selected_keys=[source_key for source_key, _ in pairs],
+        dtype_override=dtype_override,
+    )
+
+
+def _selected_component_keys_from_header(path: str, kind: str) -> dict[str, str] | None:
+    if kind not in {KIND_MODEL, KIND_CLIP, KIND_VAE}:
+        return None
+    component_maps = checkpoint_component_info_from_header(path)
+    if component_maps is None:
+        return None
+
+    pairs = component_maps.get(kind, ())
+    return dict(pairs) if pairs else None
+
+
+def _selected_component_suffix(kind: str | None) -> str | None:
+    return {
+        KIND_MODEL: "model_only",
+        KIND_CLIP: "clip_only",
+        KIND_VAE: "vae_only",
+    }.get(kind)
 
 
 def _normalize_device(device: Any | None) -> torch.device | None:
@@ -42,6 +229,40 @@ def _device_string(device: torch.device | None) -> str:
     return str(device)
 
 
+def _devices_match(device_a: Any | None, device_b: Any | None) -> bool:
+    return _shared_device_matches(device_a, device_b)
+
+
+def _cpu_offload_required(model: Any, loaded_device: Any | None) -> bool:
+    current_device = None
+    if hasattr(model, "current_loaded_device"):
+        try:
+            current_device = _normalize_device(model.current_loaded_device())
+        except Exception:
+            current_device = None
+    if current_device is None:
+        current_device = _normalize_device(loaded_device)
+    return _should_force_cpu_offload(model, active_device=current_device)
+
+
+@contextlib.contextmanager
+def _temporary_offload_device(model: Any, target_device: torch.device | None):
+    if model is None or target_device is None or not hasattr(model, "offload_device"):
+        yield False
+        return
+
+    original_device = getattr(model, "offload_device", None)
+    if _devices_match(original_device, target_device):
+        yield False
+        return
+
+    setattr(model, "offload_device", target_device)
+    try:
+        yield True
+    finally:
+        setattr(model, "offload_device", original_device)
+
+
 def _safe_open_device_arg(device: torch.device) -> Any:
     if device.type == "cuda":
         return device.index if device.index is not None else torch.cuda.current_device()
@@ -97,6 +318,125 @@ def _state_dict_device_summary(sd: dict[str, Any], requested_device: torch.devic
     return ", ".join(sorted(devices))
 
 
+def _device_summary_from_observed(observed_devices: set[str], requested_device: torch.device) -> str:
+    if not observed_devices:
+        return str(requested_device)
+    if len(observed_devices) == 1:
+        return next(iter(observed_devices))
+    return ", ".join(sorted(observed_devices))
+
+
+def infer_unet_prefix_from_keys(keys: list[str] | tuple[str, ...]) -> str:
+    counts = {candidate: 0 for candidate in _UNET_PREFIX_CANDIDATES}
+    for key in keys:
+        for candidate in _UNET_PREFIX_CANDIDATES:
+            if key.startswith(candidate):
+                counts[candidate] += 1
+                break
+    top = max(counts, key=counts.get)
+    return top if counts[top] > 5 else "model."
+
+
+def load_safetensors_state_dict(
+    ckpt: str,
+    requested_device: torch.device,
+    *,
+    return_metadata: bool = False,
+    selected_keys: dict[str, str] | None = None,
+) -> tuple[dict[str, Any], Any, str, str]:
+    import comfy.memory_management
+    import comfy.utils as comfy_utils
+
+    metadata = None
+    if comfy.memory_management.aimdo_enabled and requested_device.type == "cpu" and selected_keys is None:
+        sd, metadata = comfy_utils.load_safetensors(ckpt)
+        if not return_metadata:
+            metadata = None
+        return sd, metadata, "cpu", "aimdo_cpu"
+
+    disable_mmap = getattr(comfy_utils, "DISABLE_MMAP", False)
+
+    def read_handle(device_arg: Any, *, move_to_requested_device: bool) -> tuple[dict[str, Any], Any, str]:
+        observed_devices: set[str] = set()
+        with safe_open(ckpt, framework="pt", device=device_arg) as handle:
+            key_map = selected_keys if selected_keys is not None else {key: key for key in handle.keys()}
+            sd: dict[str, Any] = {}
+            for source_key, target_key in key_map.items():
+                tensor = handle.get_tensor(source_key)
+                loaded = _prepare_loaded_tensor(
+                    source_key,
+                    tensor,
+                    requested_device,
+                    disable_mmap=disable_mmap,
+                    move_to_requested_device=move_to_requested_device,
+                )
+                sd[target_key] = loaded
+                if torch.is_tensor(loaded):
+                    observed_devices.add(str(loaded.device))
+            handle_metadata = handle.metadata() if return_metadata else None
+        return sd, handle_metadata, _device_summary_from_observed(observed_devices, requested_device)
+
+    try:
+        safe_device = _safe_open_device_arg(requested_device)
+        sd, metadata, actual_device = read_handle(safe_device, move_to_requested_device=False)
+        return sd, metadata, actual_device, "direct"
+    except Exception as exc:
+        if requested_device.type == "cuda":
+            _LOG.warning(
+                "GPU Resident Loader: direct GPU safetensors load failed for %s; falling back to CPU path: %s",
+                ckpt,
+                exc,
+            )
+            try:
+                sd, metadata, actual_device = read_handle("cpu", move_to_requested_device=True)
+                return sd, metadata, actual_device, "cpu_then_copy"
+            except Exception as fallback_exc:
+                raise fallback_exc from exc
+        raise
+
+
+def _warn_pickle_gpu_compatibility(path: str, requested_device: torch.device) -> None:
+    if requested_device.type != "cuda" or path in _WARNED_PICKLE_GPU_PATHS:
+        return
+    _WARNED_PICKLE_GPU_PATHS.add(path)
+    _LOG.warning(
+        "GPU Resident Loader: %s is not a safetensors file, so GPU-resident loading still goes through CPU-first torch.load(). "
+        "Convert hot models with scripts/convert_checkpoint_to_safetensors.py for the narrow fast path.",
+        path,
+    )
+
+
+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:
@@ -142,36 +482,28 @@ def _patched_load_torch_file(ckpt, safe_load=False, device=None, return_metadata
 
     if lowered.endswith((".safetensors", ".sft")):
         try:
-            if comfy.memory_management.aimdo_enabled and requested_device.type == "cpu":
-                sd, metadata = comfy_utils.load_safetensors(ckpt)
+            selected_keys = None
+            selected_suffix = None
+            if ctx is not None and ctx.kind in {KIND_MODEL, KIND_CLIP, KIND_VAE}:
+                selected_keys = _selected_component_keys_from_header(ckpt, ctx.kind)
+                selected_suffix = _selected_component_suffix(ctx.kind)
+
+            sd, metadata, actual_device, load_mode = load_safetensors_state_dict(
+                ckpt,
+                requested_device,
+                return_metadata=return_metadata,
+                selected_keys=selected_keys,
+            )
+            if load_mode == "aimdo_cpu":
                 method = "safetensors_aimdo_cpu"
-                if not return_metadata:
-                    metadata = None
-                _record_generic_load(
-                    path=ckpt,
-                    method=method,
-                    requested_device=requested_device,
-                    actual_device="cpu",
+            elif selected_keys is not None and selected_suffix is not None:
+                method = f"safetensors_cpu_then_copy_to_cuda_{selected_suffix}" if load_mode == "cpu_then_copy" else (
+                    f"safetensors_gpu_direct_{selected_suffix}" if requested_device.type == "cuda" else f"safetensors_cpu_{selected_suffix}"
+                )
+            else:
+                method = "safetensors_cpu_then_copy_to_cuda" if load_mode == "cpu_then_copy" else (
+                    "safetensors_gpu_direct" if requested_device.type == "cuda" else "safetensors_cpu"
                 )
-                return (sd, metadata) if return_metadata else sd
-
-            safe_device = _safe_open_device_arg(requested_device)
-            with safe_open(ckpt, framework="pt", device=safe_device) as handle:
-                sd = {}
-                disable_mmap = getattr(comfy_utils, "DISABLE_MMAP", False)
-                for key in handle.keys():
-                    tensor = handle.get_tensor(key)
-                    sd[key] = _prepare_loaded_tensor(
-                        key,
-                        tensor,
-                        requested_device,
-                        disable_mmap=disable_mmap,
-                    )
-                if return_metadata:
-                    metadata = handle.metadata()
-
-            actual_device = _state_dict_device_summary(sd, requested_device)
-            method = "safetensors_gpu_direct" if requested_device.type == "cuda" else "safetensors_cpu"
             _record_generic_load(
                 path=ckpt,
                 method=method,
@@ -180,43 +512,6 @@ def _patched_load_torch_file(ckpt, safe_load=False, device=None, return_metadata
             )
             return (sd, metadata) if return_metadata else sd
         except Exception as exc:
-            if requested_device.type == "cuda":
-                _LOG.warning(
-                    "GPU Resident Loader: direct GPU safetensors load failed for %s; falling back to CPU path: %s",
-                    ckpt,
-                    exc,
-                )
-                try:
-                    with safe_open(ckpt, framework="pt", device="cpu") as handle:
-                        sd = {}
-                        for key in handle.keys():
-                            sd[key] = _prepare_loaded_tensor(
-                                key,
-                                handle.get_tensor(key),
-                                requested_device,
-                                disable_mmap=False,
-                                move_to_requested_device=True,
-                            )
-                        if return_metadata:
-                            metadata = handle.metadata()
-                    _record_generic_load(
-                        path=ckpt,
-                        method="safetensors_cpu_then_copy_to_cuda",
-                        requested_device=requested_device,
-                        actual_device=_state_dict_device_summary(sd, requested_device),
-                        error=str(exc),
-                    )
-                    return (sd, metadata) if return_metadata else sd
-                except Exception as fallback_exc:
-                    _record_generic_load(
-                        path=ckpt,
-                        method="safetensors_cpu_fallback_failed",
-                        requested_device=requested_device,
-                        actual_device="error",
-                        error=str(fallback_exc),
-                    )
-                    raise fallback_exc from exc
-
             if len(getattr(exc, "args", ())) > 0:
                 message = exc.args[0]
                 if isinstance(message, str):
@@ -245,6 +540,7 @@ def _patched_load_torch_file(ckpt, safe_load=False, device=None, return_metadata
     if getattr(comfy_utils, "MMAP_TORCH_FILES", False):
         torch_args["mmap"] = True
 
+    _warn_pickle_gpu_compatibility(ckpt, requested_device)
     torch_load_device = torch.device("cpu") if requested_device.type == "cuda" else requested_device
     pl_sd = torch.load(ckpt, map_location=torch_load_device, weights_only=True, **torch_args)
     if "state_dict" in pl_sd:
@@ -345,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)
@@ -357,28 +868,97 @@ def _wrap_load_models_gpu(func: Callable[..., Any]) -> Callable[..., Any]:
     return wrapper
 
 
+def _wrap_loaded_model_unload(func: Callable[..., Any]) -> Callable[..., Any]:
+    @functools.wraps(func)
+    def wrapper(self, memory_to_free=None, unpatch_weights=True):
+        model = getattr(self, "model", None)
+        loaded_device = getattr(self, "device", None)
+        if not _cpu_offload_required(model, loaded_device):
+            return func(self, memory_to_free=memory_to_free, unpatch_weights=unpatch_weights)
+
+        with _temporary_offload_device(model, torch.device("cpu")) as redirected:
+            if redirected:
+                _LOG.debug(
+                    "GPU Resident Loader: redirecting unload of %s from %s to CPU to reclaim VRAM",
+                    type(getattr(model, "model", model)).__name__,
+                    _device_string(_normalize_device(loaded_device)),
+                )
+            return func(self, memory_to_free=memory_to_free, unpatch_weights=unpatch_weights)
+
+    return wrapper
+
+
+def _wrap_model_patcher_detach(func: Callable[..., Any]) -> Callable[..., Any]:
+    @functools.wraps(func)
+    def wrapper(self, unpatch_all=True):
+        if not _cpu_offload_required(self, getattr(self.model, "device", None)):
+            return func(self, unpatch_all=unpatch_all)
+
+        with _temporary_offload_device(self, torch.device("cpu")) as redirected:
+            if redirected:
+                _LOG.debug(
+                    "GPU Resident Loader: redirecting detach of %s from %s to CPU to reclaim VRAM",
+                    type(getattr(self, "model", self)).__name__,
+                    _device_string(_normalize_device(getattr(self.model, "device", None))),
+                )
+            return func(self, unpatch_all=unpatch_all)
+
+    return wrapper
+
+
 def _wrap_free_memory(func: Callable[..., Any]) -> Callable[..., Any]:
     @functools.wraps(func)
     def wrapper(memory_required, device, keep_loaded=None, *args, **kwargs):
         import comfy.model_management as model_management
 
         keep_loaded = list(keep_loaded or [])
-        sticky_wrappers = []
-        if REGISTRY.get_policy() == "sticky_gpu":
+        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]
-
-        unloaded = func(memory_required, device, keep_loaded + sticky_wrappers, *args, **kwargs)
-
         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:
-                free_after = model_management.get_free_memory(device)
-            except Exception:
-                free_after = None
-            if free_after is not None and free_after < memory_required:
-                _LOG.warning(
-                    "GPU Resident Loader: sticky set exceeded VRAM budget; allowing fallback eviction to satisfy request"
+                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,
                 )
-                unloaded = func(memory_required, device, keep_loaded, *args, **kwargs)
+            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:
+                free_now = None
+            if free_now is None:
+                protected_wrappers = sticky_wrappers
+            else:
+                unloadable_wrappers = []
+                for loaded in list(model_management.current_loaded_models):
+                    if loaded.device == device and loaded not in keep_loaded and not loaded.is_dead():
+                        unloadable_wrappers.append(loaded)
+                available_for_protection = max(
+                    0,
+                    free_now + sum(max(0, loaded.model_loaded_memory()) for loaded in unloadable_wrappers) - protection_target,
+                )
+                protected_memory = 0
+                for loaded in sticky_wrappers:
+                    estimated_memory = max(0, loaded.model_loaded_memory())
+                    if protected_memory + estimated_memory <= available_for_protection:
+                        protected_wrappers.append(loaded)
+                        protected_memory += estimated_memory
+
+        unloaded = func(memory_required, device, keep_loaded + protected_wrappers, *args, **kwargs)
 
         REGISTRY.refresh_runtime_state()
         return unloaded
@@ -393,6 +973,15 @@ def _remember_original(key: str, value: Callable[..., Any]) -> Callable[..., Any
 def _patch_model_management_devices() -> None:
     import comfy.model_management as model_management
 
+    kind_by_function = {
+        "unet_offload_device": KIND_MODEL,
+        "unet_inital_load_device": KIND_MODEL,
+        "text_encoder_offload_device": KIND_CLIP,
+        "text_encoder_device": KIND_CLIP,
+        "vae_offload_device": KIND_VAE,
+        "vae_device": KIND_VAE,
+    }
+
     def wrap_device_func(name: str) -> None:
         key = f"model_management.{name}"
         original = _remember_original(key, getattr(model_management, name))
@@ -402,7 +991,8 @@ def _patch_model_management_devices() -> None:
         @functools.wraps(original)
         def wrapper(*args, **kwargs):
             result = original(*args, **kwargs)
-            if not REGISTRY.wants_gpu_offload(name):
+            kind = kind_by_function.get(name)
+            if not REGISTRY.wants_gpu_offload(kind):
                 return result
             gpu_device = model_management.get_torch_device()
             if getattr(gpu_device, "type", None) == "cpu":
@@ -430,6 +1020,10 @@ def _patch_model_management_devices() -> None:
     if model_management.load_models_gpu is original_load_models_gpu:
         model_management.load_models_gpu = _wrap_load_models_gpu(original_load_models_gpu)
 
+    original_model_unload = _remember_original("model_management.LoadedModel.model_unload", model_management.LoadedModel.model_unload)
+    if model_management.LoadedModel.model_unload is original_model_unload:
+        model_management.LoadedModel.model_unload = _wrap_loaded_model_unload(original_model_unload)
+
 
 def install_patches() -> None:
     global _PATCHED
@@ -440,6 +1034,7 @@ def install_patches() -> None:
     import comfy.controlnet as controlnet
     import comfy.diffusers_load as diffusers_load
     import comfy.model_management as model_management
+    import comfy.model_patcher as model_patcher
     import comfy.sd as comfy_sd
     import comfy.utils as comfy_utils
 
@@ -453,6 +1048,10 @@ def install_patches() -> None:
 
     _patch_model_management_devices()
 
+    original_model_patcher_detach = _remember_original("model_patcher.ModelPatcher.detach", model_patcher.ModelPatcher.detach)
+    if model_patcher.ModelPatcher.detach is original_model_patcher_detach:
+        model_patcher.ModelPatcher.detach = _wrap_model_patcher_detach(original_model_patcher_detach)
+
     original_load_checkpoint_guess_config = _remember_original(
         "sd.load_checkpoint_guess_config",
         comfy_sd.load_checkpoint_guess_config,
@@ -477,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

@@ -40,6 +40,7 @@ class LoadContext:
     source_path: str | None = None
     explicit_device: torch.device | None = None
     note: str | None = None
+    cache_key: str | None = None
 
 
 @dataclasses.dataclass(slots=True)
@@ -82,8 +83,10 @@ class ResidencyEntry:
     current_device: str | None = None
     last_method: str | None = None
     last_report: dict[str, Any] | None = None
+    loader_key: str | None = None
     notes: list[str] = dataclasses.field(default_factory=list)
     object_ref: weakref.ReferenceType[Any] | None = None
+    cached_object_ref: weakref.ReferenceType[Any] | None = None
 
     def is_alive(self) -> bool:
         return self.object_ref is not None and self.object_ref() is not None
@@ -91,6 +94,9 @@ class ResidencyEntry:
     def object(self) -> Any | None:
         return None if self.object_ref is None else self.object_ref()
 
+    def cached_object(self) -> Any | None:
+        return None if self.cached_object_ref is None else self.cached_object_ref()
+
     def as_dict(self) -> dict[str, Any]:
         basename = os.path.basename(self.source_path) if self.source_path else None
         return {
@@ -109,6 +115,7 @@ class ResidencyEntry:
             "current_device": self.current_device,
             "last_method": self.last_method,
             "last_report": self.last_report,
+            "loader_key": self.loader_key,
             "notes": list(self.notes),
             "alive": self.is_alive(),
         }
@@ -120,9 +127,37 @@ class ResidencyRegistry:
         self._entries: dict[str, ResidencyEntry] = {}
         self._reports_by_path: dict[str, LoadReport] = {}
         self._path_to_entry: dict[tuple[str, str], str] = {}
+        self._loader_key_to_entry: dict[tuple[str, str, str], str] = {}
         self._object_to_entry: weakref.WeakKeyDictionary[Any, str] = weakref.WeakKeyDictionary()
         self._policy = self._default_policy()
 
+    def _gpu_ingest_kinds(self, policy: str) -> set[str]:
+        if policy in {"prefer_gpu", "sticky_gpu"}:
+            return {KIND_MODEL, KIND_CHECKPOINT, KIND_CLIP, KIND_VAE, KIND_CONTROLNET, KIND_CLIP_VISION}
+        return set()
+
+    def _gpu_offload_kinds(self, policy: str) -> set[str]:
+        if policy == "sticky_gpu":
+            return {KIND_MODEL, KIND_CLIP, KIND_VAE, KIND_CONTROLNET}
+        if policy == "prefer_gpu":
+            return {KIND_MODEL, KIND_CLIP, KIND_CONTROLNET}
+        return set()
+
+    def _autopin_kinds(self, policy: str) -> set[str]:
+        if policy == "sticky_gpu":
+            return {KIND_MODEL, KIND_CLIP}
+        return set()
+
+    def default_priority(self, kind: str) -> int:
+        return {
+            KIND_MODEL: 300,
+            KIND_CHECKPOINT: 300,
+            KIND_CONTROLNET: 200,
+            KIND_CLIP: 150,
+            KIND_VAE: 100,
+            KIND_CLIP_VISION: 50,
+        }.get(kind, 0)
+
     def _default_policy(self) -> str:
         env_value = os.environ.get("COMFYUI_GPU_RESIDENT_POLICY", "").strip().lower()
         if env_value in POLICIES:
@@ -154,14 +189,21 @@ class ResidencyRegistry:
 
     def wants_gpu_ingest(self, kind: str | None = None) -> bool:
         policy = self.get_policy()
-        return policy in {"prefer_gpu", "sticky_gpu"}
+        if kind is None:
+            return bool(self._gpu_ingest_kinds(policy))
+        return kind in self._gpu_ingest_kinds(policy)
 
     def wants_gpu_offload(self, kind: str | None = None) -> bool:
         policy = self.get_policy()
-        return policy in {"prefer_gpu", "sticky_gpu"}
+        if kind is None:
+            return bool(self._gpu_offload_kinds(policy))
+        return kind in self._gpu_offload_kinds(policy)
 
     def autopin_on_bind(self, kind: str | None = None) -> bool:
-        return self.get_policy() == "sticky_gpu"
+        policy = self.get_policy()
+        if kind is None:
+            return bool(self._autopin_kinds(policy))
+        return kind in self._autopin_kinds(policy)
 
     def explicit_load_device(self, kind: str, source_path: str | None = None) -> torch.device | None:
         if not self.wants_gpu_ingest(kind):
@@ -183,6 +225,7 @@ class ResidencyRegistry:
         source_path: str | None = None,
         explicit_device: torch.device | None = None,
         note: str | None = None,
+        cache_key: str | None = None,
     ) -> Iterator[None]:
         token = _LOAD_CONTEXT.set(
             LoadContext(
@@ -190,6 +233,7 @@ class ResidencyRegistry:
                 source_path=source_path,
                 explicit_device=explicit_device,
                 note=note,
+                cache_key=cache_key,
             )
         )
         try:
@@ -222,7 +266,12 @@ class ResidencyRegistry:
         )
         with self._lock:
             self._reports_by_path[path] = report
-            entry_id = self._path_to_entry.get((kind, path))
+            entry_id = None
+            ctx = self.current_context()
+            if ctx is not None and ctx.cache_key is not None:
+                entry_id = self._loader_key_to_entry.get((kind, path, ctx.cache_key))
+            else:
+                entry_id = self._path_to_entry.get((kind, path))
             if entry_id is not None:
                 entry = self._entries.get(entry_id)
                 if entry is not None:
@@ -242,6 +291,40 @@ class ResidencyRegistry:
         basename = os.path.basename(source_path) or "anonymous"
         return f"{kind}:{basename}:{len(self._entries) + 1}"
 
+    def _clear_object_binding(self, obj: Any | None, entry_id: str) -> None:
+        if obj is None:
+            return
+        try:
+            if self._object_to_entry.get(obj) == entry_id:
+                self._object_to_entry.pop(obj, None)
+        except TypeError:
+            pass
+        try:
+            if getattr(obj, "__gpu_resident_loader_entry_id__", None) == entry_id:
+                delattr(obj, "__gpu_resident_loader_entry_id__")
+        except (AttributeError, TypeError):
+            pass
+
+    def _tag_object_with_entry(self, obj: Any, entry_id: str) -> None:
+        try:
+            setattr(obj, "__gpu_resident_loader_entry_id__", entry_id)
+            try:
+                self._object_to_entry.pop(obj, None)
+            except TypeError:
+                pass
+            return
+        except (AttributeError, TypeError):
+            _LOG.debug(
+                "GPU Resident Loader: could not tag object %r with residency entry id %s",
+                type(obj),
+                entry_id,
+            )
+
+        try:
+            self._object_to_entry[obj] = entry_id
+        except TypeError:
+            pass
+
     def bind_object(
         self,
         obj: Any,
@@ -249,15 +332,22 @@ class ResidencyRegistry:
         source_path: str,
         kind: str,
         sticky: bool | None = None,
-        priority: int = 0,
+        priority: int | None = None,
         note: str | None = None,
+        loader_key: str | None = None,
+        reusable_obj: Any | None = None,
     ) -> ResidencyEntry:
         if obj is None:
             raise ValueError("Cannot bind None into residency registry")
 
         with self._lock:
             old_key: tuple[str, str] | None = None
-            entry_id = getattr(obj, "__gpu_resident_loader_entry_id__", None)
+            old_loader_key: tuple[str, str, str] | None = None
+            previous_obj: Any | None = None
+            previous_cached_obj: Any | None = None
+            entry_id = self._loader_key_to_entry.get((kind, source_path, loader_key)) if loader_key is not None else None
+            if entry_id is None:
+                entry_id = getattr(obj, "__gpu_resident_loader_entry_id__", None)
             if entry_id is None:
                 try:
                     entry_id = self._object_to_entry.get(obj)
@@ -266,6 +356,10 @@ class ResidencyRegistry:
             if entry_id is not None and entry_id in self._entries:
                 entry = self._entries[entry_id]
                 old_key = (entry.kind, entry.source_path)
+                previous_obj = entry.object()
+                previous_cached_obj = entry.cached_object()
+                if entry.loader_key is not None:
+                    old_loader_key = (entry.kind, entry.source_path, entry.loader_key)
             else:
                 entry_id = self._make_entry_id(kind, source_path)
                 entry = ResidencyEntry(
@@ -273,45 +367,46 @@ class ResidencyRegistry:
                     kind=kind,
                     source_path=source_path,
                     sticky=self.autopin_on_bind(kind) if sticky is None else bool(sticky),
-                    priority=int(priority),
+                    priority=self.default_priority(kind) if priority is None else int(priority),
                 )
                 self._entries[entry_id] = entry
                 self._path_to_entry[(kind, source_path)] = entry_id
-                try:
-                    setattr(obj, "__gpu_resident_loader_entry_id__", entry_id)
-                    try:
-                        self._object_to_entry.pop(obj, None)
-                    except TypeError:
-                        pass
-                except (AttributeError, TypeError):
-                    _LOG.debug(
-                        "GPU Resident Loader: could not tag object %r with residency entry id %s",
-                        type(obj),
-                        entry_id,
-                    )
+
+            cache_obj = obj if reusable_obj is None else reusable_obj
+            for stale_obj in (previous_obj, previous_cached_obj):
+                if stale_obj is None or stale_obj is obj or stale_obj is cache_obj:
+                    continue
+                self._clear_object_binding(stale_obj, entry_id)
 
             try:
                 entry.object_ref = weakref.ref(obj)
-                if getattr(obj, "__gpu_resident_loader_entry_id__", None) is None:
-                    try:
-                        self._object_to_entry[obj] = entry_id
-                    except TypeError:
-                        pass
             except TypeError:
                 entry.object_ref = None
                 _LOG.debug(
                     "GPU Resident Loader: object %r is not weak-referenceable; tracking metadata only",
                     type(obj),
                 )
+            self._tag_object_with_entry(obj, entry_id)
+            try:
+                entry.cached_object_ref = weakref.ref(cache_obj)
+            except TypeError:
+                entry.cached_object_ref = entry.object_ref
             entry.sticky = entry.sticky if sticky is None else bool(sticky)
-            entry.priority = int(priority)
+            entry.priority = entry.priority if priority is None else int(priority)
             entry.source_path = source_path
             entry.kind = kind
+            entry.loader_key = loader_key
             new_key = (entry.kind, entry.source_path)
+            new_loader_key = (entry.kind, entry.source_path, entry.loader_key) if entry.loader_key is not None else None
             if old_key is not None and old_key != new_key:
                 if self._path_to_entry.get(old_key) == entry_id:
                     self._path_to_entry.pop(old_key, None)
             self._path_to_entry[new_key] = entry_id
+            if old_loader_key is not None and old_loader_key != new_loader_key:
+                if self._loader_key_to_entry.get(old_loader_key) == entry_id:
+                    self._loader_key_to_entry.pop(old_loader_key, None)
+            if new_loader_key is not None:
+                self._loader_key_to_entry[new_loader_key] = entry_id
             entry.last_touched = _now()
             if note:
                 entry.notes.append(note)
@@ -321,21 +416,49 @@ class ResidencyRegistry:
                 entry.last_report = report.as_dict()
                 entry.current_device = report.actual_device
 
-        self.refresh_runtime_state()
         return entry
 
+    def lookup_live_object(self, *, kind: str, source_path: str, loader_key: str) -> Any | None:
+        with self._lock:
+            entry_id = self._loader_key_to_entry.get((kind, source_path, loader_key))
+            if entry_id is None:
+                return None
+            entry = self._entries.get(entry_id)
+            if entry is None:
+                return None
+            obj = entry.cached_object()
+            if obj is None:
+                if entry.cached_object_ref is not None:
+                    return None
+                obj = entry.object()
+            if obj is None:
+                return None
+            entry.last_touched = _now()
+            return obj
+
+    def _entry_id_for_object(self, obj: Any) -> str | None:
+        current = obj
+        seen: set[int] = set()
+        while current is not None and id(current) not in seen:
+            seen.add(id(current))
+            entry_id = getattr(current, "__gpu_resident_loader_entry_id__", None)
+            if entry_id is None:
+                try:
+                    entry_id = self._object_to_entry.get(current)
+                except TypeError:
+                    entry_id = None
+            if entry_id is not None:
+                return entry_id
+            current = getattr(current, "parent", None)
+        return None
+
     def entry_for_object(self, obj: Any) -> ResidencyEntry | None:
         if obj is None:
             return None
-        entry_id = getattr(obj, "__gpu_resident_loader_entry_id__", None)
-        if entry_id is None:
-            try:
-                entry_id = self._object_to_entry.get(obj)
-            except TypeError:
-                entry_id = None
-        if entry_id is None:
-            return None
         with self._lock:
+            entry_id = self._entry_id_for_object(obj)
+            if entry_id is None:
+                return None
             return self._entries.get(entry_id)
 
     def set_sticky(self, obj: Any, sticky: bool, priority: int | None = None) -> ResidencyEntry | None:
@@ -369,8 +492,9 @@ class ResidencyRegistry:
                     continue
                 entry = self.entry_for_object(loaded.model)
                 if entry is not None and entry.sticky:
-                    output.append(loaded)
-        return output
+                    output.append((entry.priority, entry.last_touched, loaded))
+        output.sort(key=lambda item: (item[0], item[1]), reverse=True)
+        return [item[2] for item in output]
 
     def refresh_runtime_state(self) -> None:
         try:
@@ -401,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()