Add SeaCache support for Cosmos3 pipelines

docs/source/en/api/cache.md

@@ -46,3 +46,9 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate
 [[autodoc]] MagCacheConfig
 
 [[autodoc]] apply_mag_cache
+
+## SeaCacheConfig
+
+[[autodoc]] SeaCacheConfig
+
+[[autodoc]] apply_sea_cache

docs/source/en/api/pipelines/cosmos3.md

@@ -692,6 +692,39 @@ if result.action is not None:
 </hfoption>
 </hfoptions>
 
+## SeaCache
+
+SeaCache is disabled by default. Cosmos 3 supports enabling it explicitly with [`SeaCacheConfig`]. SeaCache reuses
+transformer residuals when the Spectral-Evolution-Aware indicator changes slowly, reducing the number of full
+transformer executions. It computes the indicator from the raw vision latents, including clean conditioning frames for
+image-to-video generation. Enable it on the transformer with scheduler metadata callbacks from the pipeline:
+
+```python
+import torch
+from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda"
+)
+
+pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))
+```
+
+The same model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and
+[`Cosmos3DistilledModularPipeline`]. SeaCache is approximate and can change generated outputs. Disable it with
+`pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer. Cache state is
+reset after each pipeline call, and conditional and unconditional guidance branches keep independent histories.
+
+Cosmos 3 keeps the SeaCache gate outside its repeated decoder layers, so it is compatible with regional compilation.
+Compile the layers after enabling the cache:
+
+```python
+pipe.transformer.compile_repeated_blocks(fullgraph=True)
+```
+
+SeaCache also supports the Cosmos 3 Ulysses context-parallel and DTensor-based tensor-parallel helpers documented
+below. Cache decisions are synchronized across ranks. Full-model compilation and other model-sharding strategies are not claimed.
+
 ## Context parallelism
 
 For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens.

docs/source/en/optimization/cache.md

@@ -68,6 +68,44 @@ config = FasterCacheConfig(
 pipeline.transformer.enable_cache(config)
 ```
 
+## SeaCache
+
+[SeaCache](https://huggingface.co/papers/2602.18993) compares Spectral-Evolution-Aware (SEA) indicators between
+successive denoising steps. When the accumulated indicator change remains below a threshold, it skips the expensive
+transformer block stack and predicts its output from cached residuals. The indicator is computed from the raw vision
+latents, including clean conditioning frames for image-to-video generation.
+
+The implementation provides built-in adapters for the following models:
+
+- **Cosmos 3** is the primary optimized and benchmarked integration. It caches the complete decoder stack through a
+  post-normalization boundary.
+- **Wan T2V** uses the generic repeated-block path in eager mode. This integration demonstrates how another
+  single-stream video transformer can provide raw vision latents to SeaCache; it is not a claim that the same cache
+  parameters are optimal for Wan or that other Wan variants are supported.
+
+Other video transformers can integrate with the generic path when they use `CacheMixin`, expose a recognized repeated
+block list, and register the block input/output layout in `TransformerBlockRegistry`. The pipeline must enter a
+`cache_context` for every transformer call, attach `step_index`, `sigma`, and `num_inference_steps`, and use separate
+context names for independent trajectories such as conditional and unconditional guidance. Pass a `raw_vision_callback`
+that returns the noisy vision latents when no built-in adapter is available. Validate output quality and tune the cache
+parameters for each model and scheduler; support and benchmark results do not transfer automatically from Cosmos 3.
+
+### Cosmos 3
+
+SeaCache is disabled by default. Enable it on the transformer; the Cosmos 3 denoising loop attaches the active
+scheduler step, sigma, and step count to each `cache_context` call, so no extra wiring is needed:
+
+```python
+from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
+
+pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
+pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))
+```
+
+This model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and
+[`Cosmos3DistilledModularPipeline`]. SeaCache is an approximate optimization and may change generated outputs. Call
+`pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer.
+
 ## FirstBlockCache
 
 [FirstBlock Cache](https://huggingface.co/docs/diffusers/main/en/api/cache#diffusers.FirstBlockCacheConfig) checks how much the early layers of the denoiser changes from one timestep to the next. If the change is small, the model skips the expensive later layers and reuses the previous output.

src/diffusers/__init__.py

@@ -204,6 +204,7 @@
             "LayerSkipConfig",
             "MagCacheConfig",
             "PyramidAttentionBroadcastConfig",
+            "SeaCacheConfig",
             "SmoothedEnergyGuidanceConfig",
             "TaylorSeerCacheConfig",
             "TextKVCacheConfig",
@@ -212,6 +213,7 @@
             "apply_layer_skip",
             "apply_mag_cache",
             "apply_pyramid_attention_broadcast",
+            "apply_sea_cache",
             "apply_taylorseer_cache",
             "apply_text_kv_cache",
         ]
@@ -1083,6 +1085,7 @@
             LayerSkipConfig,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
+            SeaCacheConfig,
             SmoothedEnergyGuidanceConfig,
             TaylorSeerCacheConfig,
             TextKVCacheConfig,
@@ -1091,6 +1094,7 @@
             apply_layer_skip,
             apply_mag_cache,
             apply_pyramid_attention_broadcast,
+            apply_sea_cache,
             apply_taylorseer_cache,
             apply_text_kv_cache,
         )

src/diffusers/hooks/__init__.py

@@ -25,6 +25,7 @@
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
     from .mag_cache import MagCacheConfig, apply_mag_cache
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
+    from .sea_cache import SeaCacheConfig, apply_sea_cache
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
     from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache
     from .tensor_parallel import apply_tensor_parallel

src/diffusers/hooks/_helpers.py

@@ -27,6 +27,7 @@ class TransformerBlockMetadata:
     return_hidden_states_index: int = None
     return_encoder_hidden_states_index: int = None
     hidden_states_argument_name: str = "hidden_states"
+    encoder_hidden_states_argument_name: str = "encoder_hidden_states"
 
     _cls: Type = None
     _cached_parameter_indices: dict[str, int] = None
@@ -174,6 +175,7 @@ def _register_transformer_blocks_metadata():
     from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
     from ..models.transformers.transformer_bria import BriaTransformerBlock
     from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
+    from ..models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
     from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
     from ..models.transformers.transformer_hunyuan_video import (
         HunyuanVideoSingleTransformerBlock,
@@ -230,6 +232,17 @@ def _register_transformer_blocks_metadata():
         ),
     )
 
+    # Cosmos 3
+    TransformerBlockRegistry.register(
+        model_class=Cosmos3VLTextMoTDecoderLayer,
+        metadata=TransformerBlockMetadata(
+            return_hidden_states_index=1,
+            return_encoder_hidden_states_index=0,
+            hidden_states_argument_name="gen_seq",
+            encoder_hidden_states_argument_name="und_seq",
+        ),
+    )
+
     # Flux
     TransformerBlockRegistry.register(
         model_class=FluxTransformerBlock,

src/diffusers/hooks/hooks.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import functools
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -31,29 +32,64 @@ def reset(self, *args, **kwargs) -> None:
         )
 
 
+@dataclass(frozen=True)
+class CacheContext:
+    """Information a pipeline attaches to a denoising call for the cache hooks, via `cache_context`.
+
+    `name` identifies the call (usually `"cond"` or `"uncond"`) and keys the per-context hook state; the remaining
+    fields describe where the denoising loop is. Fields default to `None`.
+    """
+
+    name: str
+    step_index: int | None = None
+    num_inference_steps: int | None = None
+    timestep: float | torch.Tensor | None = None
+    sigma: float | None = None
+
+
+def _set_cache_context(module: torch.nn.Module, context: CacheContext | None) -> None:
+    """Set (or clear, with `None`) the cache context on every stateful hook's `StateManager` under `module`."""
+    registry = HookRegistry.check_if_exists_or_initialize(module)
+    for child in (registry, *registry._get_child_registries()):
+        for hook_name in reversed(child._hook_order):
+            hook = child.hooks[hook_name]
+            if not hook._is_stateful:
+                continue
+            for attr in vars(hook).values():
+                if isinstance(attr, StateManager):
+                    attr.set_context(context)
+
+
 class StateManager:
     def __init__(self, state_cls: BaseState, init_args=None, init_kwargs=None):
         self._state_cls = state_cls
         self._init_args = init_args if init_args is not None else ()
         self._init_kwargs = init_kwargs if init_kwargs is not None else {}
         self._state_cache = {}
-        self._current_context = None
+        self._context: CacheContext | None = None
+
+    @property
+    def context(self) -> CacheContext:
+        if self._context is None:
+            raise ValueError(
+                "No cache context is set. Wrap the denoiser call in `model.cache_context(name, ...)` before calling it."
+            )
+        return self._context
 
     def get_state(self):
-        if self._current_context is None:
-            raise ValueError("No context is set. Please set a context before retrieving the state.")
-        if self._current_context not in self._state_cache.keys():
-            self._state_cache[self._current_context] = self._state_cls(*self._init_args, **self._init_kwargs)
-        return self._state_cache[self._current_context]
+        name = self.context.name
+        if name not in self._state_cache.keys():
+            self._state_cache[name] = self._state_cls(*self._init_args, **self._init_kwargs)
+        return self._state_cache[name]
 
-    def set_context(self, name: str) -> None:
-        self._current_context = name
+    def set_context(self, context: CacheContext | None) -> None:
+        self._context = context
 
     def reset(self, *args, **kwargs) -> None:
         for name, state in list(self._state_cache.items()):
             state.reset(*args, **kwargs)
             self._state_cache.pop(name)
-        self._current_context = None
+        self._context = None
 
 
 class ModelHook:
@@ -132,14 +168,6 @@ def reset_state(self, module: torch.nn.Module):
             raise NotImplementedError("This hook is stateful and needs to implement the `reset_state` method.")
         return module
 
-    def _set_context(self, module: torch.nn.Module, name: str) -> None:
-        # Iterate over all attributes of the hook to see if any of them have the type `StateManager`. If so, call `set_context` on them.
-        for attr_name in dir(self):
-            attr = getattr(self, attr_name)
-            if isinstance(attr, StateManager):
-                attr.set_context(name)
-        return module
-
 
 class HookFunctionReference:
     def __init__(self) -> None:
@@ -267,15 +295,6 @@ def check_if_exists_or_initialize(cls, module: torch.nn.Module) -> "HookRegistry
             module._diffusers_hook = cls(module)
         return module._diffusers_hook
 
-    def _set_context(self, name: str | None = None) -> None:
-        for hook_name in reversed(self._hook_order):
-            hook = self.hooks[hook_name]
-            if hook._is_stateful:
-                hook._set_context(self._module_ref, name)
-
-        for registry in self._get_child_registries():
-            registry._set_context(name)
-
     def invalidate_child_registries_cache(self) -> None:
         """Invalidate the cached child-registry list across this module's tree.
 
@@ -294,8 +313,8 @@ def _get_child_registries(self) -> list["HookRegistry"]:
         """Return registries of child modules, using a cached list when available.
 
         The cache is built on first call and reused for subsequent calls. This avoids the cost of walking the full
-        module tree via named_modules() on every _set_context call, which is significant for large models (e.g. ~2.7ms
-        per call on Flux2).
+        module tree via named_modules() on every cache context update, which is significant for large models (e.g.
+        ~2.7ms per call on Flux2).
         """
         if not hasattr(self, "_child_registries_cache"):
             self._child_registries_cache = None

src/diffusers/hooks/mag_cache.py

@@ -21,7 +21,7 @@
 from ..utils.torch_utils import unwrap_module
 from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS
 from ._helpers import TransformerBlockRegistry
-from .hooks import BaseState, HookRegistry, ModelHook, StateManager
+from .hooks import BaseState, CacheContext, HookRegistry, ModelHook, StateManager
 
 
 logger = get_logger(__name__)  # pylint: disable=invalid-name
@@ -183,8 +183,8 @@ def initialize_hook(self, module):
 
     @torch.compiler.disable
     def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if self.state_manager._current_context is None:
-            self.state_manager.set_context("inference")
+        if self.state_manager._context is None:
+            self.state_manager.set_context(CacheContext(name="inference"))
 
         arg_name = self._metadata.hidden_states_argument_name
         hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
@@ -299,8 +299,8 @@ def initialize_hook(self, module):
 
     @torch.compiler.disable
     def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if self.state_manager._current_context is None:
-            self.state_manager.set_context("inference")
+        if self.state_manager._context is None:
+            self.state_manager.set_context(CacheContext(name="inference"))
         state: MagCacheState = self.state_manager.get_state()
 
         if not state.should_compute: