Address PR review feedback

[codex] Fix same-device VRAM eviction

cleanup.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from contextlib import contextmanager
 from typing import Any
 
 import comfy.model_management as model_management
@@ -33,6 +34,69 @@ def _safe_is_dead(loaded) -> bool:
         return False
 
 
+def _device_matches(device_a, device_b) -> bool:
+    normalized_a = _normalize_trim_device(device_a)
+    normalized_b = _normalize_trim_device(device_b)
+    if normalized_a is None or normalized_b is None:
+        return False
+    if isinstance(normalized_a, torch.device) and isinstance(normalized_b, torch.device):
+        if normalized_a.type != normalized_b.type:
+            return False
+        if normalized_a.type == "cuda":
+            index_a = 0 if normalized_a.index is None else normalized_a.index
+            index_b = 0 if normalized_b.index is None else normalized_b.index
+            return index_a == index_b
+        return True
+    return str(normalized_a) == str(normalized_b)
+
+
+def _should_force_cpu_offload(model: Any, *, active_device=None, force: bool = False) -> bool:
+    if force:
+        return True
+    active_device = _normalize_trim_device(active_device)
+    if not isinstance(active_device, torch.device) or active_device.type != "cuda":
+        return False
+    return _device_matches(active_device, getattr(model, "offload_device", None))
+
+
+@contextmanager
+def _temporary_offload_device(model: Any, device):
+    if model is None or device is None or not hasattr(model, "offload_device"):
+        yield
+        return
+
+    original_device = model.offload_device
+    model.offload_device = device
+    try:
+        yield
+    finally:
+        model.offload_device = original_device
+
+
+def unload_loaded_model(
+    loaded,
+    *,
+    active_device: str | torch.device | None = None,
+    force_offload_to_cpu: bool = False,
+    unpatch_weights: bool = True,
+) -> bool:
+    # This helper is intentionally scoped to full unloads owned by this plugin.
+    model = getattr(loaded, "model", None)
+    if model is None:
+        return False
+
+    active_device = _normalize_trim_device(active_device if active_device is not None else getattr(loaded, "device", None))
+    force_cpu_offload = _should_force_cpu_offload(
+        model,
+        active_device=active_device,
+        force=force_offload_to_cpu,
+    )
+    unload_target = torch.device("cpu") if force_cpu_offload else None
+
+    with _temporary_offload_device(model, unload_target):
+        return loaded.model_unload(None, unpatch_weights=unpatch_weights)
+
+
 def _sort_key_for_candidate(entry, *, sticky_respected: bool) -> tuple[int, int, float]:
     if not sticky_respected:
         return (0, 0, getattr(entry, "last_touched", 0.0) if entry is not None else 0.0)

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,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import functools
 import json
 import logging
@@ -10,6 +11,11 @@ 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,
@@ -219,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()
@@ -569,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)
@@ -578,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):
@@ -674,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
@@ -684,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
 
@@ -697,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,