Invalidate HookRegistry child-registries cache on enable/disable cache

src/diffusers/hooks/hooks.py

@@ -274,6 +274,20 @@ def _set_context(self, name: str | None = None) -> None:
         for registry in self._get_child_registries():
             registry._set_context(name)
 
+    def invalidate_child_registries_cache(self) -> None:
+        """Invalidate the cached child-registry list across this module's tree.
+
+        `_get_child_registries` caches the registries it finds by walking `named_modules()`, keyed on the registry that
+        built it. Registering or removing hooks on descendant modules (e.g. block hooks added by `enable_cache`)
+        changes which modules carry a `_diffusers_hook`, which stales that cache. Call this after any operation that
+        adds or removes hooks in the subtree so the list is rebuilt on next use. Clears the cache on every registry in
+        the tree, since the same registries appear in ancestor caches.
+        """
+        for _, module in unwrap_module(self._module_ref).named_modules():
+            module = unwrap_module(module)
+            if hasattr(module, "_diffusers_hook"):
+                module._diffusers_hook._child_registries_cache = None
+
     def _get_child_registries(self) -> list["HookRegistry"]:
         """Return registries of child modules, using a cached list when available.
 

src/diffusers/models/cache_utils.py

@@ -69,6 +69,7 @@ def enable_cache(self, config) -> None:
         from ..hooks import (
             FasterCacheConfig,
             FirstBlockCacheConfig,
+            HookRegistry,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
@@ -101,6 +102,11 @@ def enable_cache(self, config) -> None:
         else:
             raise ValueError(f"Cache config {type(config)} is not supported.")
 
+        # Applying a cache technique registers hooks on child blocks, which stales any
+        # `_child_registries_cache` built earlier (e.g. by a prior `cache_context`). Invalidate it
+        # so `_set_context` reaches the freshly-registered block hooks.
+        HookRegistry.check_if_exists_or_initialize(self).invalidate_child_registries_cache()
+
         self._cache_config = config
 
     def disable_cache(self) -> None:
@@ -144,6 +150,9 @@ def disable_cache(self) -> None:
         else:
             raise ValueError(f"Cache config {type(self._cache_config)} is not supported.")
 
+        # Removing the cache hooks stales any `_child_registries_cache` that included them.
+        registry.invalidate_child_registries_cache()
+
         self._cache_config = None
 
     def _reset_stateful_cache(self, recurse: bool = True) -> None:

tests/hooks/test_hooks.py

@@ -17,6 +17,7 @@
 import pytest
 import torch
 
+from diffusers import FirstBlockCacheConfig, FluxTransformer2DModel
 from diffusers.hooks import HookRegistry, ModelHook
 from diffusers.training_utils import free_memory
 from diffusers.utils.logging import get_logger
@@ -200,6 +201,26 @@ def test_stateful_hook(self):
         assert registry.get_hook("stateful_add_hook").increment == 1
         assert torch.allclose(output1, output2)
 
+    def test_child_registries_cache_invalidation(self):
+        # Unit-level part of the regression test for
+        # https://github.com/huggingface/diffusers/issues/14037: the parent registry caches its
+        # child registries, so a hook registered on a child block after the cache was built is
+        # invisible to the parent until the cache is invalidated.
+        parent = HookRegistry.check_if_exists_or_initialize(self.model)
+
+        # Build the parent's child-registry cache while no block carries a hook yet.
+        assert parent._get_child_registries() == []
+
+        # Register a hook on a child block. The parent's cached (empty) list is now stale.
+        block = self.model.blocks[0]
+        child = HookRegistry.check_if_exists_or_initialize(block)
+        child.register_hook(AddHook(1), "add_hook")
+        assert parent._get_child_registries() == []  # still stale before invalidation
+
+        # Invalidating across the tree makes the new child registry reachable from the parent.
+        parent.invalidate_child_registries_cache()
+        assert child in parent._get_child_registries()
+
     def test_inference(self):
         registry = HookRegistry.check_if_exists_or_initialize(self.model)
         registry.register_hook(AddHook(1), "add_hook")
@@ -372,3 +393,46 @@ def test_invocation_order_stateful_last(self):
             .replace("\n", "")
         )
         assert output == expected_invocation_order_log
+
+
+def test_cache_context_after_enable_cache_with_prior_context():
+    # End-to-end regression test for https://github.com/huggingface/diffusers/issues/14037. Entering
+    # cache_context() before enable_cache() builds the model's child-registry cache without the block
+    # hooks. enable_cache() then registers FirstBlockCache hooks on the blocks; if the cache is not
+    # invalidated, _set_context() iterates the stale list and the new block hooks never receive a
+    # context, so the next cached forward raises "No context is set".
+    torch.manual_seed(0)
+    heads, head_dim = 2, 16
+    hidden = heads * head_dim
+    model = FluxTransformer2DModel(
+        patch_size=1,
+        in_channels=hidden,
+        num_layers=2,
+        num_single_layers=2,
+        attention_head_dim=head_dim,
+        num_attention_heads=heads,
+        joint_attention_dim=32,
+        pooled_projection_dim=16,
+        guidance_embeds=False,
+        axes_dims_rope=(2, 6, 8),
+    ).eval()
+
+    img_seq_len, txt_seq_len = 8, 4
+    inputs = {
+        "hidden_states": torch.randn(1, img_seq_len, hidden),
+        "encoder_hidden_states": torch.randn(1, txt_seq_len, 32),
+        "pooled_projections": torch.randn(1, 16),
+        "timestep": torch.tensor([1.0]),
+        "img_ids": torch.zeros(img_seq_len, 3),
+        "txt_ids": torch.zeros(txt_seq_len, 3),
+        "return_dict": False,
+    }
+
+    # Warmup pass inside a cache_context() while caching is disabled, then enable caching.
+    with torch.no_grad(), model.cache_context("cond"):
+        model(**inputs)
+    model.enable_cache(FirstBlockCacheConfig(threshold=0.1))
+
+    # Previously raised "No context is set"; the cache invalidation in enable_cache() fixes it.
+    with torch.no_grad(), model.cache_context("cond"):
+        model(**inputs)