Repository navigation
StateManager, Hook, State: missing context setting (cond or uncond) in pipelines #12760
Description
In diffusion denoising model, we usually use classifier free guidance to control qualities. With State Manager, we want each context hold different state, so we have to set which context when calling self.transformer, here is a good example
But some popular pipelines (Flux2) or pipeline variants (Kontext, Image to image, Inpainting) don't implement this. For example, Flux T2I has it while Flux Kontext is missing
A Fix should be simple but require us to check all pipelines
with self.transformer.cache_context("cond"):
noise_pred = self.transformer(...)
with self.transformer.cache_context("uncond"):
neg_noise_pred = self.transformer(...)
Activity
toilaluan commentedon Dec 1, 2025
By using a simple script that counts
self.transformer(andself.transformer.cache_contextoccurrences, there are pipelines need to be fixed:Check Script
omkar-334 commentedon Dec 14, 2025
@DN6 @toilaluan can you review my PR
github-actions commentedon Jan 9, 2026
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
github-actions commentedon Feb 8, 2026
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
fusheng-ji commentedon Oct 5, 2026
@toilaluan @DN6 I'd like to pick this up, since #12833 was closed without merging.
On current
main, enabling a stateful cache (TaylorSeer, FirstBlockCache, MagCache) on these pipelines still fails withValueError: No cache context is set:I'd wrap their transformer calls in
cache_context("cond")/cache_context("uncond")the same wayFluxPipelinedoes, and add a small test mixin that checks every transformer call runs inside a named context.I tested this locally on FLUX.1-dev, FLUX.1-Kontext-dev, Chroma1-HD, FLUX.2-dev and Wan 2.1 1.3B. Without a cache, the outputs are bitwise identical to
main. With a cache, they now run instead of raising.I would leave out ControlNet SD3 and OmniGen, because their transformers don't support caching yet. The remaining pipelines from the list above could follow in separate PRs. #14900 also touches Flux Kontext, so I can rebase on it if it lands first.
Does this scope work for you?
sayakpaul commentedon Oct 5, 2026
I think this works for us, but in general, we'd prefer smaller PRs so that it's easier to review.