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

@@ -660,6 +660,36 @@ 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. 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,
+        current_step_callback=lambda: pipe.current_step_index,
+        current_sigma_callback=lambda: pipe.current_sigma,
+        num_inference_steps_callback=lambda: pipe.num_timesteps,
+    )
+)
+```
+
+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.
+
 ## 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,34 @@ 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.
+
+SeaCache is disabled by default. Enable it on the transformer and provide callbacks for the active scheduler step,
+sigma, and number of inference steps:
+
+```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,
+        current_step_callback=lambda: pipe.current_step_index,
+        current_sigma_callback=lambda: pipe.current_sigma,
+        num_inference_steps_callback=lambda: pipe.num_timesteps,
+    )
+)
+```
+
+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,8 @@ 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"
+    hidden_states_norm_module_name: str = None
 
     _cls: Type = None
     _cached_parameter_indices: dict[str, int] = None
@@ -174,6 +176,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 +233,18 @@ 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",
+            hidden_states_norm_module_name="input_layernorm_moe_gen",
+        ),
+    )
+
     # Flux
     TransformerBlockRegistry.register(
         model_class=FluxTransformerBlock,