* [hooks] Support Flux2 in MagCache
MagCache could not be enabled on `Flux2Transformer2DModel` (Flux2 and Flux2 Klein):
1. `Flux2TransformerBlock` and `Flux2SingleTransformerBlock` were not registered in
`TransformerBlockRegistry`, so `apply_mag_cache` raised `ValueError: Model class ... not registered`.
2. Once registered, the tail hook raised a shape error. The head hook records the image-only input of
`transformer_blocks[0]`, but the Flux2 model concatenates the text and image streams before the
single-stream blocks and calls them with `encoder_hidden_states=None`, so the last single block returns
the fused `[text; image]` sequence. The residual `tail_output - head_input` could not be formed.
Register the two Flux2 blocks (the single block with `return_encoder_hidden_states_index=None`, matching
how the model calls it) and add `_align_to_hidden_states`, which keeps only the trailing image tokens of
the fused tensor when computing the residual. The cached residual is therefore image-only, which is what
the reference MagCache4FLUX implementation does, and it matches the head input on skipped steps. This also replaces
a dead "fallback" branch in the tail hook that subtracted mismatched shapes anyway.
Tests: `MagCacheTesterMixin` for `Flux2Transformer2DModel`, `Flux2Pipeline` and `Flux2KleinPipeline`.
FirstBlockCache has the same head/tail structure and still fails on Flux2; left for a follow-up.
* [hooks] Slice the fused tail output in the existing mismatch branch
Implement the Flux2 residual fix inside the tail hook's sequence-length mismatch branch instead of a
separate `_align_to_hidden_states` helper. The surrounding structure stays as it was: the `diff == 0`
case still covers a batch-size mismatch via broadcasting, and only the mismatched-length case changes
to subtract the trailing image tokens.
---------
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
* Add optional FP32 sampling state for Cosmos3
* Add cache execution contexts to Cosmos3 pipelines
* Add SeaCache support for Cosmos3
* Cap consecutive SeaCache reuses
* Align SeaCache defaults with Cosmos3 inference
* Enable SeaCache by default for Cosmos3
* Document SeaCache for Cosmos3
* Use FP32 sampling state by default for Cosmos3
* Fix SeaCache reconfiguration, transfer state isolation, and FP32 sampling
* Remove SeaCache stats and ablation instrumentation
* Move SeaCache to transformer-level and disable by default
* Align SeaCache and tests with model-level cache patterns
* Make Cosmos3 sampling state always FP32
* Add pipe definition to the docs snippet
* Remove redundant hook registry cache resets
* Use raw vision latents for SeaCache indicators
* Make Cosmos3 SeaCache safe for regional compilation
* Fix SeaCache with Cosmos3 model parallelism
* Support single-stream models in SeaCache
* Add Wan T2V support to SeaCache
* pass denoising step info to cache hooks via cache_context
Replaces SeaCache's pipeline-state callbacks: the pipeline attaches the loop
coordinates (step_index, sigma, num_inference_steps) to the denoiser call with
`cache_context(name, ...)`, carried by a frozen CacheContext dataclass and stored
on each stateful hook's StateManager (the registry no longer holds context).
Co-Authored-By: Dhruv Nair <dhruv.nair@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Fix Wan cache metadata and Cosmos modular cleanup
Pass denoising step metadata to Wan T2V conditional and unconditional cache contexts.
Restore stateful cache reset after Cosmos modular pipeline calls, restore the required Literal import, and update regression tests and documentation.
* Refactor Cosmos3 SeaCache decoder hook
* Apply CI formatting updates
* Fix distilled system prompt fallback
* Preserve Cosmos3 fp32 sampling state
---------
Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
