feat: support Cosmos3 (Omni) via PatchFunctor#1088
Merged
Conversation
Cosmos3OmniPipeline (nvidia/Cosmos3-Nano, Cosmos3-Super) uses
Cosmos3OmniTransformer, a MoT dual-stream DiT whose decoder layers
take/return (und_seq, gen_seq) instead of the usual
(hidden_states, encoder_hidden_states) contract, so the existing
'Cosmos' adapter (CosmosTransformer3DModel) rejects it.
- Add Cosmos3OmniPatchFunctor: reorders the block contract to
(hidden_states=gen_seq, encoder_hidden_states=und_seq) =
ForwardPattern.Pattern_0, so the DBCache residual-diff decision
tracks the DENOISED gen stream. The und/text stream is quasi-static
across denoising steps and would otherwise turn every step into a
false cache hit.
- Extend the Cosmos adapter to detect Cosmos3OmniTransformer and route
it through the functor (has_separate_cfg=True: the pipeline runs
cond/uncond as separate transformer calls).
- README: add Cosmos3 row.
Measured on nvidia/Cosmos3-Nano, RTX PRO 6000 Blackwell (sm_120),
torch 2.12.1+cu132, diffusers main (0.39.0.dev0), 189 frames/35 steps,
warmup excluded, median of 3:
480p: DBCache(rdt=0.20) 76.75s vs 138.59s baseline (1.81x);
DBCache + torch.compile 67.86s (2.04x), 18/35 steps cached,
residual diffs P50 0.148.
256p: 7.75s vs 11.87s baseline (1.53x).
DefTruth
requested changes
Jul 10, 2026
DefTruth
left a comment
DefTruth
left a comment
Member
There was a problem hiding this comment.
Thanks for your contribution!
| | Video | EasyAnimate | [`EasyAnimateTransformer3DModel`](src/cache_dit/caching/block_adapters/adapters.py) | 3 | ✔️/✖️/✔️/✔️ | | ||
| | Video | ConsisID | [`ConsisIDTransformer3DModel`](src/cache_dit/caching/block_adapters/adapters.py) | 1 | ✔️/✔️/✔️/✔️ | | ||
| | Video | Cosmos | [`CosmosTransformer3DModel`](src/cache_dit/caching/block_adapters/adapters.py) | 7+ | ✔️/✖️/✔️/✔️ | | ||
| | Video | Cosmos3 (Omni: T2V/T2I/I2V/V2V) | [`Cosmos3OmniTransformer`](src/cache_dit/caching/patch_functors/functor_cosmos3.py) | 2 | ✔️/✖️/✔️/✔️ | |
Member
There was a problem hiding this comment.
Prefer Cosmos3 (Omni) not Cosmos3 (Omni: T2V/T2I/I2V/V2V) to keep the docs simple.
Member
|
@johnnynunez It would better if you can also add an example into cache_dit.generate CLI (The cache-dit-model-integration skill can help you do this), after that, we can test it via: # For example:
python3 -m cache_dit.generate cosmos3_omni_t2i --cache --summarycache-dit/src/cache_dit/_utils/examples.py Lines 23 to 100 in 02e7acd |
Member
|
The pre-commit check failed, it would be better to refer to our docs PRE_COMMIT and make a fix. |
… examples
- functor_cosmos3.py: docstring re-indented by docformatter (the pre-commit
failure on the PR run).
- README: 'Cosmos3 (Omni)' per review.
- Add cosmos3_omni_t2i / cosmos3_omni_t2v to cache_dit.generate:
ExampleRegister entries, COSMOS3_OMNI_DIR env mapping, __all__ +
_utils/__init__ registration. Output field is 'video' (not 'frames'),
aliased via post-load call wrapper; guardrail disabled by default in the
example loader (optional cosmos_guardrail dep) with license note.
Verified on RTX PRO 6000 (sm_120):
t2i base 2.92s -> cache 2.39s, PSNR 30.25, SSIM 0.764 (rdt=0.12)
rdt=0.08: PSNR 31.22, SSIM 0.844
t2v (93f) base 53.77s -> cache 45.28s
Contributor
Author
|
Thanks for the review @DefTruth! Addressed in f2aff2e:
Verified on RTX PRO 6000 Blackwell: python3 -m cache_dit.generate cosmos3_omni_t2i --cache --summary
# base 2.92s -> cache 2.39s | PSNR 30.25, SSIM 0.764 (default rdt=0.12)
# --residual-diff-threshold 0.08 -> PSNR 31.22, SSIM 0.844
python3 -m cache_dit.generate cosmos3_omni_t2v --cache --summary
# base 53.77s -> cache 45.28s (480x832, 93 frames) |
Member
|
confirm works! thanks for your contribution! |
DefTruth
approved these changes
Jul 10, 2026
DefTruth
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds Cosmos3 (Omni) support —
nvidia/Cosmos3-Nano/nvidia/Cosmos3-SuperviaCosmos3OmniPipeline(diffusersmain).The existing
Cosmosadapter only acceptsCosmosTransformer3DModel(Cosmos 1/2).Cosmos3OmniPipelineusesCosmos3OmniTransformer, a MoT dual-stream DiT whose decoder layers take/return(und_seq, gen_seq)instead of the usual(hidden_states, encoder_hidden_states)contract — soenable_cachefails with:Changes
Cosmos3OmniPatchFunctor(new): patches the decoder layers + transformer forward to reorder the block contract to(hidden_states=gen_seq, encoder_hidden_states=und_seq)=ForwardPattern.Pattern_0. This is load-bearing, not cosmetic: theund(text/understanding) stream is quasi-static across denoising steps, so if DBCache computed its residual-diff decision on it, every step would be a false cache hit and generation would break. With the gen stream first, the decision tracks the actually-denoised tokens.cosmos_adapter: detectsCosmos3OmniTransformer(guardedtry/except ImportErrorfor older diffusers) and routes through the functor withhas_separate_cfg=True(the pipeline runs cond/uncond as separate transformer calls).Measured (real hardware)
nvidia/Cosmos3-NanoT2V, RTX PRO 6000 Blackwell (sm_120, 96 GB), torch 2.12.1+cu132, diffusers 0.39.0.dev0, 189 frames / 35 steps, UniPC flow_shift=10, warmup excluded, median of 3, fixed seed:rdt=0.20rdt=0.20+torch.compilerdt=0.20Residual diffs P50 0.148 (
rdt=0.20), stable across CFG/non-CFG contexts. Output latents finite/sane (no NaN).cache_dit.summary()works and reports per-context stats (cond + CFG separately).Notes / known limitation
main(Cosmos3OmniPipeline/Cosmos3OmniTransformerare not in the 0.38 release); the functor import is_safe_import-guarded so older diffusers are unaffected.torch.compilecurrently trips a Dynamo resume-path bug (AttributeError: 'NVFP4Tensor' object has no attribute 'per_tensor_scale', torchao 0.17 + torch 2.12). DBCache+compile (bf16) and DBCache+NVFP4 (eager) both work. I'll file a separate issue with details — relevant to RFC [RFC] NVFP4 (non SVDQuant) Quantization Support #1061.