Add SeaCache support for Cosmos3 pipelines

[yzhautouskay]

What does this PR do?

• Enable SeaCache by default with linear extrapolation, threshold 0.25, and at most 2 consecutive cached steps. This provides 2-3x speedup depending on inference params with on-par quality.
• Use FP32 sampling state by default while keeping transformer execution in the model dtype, improving numerical stability and consistency, especially when caching is enabled
• Add tests and documentation, including how to disable SeaCache.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[yzhautouskay]

Claude-Opus 4.8 self-review skill run based on guidelines

Self-review: yzhautouskay/cosmos3_diffusion_caching

The branch bundles three things: (1) a new SeaCache cache hook (hooks/sea_cache.py, 989 lines), (2) a use_fp32_sampling_state sampling-precision option across the Cosmos3 pipelines, and (3) a _get_execution_device group-offload fix. Paper link (2602.18993) verified — real (SeaCache, CVPR 2026).

Blocking (fix before submitting)

1. PR is not laser-focused — three unrelated changes in one PR. Per How to write a good PR + cosine schedule and unet config #2: "A pull request should solve one problem and one problem only." SeaCache caching, FP32 sampling state, and the _get_execution_device group-offload fix (pipeline_cosmos3_omni.py:~530) are independent and each reviewable on its own. The group-offload fix in particular has nothing to do with caching. Split into separate PRs (or at minimum call out why they're coupled in the description).

2. Heavy defensive / fail-open code contradicts the style guide. Per code_style.md: "No defensive code, unused code paths, or legacy stubs — do not add fallback paths, safety checks... 'just in case'." sea_cache.py is dominated by fail-open guards: autograd check (sea_cache.py:730), missing-context check (:735), shape/device/dtype equality checks on indicators and residual history (:459, :949), _is_parameter_sharded FSDP/DTensor probing (:687), non-adjacent-step detection (:435). Some fail-open is legitimately needed for a cache that must never corrupt output, but the current volume will draw review fire. Prune the guards for states that can't actually occur under the supported call path, and keep only the ones protecting real, reachable conditions.

3. Generic "advanced adapter" config params carried for a single-model integration. metadata_callback and raw_vision_callback on SeaCacheConfig (sea_cache.py:315-322) exist only to let future non-Cosmos3 models plug in. Per code_style.md: "do not carry unused method parameters 'for API consistency'... only keep the inference path you are actually integrating." Only Cosmos3 is wired up; drop these until a second model needs them.

Non-blocking (raise with reviewer)

• indicator_source="first_block" is a whole second indicator path that the default never uses (default is "raw_vision_latents"). It pulls in _prepare_cosmos3_vision_metadata (~60 lines), hidden_states_norm_module_name plumbing in _helpers.py, and the norm-based branch in _build_indicator. If it isn't needed for the shipping config, consider dropping it. (See dead-code table.)
• use_fp32_sampling_state=True by default changes generated outputs vs. prior behavior. Documented in cosmos3.md, but flag it explicitly to the maintainer as an intentional default change (ties to pitfalls Add UNet for Latent Diffusion #5/Unet for Grad TTS and pipeline #9/Pros and cons of the configuration setup #11 on precision — this is the correct direction, just call it out).
• _cache_context helper is duplicated verbatim in denoise.py and pipeline_cosmos3_omni.py. Minor; consider a shared location.
• Paper claims "no additional hyperparameters to tune", yet the config exposes threshold, retention_steps, cache_end_steps, max_consecutive_cached, power_exp, residual_order. Worth a one-line rationale in the PR for why the port diverges.

Dead code (advisory)

location	status	reason
sea_cache.py:555 _prepare_cosmos3_vision_metadata + first_block branch in _build_indicator	Likely-dead under default	only runs when indicator_source="first_block"; default is raw_vision_latents
SeaCacheConfig.metadata_callback / raw_vision_callback	Likely-dead	no caller passes them; Cosmos3 uses native adapters
_helpers.py hidden_states_norm_module_name	Used only by first_block	tied to the non-default indicator

Process reminders (AI-contribution guidelines)

• Uncommitted/untracked files in your worktree are NOT in the pushed diff but shouldn't get added: check_cosmos3_transfer_workflow.py, run_cosmos3_modular_transfer.sh, assets/, examples/cosmos3/inference_cosmos3_modular_distilled.py, and an unstaged encoders.py edit. The dev scripts/assets are ephemeral (per review-rules "Ephemeral context") — keep them out of the PR.
• Add a coordination link (issue where a maintainer acknowledged this work) to the PR description.
• Add test commands + pasted results to the PR. Suggested: pytest tests/hooks/test_sea_cache.py tests/models/transformers/test_models_transformer_cosmos3.py tests/pipelines/cosmos/test_cosmos3.py tests/modular_pipelines/cosmos/ -q.

Verdict: NEEDS CHANGES

Primary asks: split the three features (#1), trim defensive/unused paths (#2, #3). Tests are strong — real Cosmos3 model at tiny config, per-context cache isolation, fail-open, and dtype coverage all exercised, consistent with testing.md

[sayakpaul]

Thanks! The design looks mostly good. My major comment is that we're adding utilities to enable caching at the pipeline-level which is not something we typically do in the library.

[github-actions]

Hi @yzhautouskay, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[yzhautouskay]

@yiyixuxu @sayakpaul Thanks for the review! I addressed moving caching to the transformer level instead of pipeline, and refactored tests

Few thing I will address next:

• Verify always upcast sampling state to fp32 impact once again before making the final change: Benchmarked FP32 sampling across T2I, T2V, and I2V: runtime overhead was 0.3–2.7%, with only ~111 MiB extra peak live memory for video. FP32 was perceptually better and closer to the reference PyTorch implementation, so I removed the flag and made FP32 sampling unconditional
• Add the benchmarking results for speed and quality with caching to update the docs.
• Document if SeaCache is Cosmos3 specific method or if it can be extended to other models

[sayakpaul]

Thanks for the updates! This looks much better!

[sayakpaul]

Would it make sense to provide actual values here? Or maybe even just specify what pipe is supposed to be?

[yzhautouskay]

These must remain callbacks because the step index and sigma change during every denoising iteration, while num_timesteps is determined per pipeline call. I updated the example to define pipe explicitly

[sayakpaul]

Nice, thanks for the note! From a quick skim of the paper, it doesn't look like it needs to be Cosmos3 specific no?

[sayakpaul]

What is happening here?

[yzhautouskay]

This fix is only needed for the specific case where a baseline inference runs first without caching enabled, while still entering cache_context(), and caching is enabled on the same transformer instance afterward
Before the fix: Cosmos3OmniPipeline.__call__() → CacheMixin.cache_context() → HookRegistry._set_context() → _get_child_registries() scans the model and saves [], because no cache hooks exist yet. After enable_cache() installs child hooks, the next context call reuses that old empty list, so "cond"/"uncond" does not reach them.
With the fix:* enable_cache() sets _child_registries_cache = None after installing hooks. The next context call rescans the model, finds the new child registries, and propagates the context correctly

[sayakpaul]

/diffusers-bot pytest tests/models -k "seacache"

[github-actions]

✅ pytest tests/models -k "seacache" passed on GPU — view logs.