Address PR review feedback

[codex] Fix same-device VRAM eviction

README.md

@@ -62,6 +62,8 @@ Current patch surface:
 
 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**

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

@@ -14,7 +14,14 @@ import comfy.utils
 from comfy.cli_args import PerformanceFeature, args
 from comfy.ldm.modules.attention import attention_pytorch, wrap_attn
 
-from .patches import infer_unet_prefix_from_keys, load_safetensors_state_dict
+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
 
 
@@ -328,6 +335,113 @@ def _warn_pickle_checkpoint_gpu_compatibility(loader_name: str, path: str) -> No
     )
 
 
+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,
     *,
@@ -485,6 +599,7 @@ def _load_resident_diffusion_model(
     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)
@@ -506,6 +621,18 @@ def _load_resident_diffusion_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,
@@ -568,6 +695,7 @@ 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)
@@ -576,17 +704,74 @@ def _load_checkpoint_clip_only(
         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)
-    with REGISTRY.load_context(kind=KIND_CLIP, source_path=ckpt_path, explicit_device=explicit_device, cache_key=loader_key):
-        sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
-    _, 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,
+    _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
 
@@ -596,6 +781,7 @@ 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)
@@ -604,17 +790,49 @@ def _load_checkpoint_vae_only(
         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)
-    with REGISTRY.load_context(kind=KIND_VAE, source_path=ckpt_path, explicit_device=explicit_device, cache_key=loader_key):
-        sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
-    _, _, 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,
+    _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
 
@@ -649,70 +867,37 @@ def _load_full_checkpoint(
         _LOG.info("%s: reusing live checkpoint outputs for %s", loader_name, ckpt_path)
         return model, clip, vae
 
-    if len(missing) == 1:
-        if missing[0] == "model":
-            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,
-            )
-        elif missing[0] == "clip":
-            clip = _load_checkpoint_clip_only(
-                ckpt_path=ckpt_path,
-                policy_override=policy_override,
-                loader_name=loader_name,
-            )
-        else:
-            vae = _load_checkpoint_vae_only(
-                ckpt_path=ckpt_path,
-                policy_override=policy_override,
-                loader_name=loader_name,
-            )
-        return model, clip, vae
-
-    _warn_pickle_checkpoint_gpu_compatibility(loader_name, ckpt_path)
-    model_options = _build_model_options(weight_dtype)
-    explicit_device = REGISTRY.explicit_load_device(kind=KIND_CHECKPOINT, source_path=ckpt_path)
-    with _temporary_backend_flags(
-        cublas=patch_cublaslinear,
-        fp16_accumulation=enable_fp16_accumulation,
-    ):
-        with REGISTRY.load_context(
-            kind=KIND_CHECKPOINT,
+    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,
-            explicit_device=explicit_device,
-            cache_key=_make_loader_key(
-                loader_name,
-                component="full_checkpoint",
-                weight_dtype=weight_dtype,
-                compute_dtype=compute_dtype,
-                patch_cublaslinear=patch_cublaslinear,
-                sage_attention=sage_attention,
-                enable_fp16_accumulation=enable_fp16_accumulation,
-                policy=_effective_policy_name(policy_override),
-            ),
-        ):
-            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,
+            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,
         )
-        _apply_model_postload_options(model, compute_dtype=compute_dtype, sage_attention=sage_attention)
-
-    _bind_model_for_reuse(model, source_path=ckpt_path, note="checkpoint model", loader_key=model_key)
-    _bind_clip_for_reuse(clip, source_path=ckpt_path, note="checkpoint clip", loader_key=clip_key)
-    _bind_vae_for_reuse(vae, source_path=ckpt_path, note="checkpoint vae", loader_key=vae_key)
     return model, clip, vae
 
 
@@ -876,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 reuses live checkpoint components when possible and only falls back to a full checkpoint load when multiple outputs are missing."
+        "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(

nodes.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import json
+import logging
 from typing import Any
 
 import comfy.model_management as model_management
 
+from .cleanup import unload_loaded_model
 from .kj_loader import (
     CheckpointClipLoaderResident,
     CheckpointLoaderResident,
@@ -15,6 +17,8 @@ from .kj_loader import (
 )
 from .residency import REGISTRY
 
+_LOG = logging.getLogger(__name__)
+
 
 def _entry_report_json(obj: Any) -> str:
     entry = REGISTRY.entry_for_object(obj)
@@ -59,9 +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()

patches.py

@@ -1,13 +1,21 @@
 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,
+    unload_loaded_model,
+)
 from .residency import (
     KIND_CHECKPOINT,
     KIND_CLIP,
@@ -25,6 +33,179 @@ _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
+_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:
@@ -44,6 +225,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()
@@ -233,12 +448,10 @@ def _patched_load_torch_file(ckpt, safe_load=False, device=None, return_metadata
     if lowered.endswith((".safetensors", ".sft")):
         try:
             selected_keys = None
-            if ctx is not None and ctx.kind == KIND_MODEL:
-                with safe_open(ckpt, framework="pt", device="cpu") as handle:
-                    all_keys = list(handle.keys())
-                prefix = infer_unet_prefix_from_keys(all_keys)
-                candidate_keys = {key: key[len(prefix):] for key in all_keys if key.startswith(prefix)}
-                selected_keys = candidate_keys or 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,
@@ -248,9 +461,9 @@ def _patched_load_torch_file(ckpt, safe_load=False, device=None, return_metadata
             )
             if load_mode == "aimdo_cpu":
                 method = "safetensors_aimdo_cpu"
-            elif selected_keys is not None:
-                method = "safetensors_cpu_then_copy_to_cuda_model_only" if load_mode == "cpu_then_copy" else (
-                    "safetensors_gpu_direct_model_only" if requested_device.type == "cuda" else "safetensors_cpu_model_only"
+            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 (
@@ -396,6 +609,62 @@ def _wrap_load_clip(func: Callable[..., Any]) -> Callable[..., Any]:
 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)
@@ -405,6 +674,44 @@ 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):
@@ -501,6 +808,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
@@ -511,6 +822,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
 
@@ -524,6 +836,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,

residency.py

@@ -436,18 +436,29 @@ class ResidencyRegistry:
             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: