Protect related external fallback entries

Docstrings generation was requested by @xmarre.

The following files were modified:

* `cleanup.py`
* `patches.py`

cleanup.py

@@ -130,7 +130,27 @@ def _trim_candidates(
     respect_sticky: bool,
     sticky_floor_priority: int,
     keep_models: tuple[Any, ...],
+    include_external: bool | None = None,
 ) -> list[tuple[Any, Any, bool, bool]]:
+    """
+    Collects and returns eviction/unload candidates from in-memory models and, optionally, external integrations.
+    
+    Filters currently loaded models by the given device, excludes dead or missing models, and skips models listed in `keep_models`. If `respect_sticky` is true, entries whose registry metadata mark them as sticky and whose priority meets or exceeds `sticky_floor_priority` are flagged so they are treated as higher-priority to keep. When `include_external` is true (or when `include_external` is None and external trimming is enabled), candidates from the external registry are included.
+    
+    Parameters:
+        device: Device filter for candidates; if not None only candidates matching this device are considered.
+        respect_sticky (bool): Whether to respect sticky registry entries when computing candidate priority.
+        sticky_floor_priority (int): Minimum priority value for a registry entry to be considered sticky.
+        keep_models (tuple[Any, ...]): Objects that must not be selected as candidates.
+        include_external (bool | None): If True include external-registry candidates; if False exclude them; if None defer to runtime external_trim_enabled().
+    
+    Returns:
+        list[tuple[Any, Any, bool, bool]]: A sorted list of tuples (candidate_obj, registry_entry_or_None, sticky_respected, is_external_candidate).
+        - candidate_obj: The loaded model object (internal) or the external object.
+        - registry_entry_or_None: Registry metadata for the candidate, or None if unavailable.
+        - sticky_respected: `True` when the candidate is marked sticky and meets `sticky_floor_priority`.
+        - is_external_candidate: `True` for candidates originating from the external registry.
+    """
     candidates: list[tuple[Any, Any, bool, bool]] = []
     ensure_external_integrations_installed()
     for loaded in list(model_management.current_loaded_models):
@@ -154,7 +174,8 @@ def _trim_candidates(
         )
         candidates.append((loaded, entry, sticky_respected, False))
 
-    if external_trim_enabled():
+    should_include_external = external_trim_enabled() if include_external is None else bool(include_external)
+    if should_include_external:
         for external_obj, entry, sticky_respected in EXTERNAL_REGISTRY.candidates(
             device=device,
             respect_sticky=respect_sticky,
@@ -175,7 +196,38 @@ def trim_resident_vram(
     sticky_floor_priority: int,
     allow_partial_unload: bool,
     keep_models: tuple[Any, ...] = (),
+    include_external: bool | None = None,
 ) -> dict[str, Any]:
+    """
+    Trim resident models (and optional external integration objects) until the requested amount of free VRAM is available or a stopping condition occurs.
+    
+    Tries to free GPU memory on `device` by evicting or unloading loaded models (and, optionally, external candidates). Respects sticky/priority hints, can perform partial unloads of pinned RAM when supported, and records each attempted action in the returned report.
+    
+    Parameters:
+        device (str | torch.device | None): Target device to free (defaults to model_management.get_torch_device()).
+        target_free_vram_bytes (int): Desired amount of free VRAM, in bytes.
+        respect_sticky (bool): If true, prefer protecting entries marked as sticky with sufficient priority.
+        sticky_floor_priority (int): Minimum priority required for a sticky entry to be respected.
+        allow_partial_unload (bool): If true, allow partial unloads and attempts to free pinned host RAM before full eviction.
+        keep_models (tuple[Any, ...]): Sequence of model objects that must not be unloaded (exact matches or recognized clones).
+        include_external (bool | None): If None, use external_trim_enabled() at runtime; otherwise force inclusion/exclusion of external candidates.
+    
+    Returns:
+        dict[str, Any]: Report of the trimming operation containing:
+            - status: "met_target", "partial", or "error".
+            - stopped_reason: reason the loop stopped (e.g., "target_met", "no_candidates", "no_progress", "error").
+            - target_met (bool): whether the target free VRAM was reached.
+            - device (str): string form of the device used.
+            - target_free_vram_bytes (int), free_before_bytes (int), free_after_bytes (int).
+            - freed_vram_bytes (int): total freed VRAM during this call.
+            - respect_sticky (bool), sticky_floor_priority (int), allow_partial_unload (bool).
+            - external_trim_enabled (bool): computed flag indicating whether external candidates were considered.
+            - actions (list[dict]): ordered per-candidate action records; each entry includes metadata such as
+              entry_id, basename, tracked, external_candidate, sticky_respected, priority,
+              need_before_bytes, loaded_before_bytes, freed_pinned_ram_bytes,
+              mode (e.g., "full_unload", "partial_unload", "external_evict", "error"),
+              loaded_after_bytes, freed_vram_bytes, free_after_bytes, and any warnings/errors.
+    """
     ensure_external_integrations_installed()
     cleanup_models_gc = getattr(model_management, "cleanup_models_gc", None)
     if callable(cleanup_models_gc):
@@ -203,6 +255,7 @@ def trim_resident_vram(
             respect_sticky=respect_sticky,
             sticky_floor_priority=sticky_floor_priority,
             keep_models=keep_models,
+            include_external=include_external,
         )
         if not candidates:
             stopped_reason = "no_candidates"
@@ -306,7 +359,7 @@ def trim_resident_vram(
         "respect_sticky": bool(respect_sticky),
         "sticky_floor_priority": int(sticky_floor_priority),
         "allow_partial_unload": bool(allow_partial_unload),
-        "external_trim_enabled": bool(external_trim_enabled()),
+        "external_trim_enabled": bool(external_trim_enabled() if include_external is None else include_external),
         "actions": actions,
     }
 

external_residency.py

@@ -403,6 +403,36 @@ def external_trim_enabled() -> bool:
     return value in {"1", "true", "yes", "on"}
 
 
+def external_objects_for_models(models: tuple[Any, ...] | list[Any]) -> tuple[Any, ...]:
+    """
+    Returns registered external objects that are part of any supplied model wrapper chain.
+
+    This lets callers preserve external cache entries when they are associated with a kept
+    model through wrapper indirection rather than exact object identity.
+    """
+    related_ids: set[int] = set()
+    for model in models:
+        for related in _iter_seedvr2_wrapper_chain(model):
+            related_ids.add(id(related))
+
+    if not related_ids:
+        return ()
+
+    EXTERNAL_REGISTRY.refresh_runtime_state()
+    matches: list[Any] = []
+    seen_ids: set[int] = set()
+    with EXTERNAL_REGISTRY._lock:
+        for entry in EXTERNAL_REGISTRY._entries.values():
+            obj = entry.object()
+            if obj is None:
+                continue
+            obj_id = id(obj)
+            if obj_id in related_ids and obj_id not in seen_ids:
+                matches.append(obj)
+                seen_ids.add(obj_id)
+    return tuple(matches)
+
+
 def _call_seedvr2_method_with_optional_expected_model(
     method: Callable[..., Any],
     *args: Any,
@@ -693,4 +723,4 @@ def ensure_external_integrations_installed() -> None:
                     "GPU Resident Loader: failed to install SeedVR2 external integration from %s: %s",
                     normalized_file,
                     exc,
-                )
+                )

patches.py

@@ -18,6 +18,7 @@ from .cleanup import (
     trim_resident_vram,
     unload_loaded_model,
 )
+from .external_residency import EXTERNAL_REGISTRY, external_objects_for_models, external_trim_enabled
 from .residency import (
     KIND_CHECKPOINT,
     KIND_CLIP,
@@ -668,14 +669,27 @@ def _scaled_batch_memory(total_memory_used: int, total_batch_count: int, batch_n
     return max(1, (total_memory * current_batch + total_batches - 1) // total_batches)
 
 
+def _sticky_vae_free_memory(*, device: Any, patcher: Any) -> int:
+    import comfy.model_management as model_management
+
+    get_free_memory = getattr(model_management, "get_free_memory", None)
+    if callable(get_free_memory):
+        try:
+            return max(0, int(get_free_memory(device)))
+        except Exception:
+            pass
+
+    return max(0, int(patcher.get_free_memory(device)))
+
+
 def _prepare_sticky_vae_batch(
     *,
     device: Any,
     patcher: Any,
     total_memory_used: int,
     total_batch_count: int,
-) -> tuple[int, bool]:
-    free_memory = int(patcher.get_free_memory(device))
+) -> tuple[int, int, bool]:
+    free_memory = _sticky_vae_free_memory(device=device, patcher=patcher)
     batch_number = _sticky_safe_batch_number(
         batch_count=total_batch_count,
         free_memory=free_memory,
@@ -685,7 +699,7 @@ def _prepare_sticky_vae_batch(
     batch_memory_used = _scaled_batch_memory(total_memory_used, total_batch_count, batch_number)
 
     if REGISTRY.get_policy() != "sticky_gpu" or device is None:
-        return batch_number, False
+        return batch_number, batch_memory_used, False
 
     target_free = _sticky_protection_target(batch_memory_used, device)
     if free_memory < target_free:
@@ -701,7 +715,7 @@ def _prepare_sticky_vae_batch(
         except Exception as exc:
             _LOG.debug("GPU Resident Loader: proactive VAE trim failed for %s bytes: %s", batch_memory_used, exc)
 
-        free_memory = int(patcher.get_free_memory(device))
+        free_memory = _sticky_vae_free_memory(device=device, patcher=patcher)
         batch_number = _sticky_safe_batch_number(
             batch_count=total_batch_count,
             free_memory=free_memory,
@@ -719,7 +733,7 @@ def _prepare_sticky_vae_batch(
             target_free,
             batch_memory_used,
         )
-    return batch_number, should_tile
+    return batch_number, batch_memory_used, should_tile
 
 
 def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]:
@@ -741,8 +755,7 @@ def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]:
                 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)
-            batch_number, should_tile = _prepare_sticky_vae_batch(
+            batch_number, batch_memory_used, should_tile = _prepare_sticky_vae_batch(
                 device=self.device,
                 patcher=self.patcher,
                 total_memory_used=memory_used,
@@ -751,6 +764,11 @@ def _wrap_vae_encode(func: Callable[..., Any]) -> Callable[..., Any]:
             if should_tile:
                 do_tile = True
             else:
+                model_management.load_models_gpu(
+                    [self.patcher],
+                    memory_required=batch_memory_used,
+                    force_full_load=self.disable_offload,
+                )
                 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)
@@ -803,8 +821,7 @@ def _wrap_vae_decode(func: Callable[..., Any]) -> Callable[..., Any]:
             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)
-            batch_number, should_tile = _prepare_sticky_vae_batch(
+            batch_number, batch_memory_used, should_tile = _prepare_sticky_vae_batch(
                 device=self.device,
                 patcher=self.patcher,
                 total_memory_used=memory_used,
@@ -812,17 +829,21 @@ def _wrap_vae_decode(func: Callable[..., Any]) -> Callable[..., Any]:
             )
 
             preallocated = False
-            if not should_tile and getattr(self.first_stage_model, "comfy_has_chunked_io", False):
-                pixel_samples = torch.empty(
-                    self.first_stage_model.decode_output_shape(samples_in.shape),
-                    device=self.output_device,
-                    dtype=self.vae_output_dtype(),
-                )
-                preallocated = True
-
             if should_tile:
                 do_tile = True
             else:
+                model_management.load_models_gpu(
+                    [self.patcher],
+                    memory_required=batch_memory_used,
+                    force_full_load=self.disable_offload,
+                )
+                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:
@@ -972,6 +993,28 @@ def _wrap_model_patcher_detach(func: Callable[..., Any]) -> Callable[..., Any]:
 
 
 def _wrap_free_memory(func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    Wraps a free-memory function to enforce sticky-GPU protection and external fallback trimming.
+    
+    When the registry policy is "sticky_gpu" and a device is provided, the wrapper:
+    - Reserves VRAM for sticky-loaded models by attempting a pre-trim to a computed protection target.
+    - Protects a subset of sticky-loaded wrappers from unloading when calling the original function by adding them to `keep_loaded`.
+    - After the original free-memory call, refreshes both REGISTRY and EXTERNAL_REGISTRY runtime state.
+    - If external trimming is available and still needed, attempts a fallback trim that includes external residency.
+      Dynamic free-memory calls are excluded because Comfy reduces their effective target internally.
+    
+    Parameters:
+        memory_required: Number of bytes the caller needs to free.
+        device: Target device for which memory is being freed (may be None).
+        keep_loaded: Iterable of loaded-wrapper objects that must be kept; the wrapper may extend this list with additional protected wrappers.
+    
+    Returns:
+        The value returned by the wrapped `func`.
+        
+    Notes:
+    - The wrapper may call `trim_resident_vram` and `model_management.get_free_memory`; exceptions from trimming or free-memory queries are caught and logged, not propagated.
+    - Side effects include invoking trims and refreshing runtime state on REGISTRY and EXTERNAL_REGISTRY.
+    """
     @functools.wraps(func)
     def wrapper(memory_required, device, keep_loaded=None, *args, **kwargs):
         import comfy.model_management as model_management
@@ -1026,6 +1069,42 @@ def _wrap_free_memory(func: Callable[..., Any]) -> Callable[..., Any]:
         unloaded = func(memory_required, device, keep_loaded + protected_wrappers, *args, **kwargs)
 
         REGISTRY.refresh_runtime_state()
+        EXTERNAL_REGISTRY.refresh_runtime_state()
+
+        for_dynamic = bool(kwargs.get("for_dynamic", args[0] if args else False))
+        if device is not None and not external_trim_enabled() and not for_dynamic:
+            fallback_target = memory_required
+            if REGISTRY.get_policy() == "sticky_gpu":
+                fallback_target = max(fallback_target, _sticky_protection_target(memory_required, device))
+            try:
+                free_now = model_management.get_free_memory(device)
+            except Exception:
+                free_now = None
+            if free_now is not None and int(free_now) < int(fallback_target):
+                protected_models = tuple(
+                    model
+                    for model in (getattr(loaded_wrapper, "model", None) for loaded_wrapper in keep_loaded + protected_wrappers)
+                    if model is not None
+                )
+                keep_models = protected_models + external_objects_for_models(protected_models)
+                try:
+                    trim_resident_vram(
+                        device=device,
+                        target_free_vram_bytes=int(fallback_target),
+                        respect_sticky=True,
+                        sticky_floor_priority=0,
+                        allow_partial_unload=True,
+                        keep_models=keep_models,
+                        include_external=True,
+                    )
+                except Exception as exc:
+                    _LOG.debug(
+                        "GPU Resident Loader: external fallback trim failed for free_memory(%s): %s",
+                        memory_required,
+                        exc,
+                    )
+                REGISTRY.refresh_runtime_state()
+                EXTERNAL_REGISTRY.refresh_runtime_state()
         return unloaded
 
     return wrapper