@@ -200,6 +200,26 @@ def test_stateful_hook(self):
200200 assert registry .get_hook ("stateful_add_hook" ).increment == 1
201201 assert torch .allclose (output1 , output2 )
202202
203+ def test_child_registries_cache_invalidation (self ):
204+ # Regression test for https://github.com/huggingface/diffusers/issues/14037: the parent
205+ # registry caches its child registries. A hook registered on a child block after that cache
206+ # was built is invisible to the parent until the cache is invalidated, so _set_context never
207+ # reaches the new child hook. invalidate_child_registries_cache() must rebuild it.
208+ parent = HookRegistry .check_if_exists_or_initialize (self .model )
209+
210+ # Build the parent's child-registry cache while no block carries a hook yet.
211+ assert parent ._get_child_registries () == []
212+
213+ # Register a hook on a child block. The parent's cached (empty) list is now stale.
214+ block = self .model .blocks [0 ]
215+ child = HookRegistry .check_if_exists_or_initialize (block )
216+ child .register_hook (AddHook (1 ), "add_hook" )
217+ assert parent ._get_child_registries () == [] # still stale before invalidation
218+
219+ # Invalidating across the tree makes the new child registry reachable from the parent.
220+ parent .invalidate_child_registries_cache ()
221+ assert child in parent ._get_child_registries ()
222+
203223 def test_inference (self ):
204224 registry = HookRegistry .check_if_exists_or_initialize (self .model )
205225 registry .register_hook (AddHook (1 ), "add_hook" )
0 commit comments