Add the reasoning (VLM text-generation) capability. The framework reasoner uses
only the und (causal) pathway weights (no _moe_gen) + embed_tokens / norm /
lm_head — all already present in the native DiT — so a text-only forward + lm_head
reproduces it with no new model code.
- `cosmos3_generate_reasoner_text(transformer, input_ids, max_new_tokens)`:
greedy text reasoning via the und causal backbone + lm_head (text-only prefill,
re-prefill per step; KV-cache fast path is a later optimization). Mirrors the
framework `generate_reasoner_text` (text-only path).
- Parity (`test_cosmos3_reasoning_parity.py`, framework = oracle):
* prefill logits (native text-only forward + lm_head vs framework
`reasoner_forward` + lm_head, per-position): max abs diff = 0.0, mean = 0.0;
* greedy generation vs framework `generate_reasoner_text(do_sample=False)`:
token-for-token identical across seeds.
Text reasoning is fully supported + bit-exact. Image-conditioned reasoning
additionally needs the Qwen3-VL `vision_encoder` (a transformers ViT, tracked
separately). Full suite: 146 passed, 0 skipped.
Collapse the byte-identical _split_action_latent into the generic
_split_flat_latent; all three modality specs now share one splitter.
Behavior-preserving; 139 tests green.
Add the action (multi-embodiment world-model) modality. The Nano checkpoint
ships action weights (action_gen=true, 32 embodiment domains); this activates
them, bit-exact vs the framework.
- DiT forward: action encode/decode mirroring sound but domain-aware — pack each
[T, action_dim] latent with a per-token embodiment domain id, `action_proj_in`
(= framework `action2llm`, a `DomainAwareLinear`) + `action_modality_embed` +
timestep scatter; decode via `action_proj_out` (= `llm2action`) on noisy hidden
states -> `_unpack_action`. Mirrors `_encode_action`/`_decode_action`.
- Sound/action share the vision "full" split; sequence/flat order is
[vision | action | sound] (matches the framework per-sample concat).
- Packing: `Cosmos3ActionItem` + action fields. Action tokens are `(T,)`-shaped,
domain-tagged, with 3D-MRoPE at the vision temporal offset and
`start_frame_offset=1` (framework `_pack_action_tokens`).
- CFG velocity / `Cosmos3DenoiseEngine` take optional `action_specs`
(`Cosmos3ActionSpec`); the joint denoise machinery now handles action.
Parity (`test_cosmos3_action_parity.py`, framework model+pack = oracle, 3
embodiment domains):
- action packing fields + `position_ids` exact;
- `preds_vision` + `preds_action` (forward): max abs diff = 0.0, mean = 0.0;
- combined [vision|action] CFG velocity: max abs diff = 0.0, mean = 0.0.
Action is processed bit-exactly at the model + packing + denoise-glue level. A
full action2world real-weights run is gated on real robot-action input data (not
available here); the bit-exact parity is the correctness proof. Full suite: 139
passed, 0 skipped.
Wire text-to-video+sound end-to-end, completing the audio PR.
- CFG velocity (`cosmos3_get_cfg_velocity`) + `Cosmos3DenoiseEngine` take an
optional `sound_specs`: the denoise latent becomes a combined [vision | sound]
flat vector; each step packs vision+sound, forwards the DiT, zeros the
prediction on conditioning frames for both, and returns the concatenated
velocity. The UniPC scheduler steps the joint vector. `Cosmos3SoundSpec` +
`_split_sound_latent` added.
- `Cosmos3DenoisingStage` t2vs path (gated on `COSMOS3_T2VS`): size the sound
latent from the video duration (framework `create_placeholder_audio` +
`get_latent_num_samples`), append sound noise, joint-denoise, split, AVAE-decode
the sound latent, and set `batch.extra["audio"]` / `["audio_sample_rate"]` so
the generator muxes a stereo 48 kHz AAC track into the mp4. Sound AVAE is
lazy-loaded from `<model_path>/sound_tokenizer` (`_get_sound_vae`).
- Parity (`test_cosmos3_sound_parity::test_t2vs_cfg_velocity_matches_framework`):
the combined [vision|sound] sequential-CFG velocity vs a framework-DiT oracle —
max abs diff = 0.0, mean abs diff = 0.0 across t2vs/i2vs cases.
- Example `basic_cosmos3_t2vs_new_api.py`.
Verified on B200 (real weights, 1280x704, 29f, 35 steps): coherent ocean/sunset
video WITH a real stereo 48 kHz audio track (mean -10.2 dB, peak -0.4 dB; not
silence). Full cosmos3 suite: 130 passed, 0 skipped.
Activate the dormant sound MoT heads and pack the sound modality, completing
the t2vs DiT path (audio component 2).
DiT forward (`Cosmos3VFMTransformer`):
- Encode sound: pack each [C,T] latent -> [T,C], `audio_proj_in` (= framework
`sound2llm`) + `audio_modality_embed`, scatter timestep embeds onto noisy
frames, scatter into the (full) gen split. Mirrors framework `_encode_sound`.
- Decode sound: `audio_proj_out` (= `llm2sound`) on noisy sound hidden states ->
`_unpack_sound` to per-sample [C,T] (framework `_decode_sound`).
- Reuses the general `_scatter_timestep_embeds` (works for (T,1,1) shapes).
Sound packing (`sequence_packing.py`): `Cosmos3SoundItem` + sound fields on
`Cosmos3SampleInputs`/`Cosmos3PackedSequence`/`to_dit_kwargs`. Sound shares the
vision "full" split (preserving the causal+full 2-split invariant), with
`(T,1,1)` token shapes, a `(T,1)` condition mask, mse-loss/timestep bookkeeping
per noisy frame, and 3D-MRoPE temporal positions starting at the vision temporal
offset (parallel to vision; `start_frame_offset=0`, tcf=1; does not advance the
offset). Mirrors framework `_pack_sound_tokens`.
Parity (`test_cosmos3_sound_parity.py`, framework model+pack = oracle):
- sound packing fields + `position_ids` exact vs framework `pack_input_sequence`
(has_sound) across t2vs/i2vs cases.
- `preds_vision` AND `preds_sound`: max abs diff = 0.0, mean abs diff = 0.0
(both native-pack and framework-pack inputs).
Tiny framework builder gains a `sound_gen` flag. Full suite: 127 passed.
Behavior-preserving cleanup (95 tests green):
- reuse fastvideo.layers.visual_embedding.timestep_embedding in the DiT
- single-home COSMOS3_VIDEO_NEGATIVE_PROMPT in the lightweight presets
module; drop the unused duplicate from cosmos3_pipeline
- drop the write-only Cosmos3VisionSpec.clean_latent field + its kwarg
- hoist device = next(transformer.parameters()).device out of the per-pass
_run closure in cosmos3_get_cfg_velocity
The first real-weights Cosmos3-Nano T2V run produced an all-black video.
Root cause: the checkpoint's diffusers-style scheduler_config.json sets
use_karras_sigmas=true, and FastVideo's vendored UniPC checks karras
*before* use_flow_sigmas, so it built diffusion (beta) sigmas instead of
flow-matching sigmas -> scheduler.step diverged to NaN latents -> black
frames. The DiT/CFG velocity itself was clean.
The framework samples with FlowUniPCMultistepScheduler (pure flow:
shift + num_train_timesteps only). FastVideo's vendored UniPC is the same
algorithm, so the fix is config, not a port:
- Coerce the loaded scheduler to the flow setup in
Cosmos3OmniDiffusersPipeline.initialize_pipeline (kill karras/exp/beta
sigmas, force use_flow_sigmas + flow_prediction + final_sigmas_type=zero)
and rebuild the runtime scheduler from it.
- Use FastVideo's native UniPC (not diffusers) in the pipeline and tests,
honoring the "no diffusers at runtime" constraint.
Parity:
- Add test_cosmos3_scheduler_parity.py: FastVideo vendored UniPC (flow
config) vs the framework FlowUniPCMultistepScheduler. timesteps bit-exact,
sigmas within fp32 epsilon (~1e-8), full multi-step trajectory < ~1e-6
across shift in {10 (video), 3 (t2i)} and steps in {4,10,35}.
- Fix test_cosmos3_denoise_cfg_parity to use the framework scheduler as the
oracle (it previously compared diffusers-vs-diffusers, so the scheduler
was never actually checked against the framework).
Integration fixes to run end-to-end on real weights:
- registry: map checkpoint DiT class Cosmos3OmniTransformer ->
Cosmos3VFMTransformer.
- component_loader: route the "text_tokenizer" module to TokenizerLoader;
filter scheduler config to the class's __init__ params (robust to
diffusers schema drift like shift_terminal/sigma_min/sigma_max).
- DiT: materialize_non_persistent_buffers (recompute rotary_emb.inv_freq
after meta-device FSDP load) + cast vision latents / timestep embeds to
the compute dtype (no-ops in the fp32 parity tests).
- sequence_packing: move packed tensors to the model device in to_dit_kwargs.
- config: drop the unused identity text-preprocess (Cosmos3 has no text
encoder; the DiT tokenizes in the denoise stage).
Adds examples/inference/basic/basic_cosmos3_new_api.py (new VideoGenerator
API). Verified: 1280x704, 29 frames, 35 steps on B200 -> coherent video
matching the prompt (no NaNs, real pixel variance). Full cosmos3 parity
suite green.
SolitaryThinker