Cap AutoOffloadStrategy's exhaustive offload search for large candidate counts

src/diffusers/modular_pipelines/components_manager.py

@@ -146,6 +146,13 @@ def custom_offload_with_hook(
     return user_hook
 
 
+# search_best_candidate tries every combination of models to find the smallest one
+# that frees enough memory, so it's O(2^n). fine for a few models, but ComponentsManager
+# is meant to hold components from multiple pipelines at once so n can get big. past
+# this many candidates just do a greedy pick instead of grinding through all of them.
+_MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES = 20
+
+
 # this is the class that user can customize to implement their own offload strategy
 class AutoOffloadStrategy:
     """
@@ -198,6 +205,26 @@ def search_best_candidate(module_sizes, min_memory_offload):
             larger than `min_memory_offload`
             """
             model_ids = list(module_sizes.keys())
+
+            if len(model_ids) > _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES:
+                # too many to brute-force all combinations, so fall back to a greedy
+                # pick instead. module_sizes is sorted largest-to-smallest (for the
+                # exhaustive path below), but per the TODO above we want to keep
+                # bigger models on GPU when we can, so walk it smallest-to-largest
+                # here and evict small models first.
+                logger.warning(
+                    f"{len(model_ids)} candidate models to offload exceeds "
+                    f"{_MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES}, using a greedy search instead of an exhaustive one"
+                )
+                greedy_candidate = []
+                greedy_size = 0
+                for model_id in reversed(model_ids):
+                    if greedy_size >= min_memory_offload:
+                        break
+                    greedy_candidate.append(model_id)
+                    greedy_size += module_sizes[model_id]
+                return tuple(greedy_candidate) if greedy_size >= min_memory_offload else None
+
             best_candidate = None
             best_size = float("inf")
             for r in range(1, len(model_ids) + 1):

tests/modular_pipelines/test_components_manager.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import gc
+import time
 from unittest import mock
 
 import pytest
@@ -26,7 +27,10 @@
 
 
 if is_accelerate_available():
-    from diffusers.modular_pipelines.components_manager import AutoOffloadStrategy
+    from diffusers.modular_pipelines.components_manager import (
+        _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES,
+        AutoOffloadStrategy,
+    )
 
 
 # The offload logic deals in bytes. We keep the test models tiny (a few KB of real
@@ -209,6 +213,51 @@ def test_strategy_memory_reserve_margin_changes_decision(self):
             memory_reserve_margin=3 * UNIT,
         ) == ["c"]
 
+    def test_strategy_uses_greedy_fallback_above_exhaustive_search_limit(self):
+        # past _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES, search_best_candidate stops trying
+        # every combination and grabs models greedily instead. check it still frees
+        # enough memory and does it fast - 28 models is 2^28 combos the old way.
+        num_hooks = _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES + 4
+        hook_sizes = {f"model_{i}": UNIT for i in range(num_hooks)}
+
+        start = time.perf_counter()
+        selected = self._select_offload(
+            incoming_footprint=15 * UNIT,
+            free_bytes=10 * UNIT,
+            hook_sizes=hook_sizes,
+            memory_reserve_margin=0,
+        )
+        elapsed = time.perf_counter() - start
+
+        # usable = 10 - 0 = 10, incoming needs 15 -> must free 5. All 28 models are the
+        # same size here so direction doesn't matter for this assertion, just that it
+        # picks the right count and stays fast - see the next test for direction.
+        assert len(selected) == 5
+        assert sum(hook_sizes[model_id] for model_id in selected) == 5 * UNIT
+        assert elapsed < 1.0
+
+    def test_strategy_greedy_fallback_evicts_smallest_models_first(self):
+        # the greedy fallback should keep bigger models resident and evict smaller
+        # ones first - matches the YiYi/Dhruv TODO above search_best_candidate ("we
+        # would tend to keep the larger models on GPU more often"). Use one big model
+        # plus enough small ones to push past the exhaustive-search cap.
+        num_small = _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES + 3
+        hook_sizes = {"big": 5 * UNIT}
+        hook_sizes.update({f"small_{i}": UNIT for i in range(num_small)})
+
+        selected = self._select_offload(
+            incoming_footprint=8 * UNIT,
+            free_bytes=5 * UNIT,
+            hook_sizes=hook_sizes,
+            memory_reserve_margin=0,
+        )
+
+        # usable = 5, incoming needs 8 -> must free 3. Three small models cover that
+        # on their own, so "big" should be left alone.
+        assert "big" not in selected
+        assert len(selected) == 3
+        assert all(model_id.startswith("small_") for model_id in selected)
+
     # ------------------------------------------------------------------
     # Registry tests (hardware-independent)
     # ------------------------------------------------------------------