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,47 @@ 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,
+ 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.
+
+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,52 @@ 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 and supports eager inference and regional compilation.
+- **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, using separate context names for independent trajectories such as
+conditional and unconditional guidance. Pass a `raw_vision_callback` that returns the noisy vision latents, in addition
+to the scheduler metadata callbacks shown below. 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 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,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,