Add SeaCache support for Cosmos3 pipelines by yzhautouskay · Pull Request #14663 · huggingface/diffusers

yzhautouskay · 2026-08-31T14:15:12Z

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

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.

…ling

yzhautouskay · 2026-08-31T20:37:55Z

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)

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).
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.
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.

sayakpaul · 2026-09-01T02:48:40Z

+[`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and [`Cosmos3DistilledModularPipeline`] enable SeaCache
+automatically when inference starts. The default [`SeaCacheConfig`] filters raw vision latents, linearly extrapolates
+the cached residual, uses `threshold=0.25`, and allows at most two consecutive cached steps before forcing a full
+transformer execution.


It is not clear if SeaCache is specific to the Cosmos3 family or if it can generally be applied to other pipelines as well. If the latter is true, then let's make it clearer.

sayakpaul · 2026-09-01T02:52:28Z


 class TestCosmos3OmniTransformerModel(Cosmos3OmniTransformerTesterConfig, ModelTesterMixin):
+    @pytest.mark.parametrize("indicator_source", ["first_block", "raw_vision_latents"])
+    def test_cosmos3_supports_sea_cache_without_changing_state_dict_keys(self, indicator_source):


We should add the caching tests similar to:

diffusers/tests/models/transformers/test_models_transformer_flux.py

Line 602 in a262497

class TestFluxTransformerFasterCache(FluxTransformerTesterConfig, FasterCacheTesterMixin):

github-actions · 2026-09-01T07:38:23Z

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.

yiyixuxu · 2026-09-02T00:58:03Z

        add_resolution_template: bool = True,
        add_duration_template: bool = True,
        enable_safety_check: bool = True,
+        use_fp32_sampling_state: bool = True,


ohhh can we just drop this flag and always upcost to fp32 if it improve the accuracy, i mean how much additional memory it would cost? should not be much , no?

yzhautouskay · 2026-09-02T16:59:54Z

@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
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

yzhautouskay added 8 commits August 31, 2026 03:17

Add optional FP32 sampling state for Cosmos3

981fe04

Add cache execution contexts to Cosmos3 pipelines

1dac4d4

Add SeaCache support for Cosmos3

613e441

Cap consecutive SeaCache reuses

12f66f6

Align SeaCache defaults with Cosmos3 inference

54c4cc5

Enable SeaCache by default for Cosmos3

058c0e9

Document SeaCache for Cosmos3

9985bf7

Use FP32 sampling state by default for Cosmos3

e185225

github-actions Bot added size/L PR with diff > 200 LOC documentation Improvements or additions to documentation models tests modular-pipelines utils pipelines hooks and removed size/L PR with diff > 200 LOC labels Aug 31, 2026

Fix SeaCache reconfiguration, transfer state isolation, and FP32 samp…

2cf229c

…ling

github-actions Bot added the size/L PR with diff > 200 LOC label Aug 31, 2026

Remove SeaCache stats and ablation instrumentation

3e68be9

sayakpaul reviewed Sep 1, 2026

View reviewed changes

yiyixuxu added the no-issue-needed for PRs that do not require link to an issue label Sep 2, 2026

yiyixuxu reviewed Sep 2, 2026

View reviewed changes

yzhautouskay added 2 commits September 2, 2026 07:30

Move SeaCache to transformer-level and disable by default

21c75ca

Align SeaCache and tests with model-level cache patterns

0d3c8d6

Navigation Menu

Uh oh!

Add SeaCache support for Cosmos3 pipelines - #14663

Add SeaCache support for Cosmos3 pipelines#14663
yzhautouskay wants to merge 12 commits into
huggingface:mainfrom
yzhautouskay:yzhautouskay/cosmos3_diffusion_caching

yzhautouskay commented Aug 31, 2026 •

edited by yiyixuxu

Loading

Uh oh!

yzhautouskay commented Aug 31, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Sep 1, 2026

Uh oh!

Uh oh!

sayakpaul Sep 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiyixuxu Sep 2, 2026

Uh oh!

yzhautouskay commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

yzhautouskay commented Aug 31, 2026 • edited by yiyixuxu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

yzhautouskay commented Aug 31, 2026

Self-review: yzhautouskay/cosmos3_diffusion_caching

Blocking (fix before submitting)

Non-blocking (raise with reviewer)

Dead code (advisory)

Process reminders (AI-contribution guidelines)

Verdict: NEEDS CHANGES

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiyixuxu Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

yzhautouskay commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yzhautouskay commented Aug 31, 2026 •

edited by yiyixuxu

Loading

Self-review: `yzhautouskay/cosmos3_diffusion_caching`