Caching implementation to speedup video diffusion

[lfengad]

Cherry-pick of imaginaire4 7e54d99e (MR !10482) onto main.

Originally opened against release/2026-08-18-ee5f5a61; rebased onto main (cb76f47, the 2026-08-19 i4 release) after that release branch was superseded. The picked hunks are byte-identical to the previous revision of this PR.

Why this is smaller than the upstream MR

The model-side implementation and the TaylorSeer removal from unified_mot.py already landed on main via the i4 release syncs. Of the 8 files the MR touched, only these still needed picking:

upstream (i4)	here	status
packages/cosmos3/cosmos3/common/args.py	cosmos_framework/inference/common/args.py	picked
packages/cosmos3/cosmos3/inference.py	cosmos_framework/inference/inference.py	picked
packages/cosmos3/cosmos3/scripts/inference.py	cosmos_framework/scripts/inference.py	picked
projects/.../models/utils/taylorseer.py (deleted)	same	picked (+ its NOTICE entry)
projects/.../models/mot/diffusion_cache.py	cosmos_framework/model/generator/mot/diffusion_cache.py	already synced
projects/.../models/mot/unified_mot.py	same	already synced
projects/cosmos3/cosmos3/CONTRIBUTING.md	CONTRIBUTING.md	no-op, never mentioned TaylorSeer here
projects/.../models/mot/diffusion_cache_test.py	—	intentionally not picked

What it does

Exposes SeaCache (the diffusion-time inference cache) through the inference args:

• --diffusion-cache / --no-diffusion-cache
• --diffusion-cache-thresh (default 0.35)
• --diffusion-cache-residual-order (default 1)

Installed from OmniInference._create rather than from scripts/inference.py, so any caller that opts in — CLI, Ray serve, action policy servers, cosmos-benchmarks EvalModel — goes through the same path.

Behavior change

The inference CLI now has the cache on by default. This is a lossy approximation, not just a speedup: when the accumulated relative-L1 SEA indicator stays under the threshold, the language_model evaluation is skipped and the generation residual is extrapolated via Newton divided differences. Generated video will differ from before. --no-diffusion-cache restores the previous behavior exactly.

Everything else is unchanged. SetupOverrides.diffusion_cache defaults to False, so Ray serve, the action policy servers, examples/, and upsample_prompts behave identically unless they opt in — and with the flag off, _maybe_install_diffusion_cache returns before importing the module or patching anything. AR / KV-cache (reasoner) paths bypass the cache automatically even when it is on.

Notes

• The three picked code regions are byte-identical to upstream apart from the one rewritten import path; the install point sits in the same enclosing method (OmniInference._create) with the same surrounding context.
• The upstream comment says all callers "get the default-on cache". In practice only the CLI sets it — that is also true in i4, so the comment is kept verbatim to keep future diffs against i4 clean.
• taylorseer.py has zero importers on main (the unified_mot.py import is already gone); its NOTICE entry is removed alongside it so the file list stays accurate.

Testing

• Rebase onto main applied with no conflicts; git show of the rebased commit is identical to the pre-rebase diff.
• ruff check passes on all touched files, re-run after the rebase.
• On the previous base: inference/args_test.py, inference/common/args_test.py, inference/inference_test.py, inference/common/inference_test.py → 51 passed, 2 failed. Both failures reproduced identically on the unmodified base and are CUDA-driver-too-old errors on the test node, not regressions. Not re-run on the new base.
• Verified manually: CLI default True, library default False, fields propagate through _build's model_dump(), and the override keys are accepted by DiffusionCache.Config.

🤖 Generated with Claude Code