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