Invalidate HookRegistry child-registries cache on enable/disable cache · huggingface/diffusers@42f7a74 · GitHub
Skip to content

Commit 42f7a74

Browse files
Invalidate HookRegistry child-registries cache on enable/disable cache
1 parent 1f9c201 commit 42f7a74

3 files changed

Lines changed: 44 additions & 0 deletions

File tree

‎src/diffusers/hooks/hooks.py‎

Lines changed: 15 additions & 0 deletions

‎src/diffusers/models/cache_utils.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ def enable_cache(self, config) -> None:
6969
from ..hooks import (
7070
FasterCacheConfig,
7171
FirstBlockCacheConfig,
72+
HookRegistry,
7273
MagCacheConfig,
7374
PyramidAttentionBroadcastConfig,
7475
TaylorSeerCacheConfig,
@@ -101,6 +102,11 @@ def enable_cache(self, config) -> None:
101102
else:
102103
raise ValueError(f"Cache config {type(config)} is not supported.")
103104

105+
# Applying a cache technique registers hooks on child blocks, which stales any
106+
# `_child_registries_cache` built earlier (e.g. by a prior `cache_context`). Invalidate it
107+
# so `_set_context` reaches the freshly-registered block hooks.
108+
HookRegistry.check_if_exists_or_initialize(self).invalidate_child_registries_cache()
109+
104110
self._cache_config = config
105111

106112
def disable_cache(self) -> None:
@@ -144,6 +150,9 @@ def disable_cache(self) -> None:
144150
else:
145151
raise ValueError(f"Cache config {type(self._cache_config)} is not supported.")
146152

153+
# Removing the cache hooks stales any `_child_registries_cache` that included them.
154+
registry.invalidate_child_registries_cache()
155+
147156
self._cache_config = None
148157

149158
def _reset_stateful_cache(self, recurse: bool = True) -> None:

‎tests/hooks/test_hooks.py‎

Lines changed: 20 additions & 0 deletions

0 commit comments

Comments
 (0)