Add the one new native piece for image-conditioned reasoning: a deepstack-capable
reasoner backbone forward. The other parts are already proven/reusable — the
Qwen3-VL vision_encoder (transformers, bit-exact vs framework) and the multimodal
prefill (masked-scatter of image features + get_rope_index, standard transformers
Qwen3-VL boilerplate, reusable like the tokenizer/encoder).
- `Cosmos3VFMTransformer.reason_forward(inputs_embeds, position_ids,
deepstack_embeds, visual_pos_mask)`: runs the und (causal) layers over a single
text(+image) sequence, injecting per-layer deepstack visual embeds at the visual
positions for the first `len(deepstack_embeds)` layers (Qwen3-VL deepstack),
then final `norm`. Mirrors the framework `reasoner_forward`.
- Parity (`test_cosmos3_reasoning_parity::test_deepstack_reasoner_forward_matches_framework`):
native `reason_forward` vs framework `reasoner_forward` given identical
inputs_embeds + positions + per-layer deepstack + visual mask — hidden states
max abs diff = 0.0, mean = 0.0.
All five omni modalities are now bit-exact vs the framework (video/audio/action
gen, text reasoning, vision encoding, deepstack reasoner). Full suite: 150 passed.
Image-conditioned reasoning needs the Qwen3-VL vision_encoder. The checkpoint's
vision_encoder is a standard transformers `Qwen3VLVisionModel` (the framework
ships its own copy of the same model). Like the Qwen2 tokenizer, FastVideo reuses
the transformers model (no diffusers) rather than re-porting a 27-layer ViT.
Parity (`test_cosmos3_vision_encoder_parity.py`): transformers `Qwen3VLVisionModel`
vs the framework `Qwen3VLVisionModel`, tiny config, weights copied across,
identical (hidden_states, grid_thw) forward — embeds max abs diff = 0.0, mean = 0.0
across grids. Also verified on the REAL 1.15 GB checkpoint: both strict-load and
produce identical [N, out_hidden] embeds (max=mean=0.0).
All five omni modalities now have bit-exact framework parity: video gen
(T2V/I2V/T2I), audio gen (t2vs), action gen, text reasoning, vision encoding.
Remaining: wire image-conditioned reasoning (deepstack multimodal prefill joining
the two already-proven components). Full suite: 149 passed.
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.
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.
First audio component: the Cosmos3 sound tokenizer's decode path (t2vs needs
only DECODE — generate sound latents, decode to waveform).
Finding: the shipped `sound_tokenizer` checkpoint is decoder-only (`decoder.*`,
the SpectrogramConvNeXt encoder is not exported) in diffusers AutoencoderOobleck
naming, but with SnakeBeta activations (alpha+beta, logscale) and
weight_g/weight_v weight-norm — i.e. exactly FastVideo's existing native
`OobleckVAE` decoder. So no from-scratch port: reuse it.
- `OobleckDecoderBlock`: add `output_padding=stride % 2` to the transpose conv.
A provable no-op for even strides (Stable Audio: [2,4,4,8,8]); required for the
odd stride in Cosmos3 (`[2,4,5,6,8]`), where the framework's
`output_padding=stride%2` otherwise makes the decode 1 sample longer per
odd-stride block. Without it, parity diverged (60 vs 59 samples).
- `Cosmos3SoundVAE` (`fastvideo/models/audio/cosmos3_avae.py`): decoder-only
wrapper — `decode` (Oobleck decoder + clamp to [-1,1], matching AVAEModel),
`get_latent_num_samples` (N // hop_size=1920), `from_pretrained` (reads the
checkpoint config, strict-loads `decoder.*`). 48 kHz stereo, 473M params.
- Parity: `test_cosmos3_avae_parity.py` maps the framework OobleckDecoder
(`avae_utils.models`, Sequential naming, the oracle) into the FastVideo decoder
and asserts bit-exact decode across stride patterns incl. odd (5) — max abs
diff 0.0. Verified the real 1.9GB checkpoint strict-loads and decodes
[1,64,25] -> [1,2,48000] (1s @ 48kHz stereo).
Full cosmos3 suite: 121 passed, 0 skipped.
Add and verify the text-to-image path, and fix the UniPC flow_shift selection
it exposed.
Bug: the stage chose flow_shift by task (`3.0 if is_t2i else 10.0`). The
framework selects it purely by the named resolution bucket
(`OmniSampleArgs._RESOLUTION_SHIFT_DEFAULTS`, 8B backbone: 256->3.0, 480->5.0,
720/768->10.0); T2V/I2V/T2I share a shift at a given resolution. The task-based
rule only coincidentally matched (T2V@720, T2I@256). The canonical Cosmos3 T2I
is 960x960 (the "720" bucket -> 10.0), so `is_t2i->3.0` was wrong for real T2I.
- Replace with `_flow_shift_for_resolution(h, w)` mapping the longest side to the
framework bucket (<=320:3.0, 640-832:5.0, 960-1360:10.0), applied to all tasks.
- Parity: `test_cosmos3_flow_shift_parity.py` checks the mapping against the
framework's `{VIDEO,IMAGE}_RES_SIZE_INFO` x `_RESOLUTION_SHIFT_DEFAULTS` for
every (resolution, aspect) in the 8B rows (20 cases).
- Update the call-graph mode-dispatch tests for resolution-based shift.
Also harden `_image_to_video_tensor`'s tensor branch to respect the FastVideo
[-1,1] convention for already-preprocessed conditioning tensors (the PIL path,
used by the real pipeline, stays framework-exact and bit-parity-tested).
Example: `examples/inference/basic/basic_cosmos3_t2i_new_api.py` (num_frames=1,
960x960). Verified on B200 (real weights, 35 steps): a coherent red-panda image
matching the prompt, flow_shift=10.0. Full cosmos3 suite: 118 passed, 0 skipped.
Wire up and verify the Cosmos3 image-to-video path. The denoise stage already
handled the I2V conditioning math (encode -> condition mask -> keep frame 0
clean -> velocity zeroed on condition frames, all matching the framework), but
the conditioning pixel clip was built wrong.
Bug: `_image_to_video_tensor` zero-filled every frame after frame 0. The
framework (`cosmos_framework.inference.vision.build_conditioned_video_batch`)
instead fills frame 0 with the image and **repeats the last conditioning frame**
for the rest of the clip (a static video) before VAE-encoding. Because the Wan
VAE is temporal (4x), latent frame 0 (the kept-clean condition frame) depends on
several pixel frames, so zero-filling produces a wrong conditioning latent.
Fix `_image_to_video_tensor` to be framework-faithful:
- aspect-preserving resize + center crop + uint8 quantization, then `/127.5 - 1`
(mirrors `load_conditioning_image` / `_resize_and_center_crop`);
- static-repeat fill across all frames (mirrors `build_conditioned_video_batch`).
Parity: add `test_cosmos3_i2v_conditioning_parity.py` comparing the FastVideo
conditioning clip against the framework's `load_conditioning_image` +
repeat-fill across aspect/size/frame-count cases — bit-exact (max abs diff 0.0).
Example: `examples/inference/basic/basic_cosmos3_i2v_new_api.py` (new
VideoGenerator API, image via `InputConfig(image_path=...)`,
`assets/images/cyclist.jpg` default).
Verified on B200 (1280x704, 29f, 35 steps, real weights): frame 0 reproduces
the conditioning image and the clip generates coherent forward motion following
the prompt. Full cosmos3 suite: 98 passed, 0 skipped.
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.
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.
SolitaryThinker