Invalidate HookRegistry child-registries cache on enable/disable cache

[SuryanshSS1011]

What does this PR do?

Fixes #14037.

HookRegistry._get_child_registries() caches the child-module registries it finds by walking named_modules(), and never invalidates that cache. But enable_cache() / disable_cache() add and remove block-level hooks, changing which modules carry a _diffusers_hook. If cache_context() is first entered while no block hooks exist (e.g. a warmup pass with caching disabled), the parent registry caches an incomplete child list. A later enable_cache(FirstBlockCacheConfig(...)) registers block hooks, but _set_context() still iterates the stale cache, so the new block StateManagers never receive a context and the next cached forward raises:

ValueError: No context is set. Please set a context before retrieving the state.

This adds HookRegistry.invalidate_child_registries_cache(), which clears the cached list across the module tree, and calls it from enable_cache() and disable_cache() after hooks are added/removed.

The staleness originates in register_hook / remove_hook, but those run on the child block registries, which can't reach the parent registry whose cache is stale. enable_cache / disable_cache operate on the root module, so invalidating there covers the reported scenario for every cache technique. Happy to move it into register_hook / remove_hook instead if you'd prefer it lower down.

The self-contained CPU reproduction from the issue passes after the fix, and a regression test is added in tests/hooks/test_hooks.py.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you self-review the diff against .ai/review-rules.md?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@DN6 @sayakpaul

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sergereview]

🤗 Serge says:

Clean, well-targeted fix for #14037. The stale-cache diagnosis is accurate: _get_child_registries() memoizes the child-registry walk and enable_cache/disable_cache change which modules carry a _diffusers_hook, so invalidating after hooks are added/removed is the right correction.

Correctness

• invalidate_child_registries_cache() walks the full tree and resets _child_registries_cache = None, which _get_child_registries() correctly treats as "rebuild on next use". Clearing every registry in the subtree (not just the root) is right, since a child registry can also appear in an ancestor's cache.
• Invalidation is placed in enable_cache/disable_cache on the root module rather than at the true source (register_hook/remove_hook), which the author calls out. Given a child registry can't reach a stale ancestor cache, this is a reasonable and minimal choice that covers the reported path for every cache technique.
• The HookRegistry import added to enable_cache/disable_cache and the top-level FirstBlockCacheConfig/FluxTransformer2DModel imports in the test are all valid exports.

Tests

• Both the unit-level (test_child_registries_cache_invalidation) and end-to-end (test_cache_context_after_enable_cache_with_prior_context) tests exercise the fix and match the failure described in the issue. They use small CPU-friendly configs consistent with the rest of the file.

Matches the PR description. No blocking issues.

serge v0.1.0 · model: claude-opus-4-8 · 8 LLM turns · 9 tool calls · 36.9s · 145457 in / 2138 out tokens

[SuryanshSS1011]

Gentle bump on this one @sayakpaul. I moved the end-to-end test into FirstBlockCacheTesterMixin as you suggested. Anything else you'd like changed before this can go in?

[sayakpaul]

Are we assuming that the underlying pipeline implementation is following the following pattern (originally shown in #14037):

with torch.no_grad(), model.cache_context("cond"):
    model(**make_inputs())

There is also no presence of another

with torch.no_grad(), model.cache_context("cond"):
    model(**make_inputs())   # ValueError: No context is set

(as reported in #14037)

So, I think this test test is still incomplete.

[SuryanshSS1011]

Yes, you are right. I do realize that the test leans on the pipeline calling cache_context() internally and doesn't include the second context entry from the issue, so it's incomplete as-is.

I'll rework it to explicitly drive the two cache_context("cond") forwards from the issue's reproduction, with the second one being where No context is set was raised, rather than relying on the pipeline to do it implicitly. I'll make sure it fails on main and passes with the fix.

[SuryanshSS1011]

@sayakpaul, I have pushed the rework. The test now drives the two cache_context("cond") forwards explicitly and asserts the pipeline actually called the transformer. Confirmed failing on main and passing with the fix across all the FirstBlockCache pipelines (Flux, LTX, Mochi, CogVideoX, HunyuanVideo, HunyuanImage). Ready for another look when you have a chance.

[SuryanshSS1011]

@sayakpaul gentle bump. The test rework you asked for is pushed (drives the two cache_context("cond") forwards explicitly, fails on main and passes with the fix). Anything else before this can go in?