Skip to content

Inference

Skill: .agents/skills/cosmos3-inference/SKILL.md

Prerequisites:

For caption-and-control generation with Cosmos3-Nano-Sim-Transfer, see the Simulation Transfer recipe. For multiview driving video and optional joint LiDAR generation, see the Transfer Auto recipe. For its autoregressive Cosmos3-Nano-Sim-Auto student, see the Sim Auto recipe.

Arguments:

  • -i, --input-files: Path to the sample argument file(s) (JSON, JSONL, YAML). Accepts quoted glob patterns (e.g. "inputs/*.json").
  • -o, --output-dir: Output directory.

Outputs:

  • <sample_name>/
  • sample_args.json: Sample arguments.
  • sample_outputs.json: Generation status, action (if enabled).
  • vision.jpg, vision.mp4: Vision output (if enabled).

To see all available arguments:

python -m cosmos_framework.scripts.inference --help

Quick Start

Single-GPU

Use python -m directly. Suitable for --parallelism-preset=latency on a single GPU, or for quick experimentation:

python -m cosmos_framework.scripts.inference \
    --parallelism-preset=latency \
    -i "inputs/omni/t2v.json" \
    -o outputs/omni_nano \
    --checkpoint-path Cosmos3-Nano \
    --seed=0

Note: Cosmos3-Super (32B) does not fit on a single 80 GB H100 — see Cosmos3-Super for the multi-GPU recipes.

Cosmos3-Edge

Cosmos3-Edge is a compact 2B omni model. It fits comfortably on a single GPU and is the recommended starting point for single-GPU inference. It supports every mode except audio (enable_sound), since the checkpoint ships without a sound tokenizer — the audio-enabled examples (inputs/omni/t2vs.json, inputs/omni/i2vs.json) are therefore not supported.

python -m cosmos_framework.scripts.inference \
    --parallelism-preset=latency \
    -i "inputs/omni/t2i.json" \
    -o outputs/omni_edge \
    --checkpoint-path Cosmos3-Edge \
    --seed=0

To run every supported example in one batch (the action_*.json glob covers the action modes; the audio-enabled t2vs.json / i2vs.json are intentionally excluded):

python -m cosmos_framework.scripts.inference \
    --parallelism-preset=latency \
    -i inputs/omni/t2i.json inputs/omni/t2v.json inputs/omni/i2v.json inputs/omni/v2v.json inputs/omni/action_*.json \
    -o outputs/omni_edge \
    --checkpoint-path Cosmos3-Edge \
    --seed=0

Multi-GPU

Use torchrun --nproc-per-node=N when launching across multiple GPUs (N > 1). By default the model weights are sharded (FSDP) across all N GPUs, reducing per-GPU weight storage. The full request must also fit activations, tokenizers and output buffers. The throughput preset runs that single sharded replica over a batch; the latency preset additionally needs --dp-shard-size=1 on multiple GPUs so the ranks are free for context parallelism (see Parallelism Arguments).

Cosmos3-Nano

torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \
    --parallelism-preset=throughput \
    -i "inputs/omni/*.json" \
    -o outputs/omni_nano \
    --checkpoint-path Cosmos3-Nano \
    --seed=0

Note: The progress bar only prints on rank 0.

Note: With the default full-GPU sharding, this same command also runs Cosmos3-Super (32B) on 8×80 GB H100 — the weights are sharded (FSDP) across all 8 GPUs. See Cosmos3-Super for the explicit-axis variants and the 4-GPU recipe.

Cosmos3-Super

Cosmos3-Super (32B) must be sharded across multiple GPUs to fit in 80 GB H100 memory. The default already shards the model across every visible GPU (FSDP), so the throughput preset fits it directly; the commands below pin the axes explicitly (pure FSDP, no context- or CFG-parallelism overlay) and add the 4-GPU recipe.

4 GPUs:

torchrun --nproc-per-node=4 -m cosmos_framework.scripts.inference \
    --parallelism-preset=throughput \
    --dp-shard-size=4 --dp-replicate-size=1 \
    --cp-size=1 --cfgp-size=1 \
    -i "inputs/omni/*.json" \
    -o outputs/omni_super \
    --checkpoint-path Cosmos3-Super \
    --seed=0

8 GPUs:

torchrun --nproc-per-node=8 -m cosmos_framework.scripts.inference \
    --parallelism-preset=throughput \
    --dp-shard-size=8 --dp-replicate-size=1 \
    --cp-size=1 --cfgp-size=1 \
    -i "inputs/omni/*.json" \
    -o outputs/omni_super \
    --checkpoint-path Cosmos3-Super \
    --seed=0

The four --{dp,cp,cfgp}-*-size flags override the auto-selected values from --parallelism-preset. Super supports text2image, text2video, and image2video (see Models).

Models

Model Arguments Modes
Cosmos3-Nano --checkpoint-path=Cosmos3-Nano All
Cosmos3-Edge --checkpoint-path=Cosmos3-Edge All except audio (enable_sound)
Cosmos3-Super --checkpoint-path=Cosmos3-Super text2image, text2video, image2video

Modes

model_mode selects the generation modality. The table below lists every supported mode with its required sample fields and a paired example file.

model_mode Inputs Outputs Required sample fields Example
text2image text prompt vision.jpg prompt inputs/omni/t2i.json
text2video text prompt vision.mp4 prompt inputs/omni/t2v.json
image2video text prompt + image vision.mp4 prompt, vision_path inputs/omni/i2v.json
video2video text prompt + video vision.mp4 prompt, vision_path inputs/omni/v2v.json
forward_dynamics observation image/video + prompt + actions future visual rollout in vision.mp4 domain_name, vision_path, action_path inputs/omni/action_forward_dynamics_av.json, inputs/omni/action_forward_dynamics_camera.json, inputs/omni/action_forward_dynamics_robot.json, inputs/omni/action_forward_dynamics_batch.jsonl
inverse_dynamics observation video + prompt predicted action sequence in sample_outputs.json domain_name, vision_path inputs/omni/action_inverse_dynamics_av.json, inputs/omni/action_inverse_dynamics_robot.json, inputs/omni/action_inverse_dynamics_batch.jsonl
wam observation image/video + prompt predicted action sequence in sample_outputs.json + future visual rollout in vision.mp4 domain_name, vision_path inputs/omni/action_policy_av.json, inputs/omni/action_policy_robot.json, inputs/omni/action_policy_batch.jsonl

Set enable_sound: true on a text2video sample (see inputs/omni/t2vs.json) to also generate audio. To run every example in one batch, use -i "inputs/omni/*.json".

Parallelism Arguments

By default the model weights are sharded (FSDP) across all visible GPUs (dp_shard_size = WORLD_SIZE, dp_replicate_size = 1). Sharding reduces per-GPU weight storage but does not guarantee that the complete workload fits. Override any axis with the --dp-shard-size / --dp-replicate-size / --cp-size / --cfgp-size flags.

Increasing context parallelism can reduce per-GPU activation memory and speed up large requests, but communication overhead can offset the benefit. Compare latency on the same workload; more GPUs or a higher CP degree do not guarantee faster inference. For Transfer-Auto and Sim-Auto, follow the model-specific commands and Auto memory guidance.

  • --parallelism-preset
  • latency: Minimize wall-clock per sample by splitting each sample across GPUs with context parallelism. On multiple GPUs, also pass --dp-shard-size=1 so the ranks are used for context/CFG parallelism instead of weight sharding. Used for real-time jobs.
  • throughput: No context parallelism (cp=cfgp=1); the model is sharded across all GPUs and a single replica processes the batch. Used for batch jobs.
  • --dp-shard-size: Number of ranks the model is sharded over (FSDP). Defaults to all ranks (WORLD_SIZE).
  • --max-num-seqs: Maximum number of samples batched together per replica.
  • --attention-io-layout {sequence_sharded,replicated}: Tensor layout around attention under context parallelism. sequence_sharded (default) keeps hidden states sequence-sharded with Ulysses all-to-all around attention; replicated keeps them replicated and slices heads locally.
  • --use-batched-cfg: Run the conditional and unconditional CFG branches as one batched forward. Off by default because it raises peak activation memory; it is turned off automatically when --cfgp-size 2 is used.
  • --ar-post-saturation-mode {default,static-compile,cuda-graph}: Autoregressive execution after the rolling KV window saturates. default keeps the global --use-torch-compile / --use-cuda-graphs behaviour.

Model Arguments

  • --quantization-method {mxfp8,nvfp4,fp8}: Post-training quantization applied at load time (replicated weights only, --dp-shard-size 1). mxfp8 / nvfp4 need Blackwell; fp8 also runs on Hopper and Ada.
  • --quantization-fp8-granularity {per_row,per_tensor}: Scale granularity for fp8 (default per_row).
  • --experiment-overrides: Extra config overrides as dotlist entries. Both --experiment-overrides model.config.a=1 model.config.b=2 and the bracketed form --experiment-overrides "[model.config.a=1,model.config.b=2]" are accepted.

Sample Arguments

Sample arguments are read from multiple sources (in priority order):

  • CLI overrides (e.g. --model-mode=text2video): Overrides for all samples.
  • Input files (e.g. --input-files "inputs/omni/*t2i*.json"): Single sample per input.
  • Defaults: cosmos_framework/inference/defaults/<model_mode>: Defaults for all samples.

For debugging, the full set of sample arguments is saved to <output_dir>/<sample_name>/sample_args.json.

Common arguments:

  • model_mode: Generation modality. See Modes above for all options.
  • seed: Random seed for reproducibility.

Note: Condition file paths are relative to the input file.

Text

  • prompt: Inline text prompt.

Vision (Image/Video)

Common arguments:

  • fps: Condition and output frames per second.
  • resolution ("256", "480", "720", "768"): Condition and output resolution tier, not a pixel dimension. The output width and height are looked up from the tier and the aspect ratio — see Resolution tiers below.
  • aspect_ratio (1,1, 4,3, 3,4, 16,9, 9,16): Condition and output aspect ratio. Defaults to 16,9.

Resolution tiers

For the 480 and 720 tiers the 16,9 output height matches the tier name (832x480, 1280x720). The 256 tier does not: its widescreen entries are smaller than the name suggests.

aspect_ratio 256 480 720 768
1,1 256x256 640x640 960x960 1024x1024
4,3 320x256 736x544 1104x832 1184x880
3,4 256x320 544x736 832x1104 880x1184
16,9 320x192 832x480 1280x720 1360x768
9,16 192x320 480x832 720x1280 768x1360

So resolution: "256" with the default 16,9 produces a 320x192 output, not 256 pixels tall. Use aspect_ratio: "1,1" or "4,3" if you need a 256-pixel short edge at that tier. The authoritative tables are IMAGE_RES_SIZE_INFO and VIDEO_RES_SIZE_INFO in cosmos_framework/data/generator/utils.py.

Condition arguments:

  • vision_path: Path to an image or video file (local path or URL).

Generation arguments:

Outputs vision.jpg or vision.mp4 depending on num_frames.

Multiview camera + LiDAR transfer

A joint cosmos3_av_cam_lidar_transfer checkpoint generates the camera views and the LiDAR range map together. Cameras go in multiview.views, one entry per camera with its own control_path; the HD map goes in lidar.control_path; wsm stays empty.

Download and verify the shared Auto example first. Its input-transfer-auto.jsonl uses seven cameras, 720p, 30 FPS and 297 output frames, with vision_path per camera and multiview.condition_video_as_image: true for first-frame RGB conditioning. The sample also contains the corresponding Sim-Auto request; see Sim-Auto inference.

torchrun --standalone --nproc-per-node=4 -m cosmos_framework.scripts.inference \
  -i "$AUTO_SAMPLE_ROOT/input-transfer-auto.jsonl" \
  -o outputs/transfer_auto \
  --checkpoint-path "$TRANSFER_AUTO_CHECKPOINT" \
  --cp-size 4 --cfgp-size 1 \
  --no-use-torch-compile --no-diffusion-cache

Both Auto requests set multiview.load_caption_from_data: true and give every camera a caption_path. That is not optional here: this checkpoint tokenizes one caption per camera, and generation is refused without one caption file per camera. Each camera's caption supplies its text conditioning; the top-level prompt is empty. Captions must cover the requested generation horizon.

Set multiview.first_caption_paragraph_only: true for a checkpoint trained with first_caption_paragraph_only (the Transfer-Auto stage-3 and LiDAR-transfer recipes). Each loaded caption is then trimmed to its opening paragraph before the camera headers are added, so inference reads the same text the checkpoint was trained on. The default (false) keeps full captions.

Outputs are vision_viewXX_<camera_key>.mp4 per camera, plus a single lidar.safetensors holding range_m, intensity and frame_ids. LiDAR runs at 10 Hz against a 30 fps camera clock, so a 297-frame generation produces 99 sweeps.

To condition on measured LiDAR, set lidar.condition_path to a metric 3-channel .pt clip or a prepared lidar_rangemap .tar that starts at the same instant as the controls, and lidar.num_conditional_sweeps (default 1) to the number of its sweeps to keep as a clean prefix; the rest of the sweeps are generated. An archive prefix must carry the same source frame IDs as the start of a .tar control, and a .pt prefix is refused beside one. Prefix conditioning needs a temporally causal LiDAR tokenizer without temporal compression, and causal (rolling) checkpoints accept one prefix sweep only.

The shared sample is published in the nvidia/Cosmos3-Nano-Transfer-Auto model repository. Download the whole pinned sample folder once; the requests then read its local relative paths. Authenticate with access to the repository if required.

Action

Common arguments:

  • action_chunk_size: Number of action steps in the chunk. The action media loader reads at most action_chunk_size + 1 observation frames.
  • domain_name: Domain name passed to the action domain registry, such as bridge_orig_lerobot, camera_pose, or av.
  • view_point: Viewpoint description injected into the action prompt, such as ego_view.

Condition arguments:

  • action_path: JSON action sequence. Required for forward_dynamics; each row is one action step and each column is one raw action dimension.
  • image_size: Action input resize bucket. The value is passed as the action media resolution bucket; examples use 256 for LIBERO and 480 for AV.

The action output is written to sample_outputs.json.

See the Modes table above for the action mode inputs/outputs and example files.

For the Bimanual checkpoint, see the Bimanual inference guide, including the normalized 59D sidecars, per-row robot/camera domains, plain-text prompts, and fixed four-step sampling schedule.

Reasoner

model_mode=reasoner generates text (written to reasoner_text.txt) from a prompt and an optional vision_path. The vision_path may point to an image (.jpg/.png/…) or a video (.mp4): a video is decoded and uniformly sampled into frames that condition the reasoner.

  • video_fps: frames per second to sample from the video (default: the decoder's default of 2.0).

Examples: inputs/reasoner/reasoner.json (text), inputs/reasoner/reasoner_image.json (image), inputs/reasoner/reasoner_video.json (video).

Cosmos3-Edge vision tower: with a local --checkpoint-path, the vision tower and processor load from the checkpoint itself (vision_encoder/ + processor files, bundled by default by export_model) — such exports run fully offline (add --no-guardrails, since guardrails download their own models). If the checkpoint has no bundle (e.g. a --no-vit export), the tower is fetched from nvidia/Cosmos3-Edge on the Hub instead; when that fails (offline, missing HF_TOKEN), the error says so and suggests re-exporting with the default --vit. Nano/Super checkpoints without a vision tower (include_visual=false) reject reasoner image/video samples up front with a clear error.

Custom Defaults

To use your own default values instead of the built-in presets, pass a JSON file via the defaults_file field in your sample arguments:

{
    "defaults_file": "my_defaults.json",
    "prompt": "..."
}

The custom defaults file has the same format as the built-in presets. Fields you set explicitly in the sample argument file still take precedence over the custom defaults file.

Guardrails

Inference ships with guardrails enabled by default, sourced from nvidia/Cosmos-Guardrail1. Active filters: text blocklist (better-profanity + fuzzy match), text safety classifier (Qwen/Qwen3Guard-Gen-0.6B), video content-safety classifier, and RetinaFace face-blur post-processor. Pass --no-guardrails to disable, or --offload-guardrail-models to keep them on CPU between calls (saves GPU memory, adds latency).

Troubleshooting

Checkpoint Issue

If you encounter failures downloading checkpoints, refer to Downloading Base Checkpoints.

Checkpoint download commands are printed to the console. You can run them manually to debug issues.

Torch CUDA Out of Memory Error

Error: torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate X MiB

Optimize memory allocation:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

If that's not enough, see FAQ § OOM during inference for the full ladder (--dp-shard-size, --device-memory-utilization, --offload-guardrail-models).

NCCL Issue

Error:

[rank0]:[W415 18:57:09.249883195 ProcessGroupNCCL.cpp:5138] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()

Fatal Python error: Segmentation fault

Re-run with debugging enabled:

export NCCL_DEBUG=INFO
export TORCH_DISTRIBUTED_DEBUG=DETAIL
export CUDA_LAUNCH_BLOCKING=1

NCCL Plugin Issue

Error:

NCCL INFO Failed to initialize NET plugin Libfabric

Fatal Python error: Segmentation fault

Fix:

export NCCL_NET_PLUGIN=none