[feat]: Cosmos3 port (Tier-A scaffold; SSIM + Tier-B deferred pending weight publish) - #1382
[feat]: Cosmos3 port (Tier-A scaffold; SSIM + Tier-B deferred pending weight publish)#1382SolitaryThinker wants to merge 5 commits into
Conversation
mergify
Bot
added
type: feat
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the architectural scaffold for the Cosmos3 model and pipeline, adding configuration classes, a skeleton for the "Cosmos3VFMTransformer", and the "Cosmos3OmniDiffusersPipeline" logic. It also includes a checkpoint conversion script and several local parity tests. Feedback points out that the transformer's "init" signature does not match the "BaseDiT" contract, and both the transformer and pipeline are missing essential attributes and methods, such as "reset_cache", "cached_kv", and "progress_bar", which are necessary for the diffusion loop to execute.
| def __init__( | ||
| self, | ||
| od_config: object | None = None, | ||
| *, | ||
| temporal_compression_factor: int | None = None, | ||
| ) -> None: |
There was a problem hiding this comment.
The __init__ signature of Cosmos3VFMTransformer does not match the BaseDiT interface contract. BaseDiT (defined in fastvideo/models/dits/base.py) expects (self, config: DiTConfig, hf_config: dict[str, Any], **kwargs). The current implementation using od_config will break generic model loading and initialization logic within the FastVideo framework, such as the TransformerLoader.
| raise NotImplementedError("Phase 2b.2") | ||
|
|
||
|
|
||
| class Cosmos3VFMTransformer(BaseDiT): |
There was a problem hiding this comment.
The Cosmos3VFMTransformer class is missing the reset_cache() method and the cached_kv / cached_freqs_gen attributes. These are explicitly accessed by the Cosmos3OmniDiffusersPipeline.diffuse method (e.g., lines 373 and 401). Without these, any attempt to run the ported diffusion loop will raise an AttributeError. Even as a scaffold, the model should expose the interface required by its corresponding pipeline.
|
|
||
| # -- Denoising loop ----------------------------------------------------- | ||
|
|
||
| def diffuse( |
There was a problem hiding this comment.
The diffuse and forward methods reference self.progress_bar and self.device, neither of which are defined or initialized in this class. Since __init__ is currently a stub that raises NotImplementedError, these core dependencies are never established, rendering the ported logic unusable outside of the specific test environment where they are monkey-patched.
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
Adds 8 CPU-only parity test files + conftest stubs (StubScheduler, StubCosmos3VAE, StubCosmos3Transformer) under tests/local_tests/cosmos3/ mirroring the vllm-omni reference at tests/diffusion/models/cosmos3/conftest.py. All 15 tests skip until the FastVideo modules land (subsequent commits). Reference: vllm-omni PR #3454 @ 8536f5b1421f.
…Model)
Ports math (compute_mrope_position_ids_text/_vision, patchify, unpatchify) verbatim from vllm-omni transformer_cosmos3.py:113-177,1009-1036. Module tree (language_model.layers.*.self_attn.{q,k,v,o}_proj + gen_layers.*.cross_attention.* + vae2llm + llm2vae + time_embedder + norm_moe_gen) exposes the parameter names the checkpoint converter targets. All layer forward() raise NotImplementedError until weights publish; module instantiation + state-dict key tests pass.
Pipeline class with: _remap_ckpt_key static method (verbatim port from pipeline_cosmos3.py:319-409, 14-rule checkpoint remap UND/GEN split + lm_head skip); _set_flow_shift method with lazy UniPCMultistepScheduler construction; diffuse() with sequential 3-mode CFG denoising loop + I2V velocity_mask + image_latent re-injection (ported from pipeline_cosmos3.py:883-1033); forward() with T2I/T2V/I2V mode dispatch + flow_shift selection (ported from pipeline_cosmos3.py:1037-1206); 4 helper-method stubs raise NotImplementedError. Dual inheritance (nn.Module, ComposedPipelineBase) matches upstream pattern.
Cosmos3Config(PipelineConfig): reuses Cosmos25VAEConfig (Cosmos3 uses DistributedAutoencoderKLWan); text_encoder_configs=() since Cosmos3LanguageModel lives inside the DiT; flow_shift=1.0 default (T2V/I2V engine init; T2I uses 3.0 per-request via _set_flow_shift). Registry entry registers BEFORE the generic cosmos detector to win path precedence (same pattern as GEN3C). Resolves nvidia/Cosmos3-Nano -> Cosmos3Config without regressing Cosmos25/GEN3C/Cosmos.
scripts/checkpoint_conversion/cosmos3_convert.py wraps Cosmos3OmniDiffusersPipeline._remap_ckpt_key. convert_state_dict(src) returns (new_state, skipped, unmapped) for shard-by-shard remap; convert_checkpoint_dir reads *.safetensors shards and writes a consolidated FastVideo-format state dict (sharded index.json TODO when real weights land). smoke_test() exercises all 14 representative remap branches synthetically (no real weights needed) and returns nonzero exit on failure; runs in CI as a remap-drift guard. Usage: python scripts/checkpoint_conversion/cosmos3_convert.py --smoke-test.
ec8d74f to
96005e9
Compare
|
Rebased onto current |
PR #1382 review — Cosmos3 Tier-A scaffoldReviewed exact head: Verdict: changes required. The PR is an honest scaffold, but it registers that scaffold as a supported production model and ships a converter that no longer matches NVIDIA's released checkpoint. [P0] The production registry exposes a pipeline/config/DiT that cannot initialize or run
Reproductions: This is guaranteed runtime failure, not deferred quality work. Complete the FastVideo-native loader and stage-composed pipeline contracts, or keep the scaffold out of [P0] The converter maps only 4/814 released transformer tensors and writes the broken resultNVIDIA has now published The PR's remap expects the old pre-release The documented root snapshot invocation also finds no shards because the release stores all seven transformer shards under Redo conversion from pinned official NVIDIA weights plus official/FastVideo key-and-shape inventories. Conversion must fail closed on unexpected/missing keys and prove the emitted directory strict-loads through Official evidence:
[P1] The green suite is scaffold self-consistency, not component or pipeline parity
Exact-head result: There is no real-weight strict load, official numerical component parity, production-loader smoke, pipeline parity, runnable example, SSIM test/reference, or [P1] Released modalities and public workload scope are unresolvedThe official transformer enables both Either support every base output head and advertised workload or record an explicit accepted scope cut and make the public registry/API match it. Verification
|
1 participant
What
Tier-A architectural port of NVIDIA Cosmos3 (
Cosmos3OmniDiffusersPipeline) from vllm-project/vllm-omni#3454 @8536f5b1421finto FastVideo. T2V/I2V/T2I modalities; sound/action modes deliberately out-of-scope per the upstream PR's own scope carve-out.This PR ships scaffolding only.
nvidia/Cosmos3-Nanoweights are not yet publicly accessible on HF Hub (401 anonymous/404 authenticatedacross all plausible name variants undernvidia/andnvidia-cosmos/; vllm-omni maintainer @ywang96 confirmed publication is coordinated with #3454's merge timing). Tier-B work (real-weight forward parity, full inference SSIM gate, layer implementations) is deferred until weights publish.Why Tier-A only
The vllm-omni reference itself ships its tests using stub fixtures + tiny synthetic configs (
tests/diffusion/models/cosmos3/conftest.py:StubScheduler/StubCosmos3VAE/StubCosmos3Transformer+_tiny_cosmos3_config(hidden_size=8, num_hidden_layers=0)). The same approach works here — architectural correctness IS testable without weights. We ship that test coverage now; SSIM and per-layer forward parity activate the moment weights become accessible (no further code changes needed beyond filling in theNotImplementedErrorlayer bodies).Test status
The 2 skipped tests are tokenizer-shape contracts that require a real
nvidia/Cosmos3-Nanocheckpoint with itstext_tokenizer/subfolder. They auto-activate when the pipeline's__init__is wired to load the real tokenizer (post-weight-publish).test_cosmos3_mrope_paritytransformer_cosmos3.py:113-177test_cosmos3_patchify_unpatchify_parity(×2)[B,C,T,H,W] ↔ [B, T·hp·wp, p²·C]roundtrip + defaultpatch_size=[1,2,2]shape contracttest_cosmos3_state_dict_keys::test_fastvideo_cosmos3_remap_matches_reference_remap_ckpt_keymatch the canonical vllm-omni tabletest_cosmos3_state_dict_keys::test_fastvideo_cosmos3_dit_module_tree_param_namestest_cosmos3_scheduler_default_parity(×3)flow_shift=3.0, T2V engine-init shift, scheduler timestep determinismtest_cosmos3_pipeline_call_graph(×5)diffuse()CFG call order [cond, uncond, cond], I2V velocity_mask + image_latent re-injection, T2I/T2V mode dispatch defaults, T2I+video modality rejectiontest_cosmos3_tokenizer_chat_template(×2)Layout
Stats
feat/cosmos3-tier-a-port, each one a phase boundaryValidation
pytest tests/local_tests/cosmos3/ -v→ 13 passed, 2 skippedpython -m py_compileon all new Python files: cleanpython scripts/checkpoint_conversion/cosmos3_convert.py --smoke-test→smoke_test PASSED: 13 new keys, 1 skipped, 0 unmappedpython -c "from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config; c = Cosmos3Config()"→ instantiates cleanly withhidden_size=4096,vae=Cosmos25VAEConfig,flow_shift=1.0nvidia/Cosmos3-Nano → Cosmos3Config; no regression onCosmos25/GEN3C/CosmosforKyleShao/Cosmos-Predict2.5-2B-Diffusers,FastVideo/GEN3C-Cosmos-7B-Diffusers,nvidia/Cosmos-Predict2-2B-Video2Worldpre-commit run --files <changed-paths>on all 16 files: yapf + ruff + codespell + mypy all passTier-B follow-up (when weights publish)
Cosmos3CausalAttention,Cosmos3CrossAttention,Qwen3VLTextRotaryEmbedding,Cosmos3GatedMLP. Each currently has the module tree +q/k/v/o_projlinears at the right GQA shapes;forward()raisesNotImplementedError._format_and_tokenize_prompts,_prepare_latents,_set_scheduler_timesteps,_decode_latents. CurrentlyNotImplementedError; monkey-patched in tests.nvidia/Cosmos3-Nano/text_tokenizer/inCosmos3OmniDiffusersPipeline.__init__; auto-activates the 2 SKIP tests.fastvideo/tests/ssim/test_cosmos3_similarity.pyonce a reference video can be generated against real weights.cosmos3_convert.pywrites a single consolidatedmodel.safetensors; shardedindex.jsonis a TODO marked in the code.What this PR does NOT claim
Refs
8536f5b1421f