diffusers/examples/cosmos3/README.md

# Cosmos3 — smoke-test runner

The canonical reference for `Cosmos3OmniModularPipeline` lives in the diffusers docs:
[`docs/source/en/api/pipelines/cosmos3.md`](../../docs/source/en/api/pipelines/cosmos3.md). Use the
examples there as the source of truth for application code — they cover text-to-image,
text-to-video, image-to-video, and text+sound modes.

This directory provides two files:

- `inference_cosmos3.py` — the runnable `Cosmos3OmniModularPipeline` CLI (text-to-image/video,
  image-to-video, sound, action modes). Single-GPU by default; pass `--tp-degree` / `--cp-degree`
  and launch with `torchrun` to run any modality multi-GPU (see [Multi-GPU inference](#multi-gpu-inference-context-parallelism)
  below).
- `cosmos_parallel.py` — the importable multi-GPU helpers (context + tensor parallelism). No
  `main`; the CLI imports from it. Read it to understand or adapt the sharding.

## Setup

```bash
pip install -r examples/cosmos3/requirements.txt
```

## Usage

Text-to-image:

```bash
python examples/cosmos3/inference_cosmos3.py \
    --prompt "A medium shot of a modern robotics research laboratory…" \
    --num-frames 1
```

Text-to-video:

```bash
python examples/cosmos3/inference_cosmos3.py \
    --prompt "A waterfall cascading down a rocky cliff in a lush forest."
```

Image-to-video:

```bash
python examples/cosmos3/inference_cosmos3.py \
    --prompt "The right robotic hand picks up the red sphere…" \
    --vision-path https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/assets/robot_153.jpg
```

Video-to-video (condition on the leading frames of a clip and continue it):

```bash
python examples/cosmos3/inference_cosmos3.py \
    --prompt "A robotic arm finishes pouring liquid into the glass." \
    --video-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4" \
    --condition-frame-indexes-vision 0,1 \
    --condition-video-keep first
```

Text-to-video-with-sound (sound-capable checkpoint only):

```bash
python examples/cosmos3/inference_cosmos3.py \
    --prompt "A waterfall in a lush forest." \
    --enable-sound
```

Action forward dynamics, robot domain (predict video from an observation video and a provided action chunk):

```bash
python examples/cosmos3/inference_cosmos3.py \
    --model nano \
    --prompt "Put the pot to the left of the purple item." \
    --vision-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_0.mp4" \
    --action-mode forward_dynamics \
    --action-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_0.json" \
    --action-chunk-size 16 \
    --domain-name bridge_orig_lerobot \
    --resolution-tier 480 --fps 5 \
    --num-inference-steps 30 --guidance-scale 1.0 --flow-shift 10.0 --seed 0 \
    --output results/cosmos3_forward_dynamics_robot
```

Action forward dynamics, autonomous-vehicle domain:

```bash
python examples/cosmos3/inference_cosmos3.py \
    --model nano \
    --prompt "You are an autonomous vehicle planning system." \
    --vision-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/av_vision_25_73d01c91-51f0-46cf-9b76-5682a76fb349.mp4" \
    --action-mode forward_dynamics \
    --action-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/av_action_25.json" \
    --action-chunk-size 60 \
    --domain-name av \
    --resolution-tier 480 --fps 10 \
    --num-inference-steps 30 --guidance-scale 1.0 --flow-shift 10.0 --seed 0 \
    --output results/cosmos3_forward_dynamics_av
```

Action inverse dynamics, robot domain (predict actions from an observed video):

```bash
python examples/cosmos3/inference_cosmos3.py \
    --model nano \
    --prompt "Put the pot to the left of the purple item." \
    --vision-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_0.mp4" \
    --action-mode inverse_dynamics \
    --action-chunk-size 16 \
    --domain-name bridge_orig_lerobot \
    --resolution-tier 480 --fps 5 \
    --num-inference-steps 30 --guidance-scale 1.0 --flow-shift 10.0 --seed 0 \
    --output results/cosmos3_inverse_dynamics_robot
```

Action inverse dynamics, autonomous-vehicle domain:

```bash
python examples/cosmos3/inference_cosmos3.py \
    --model nano \
    --prompt "You are an autonomous vehicle planning system." \
    --vision-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/av_vision_25_73d01c91-51f0-46cf-9b76-5682a76fb349.mp4" \
    --action-mode inverse_dynamics \
    --action-chunk-size 60 \
    --domain-name av \
    --resolution-tier 480 --fps 10 \
    --num-inference-steps 30 --guidance-scale 1.0 --flow-shift 10.0 --seed 0 \
    --output results/cosmos3_inverse_dynamics_av
```

Action policy, robot domain (predict both future video and actions from the first observation frame):

```bash
python examples/cosmos3/inference_cosmos3.py \
    --model nano \
    --prompt "Put the pot to the left of the purple item." \
    --vision-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_0.mp4" \
    --action-mode policy \
    --action-chunk-size 16 \
    --domain-name bridge_orig_lerobot \
    --resolution-tier 480 --fps 5 \
    --num-inference-steps 30 --guidance-scale 1.0 --flow-shift 10.0 --seed 0 \
    --output results/cosmos3_policy_robot
```

Action policy, autonomous-vehicle domain:

```bash
python examples/cosmos3/inference_cosmos3.py \
    --model nano \
    --prompt "You are an autonomous vehicle planning system. Please go backward." \
    --vision-path "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/av_vision_25_73d01c91-51f0-46cf-9b76-5682a76fb349.mp4" \
    --action-mode policy \
    --action-chunk-size 60 \
    --domain-name av \
    --resolution-tier 480 --fps 10 \
    --num-inference-steps 30 --guidance-scale 1.0 --flow-shift 10.0 --seed 0 \
    --output results/cosmos3_policy_av
```

Action modes use `action_chunk_size + 1` conditioning frames. `forward_dynamics` consumes `--action-path`; `inverse_dynamics` and `policy` write predicted actions to `sample_action.json` in model-normalized action space. This script loads `--vision-path` as a video for all action modes; `policy` and `forward_dynamics` condition only on the first frame, while `inverse_dynamics` uses the whole video.

Pass `--prompt` as a plain task description and select the camera perspective with `--view-point` (default `ego_view`); the pipeline builds the structured action caption (task, viewpoint, duration, FPS, resolution) the model was trained on. Do not hand-write the viewpoint sentence into `--prompt`.

`--resolution-tier` is a resolution *tier* (`256`/`480`/`704`/`720`). The tier keys a table of predefined aspect-ratio canvases; the one closest to the input aspect ratio becomes the padded conditioning canvas. It is not the output frame size: the input is downscaled (never upscaled) and padded to fill the canvas, then the padding is cropped from the latents so the decoded output follows the downscaled input content. `--height` / `--width` (and `--num-frames`) are ignored for action modes.

Pick the tier that matches the native resolution of your conditioning input (`480` for ~480p, `720` for ~720p). A tier below your input downscales it and discards detail; a tier above your input gains no resolution (content is never upscaled), wastes compute on padding, and is a train/inference distribution mismatch that can degrade quality.

### Useful flags

| Flag | Default | Description |
|---|---|---|
| `--prompt` | (required) | Text prompt. |
| `--vision-path` | `None` | URL or local path for an image-conditioning frame (image-to-video), or the image/video conditioning for action modes. |
| `--num-frames` | `189` | `1` = image, otherwise number of video frames (`189` ≈ 7.9 s @ 24 FPS). Ignored for action modes (derived from `--action-chunk-size`). |
| `--height` / `--width` | `720` / `1280` | Output resolution (must be a multiple of the VAE spatial scale factor). Ignored for action modes; use `--resolution-tier`. |
| `--resolution-tier` | `480` | Action resolution tier (`256`/`480`/`704`/`720`): selects the aspect bin / padded conditioning canvas, not the output size. |
| `--fps` | `24.0` | Frame rate of the generated video. |
| `--flow-shift` | `None` | Override `UniPCMultistepScheduler.flow_shift` (and force `use_karras_sigmas=False`); left at the checkpoint default when unset. Cosmos3 runs use `10.0`. |
| `--enable-sound` | off | Generate a synchronized audio track. |
| `--action-mode` | `None` | Enable action conditioning/generation. One of `forward_dynamics`, `inverse_dynamics`, or `policy`. |
| `--action-path` | `None` | URL or local JSON action path for `forward_dynamics`. |
| `--action-chunk-size` | `None` | Number of action tokens. Action runs generate/use `action_chunk_size + 1` video frames. |
| `--domain-name` | `None` | Action embodiment domain, for example `bridge_orig_lerobot` or `av`. |
| `--view-point` | `ego_view` | Camera perspective for the action caption's framing (`ego_view`, `third_person_view`, `wrist_view`, `concat_view`). Action only. |
| `--no-duration-template` | off | Skip the duration metadata sentence appended to the prompt and negative prompt. Ignored for `--num-frames 1` and for action modes (which build a structured caption instead). |
| `--no-resolution-template` | off | Skip the resolution metadata sentence appended to the prompt and negative prompt. Ignored for action modes. |
| `--output` | `.` | Directory to write `sample.jpg` or `sample.mp4`. |

## Multi-GPU inference (context parallelism)

Cosmos 3 can be sharded across GPUs on two orthogonal axes (implemented in `cosmos_parallel.py`):

- **Context parallelism (CP)** — `enable_cosmos3_context_parallel`. The *sequence* is sharded
  across GPUs and attention runs with two Ulysses all-to-all collectives per layer, cutting
  per-step latency for long videos / high resolutions. Weights are replicated, so this is for
  models that already fit one GPU (`Nano`).
- **Tensor parallelism (TP)** — `enable_cosmos3_tensor_parallel`. The attention and MLP *weight*
  matrices are sharded across GPUs (Megatron-style), so a checkpoint that doesn't fit one GPU
  (`Super`, ~120 GB) loads. The sequence is not sharded.
- **TP + CP** — both at once on a 2-D `(tp, cp)` mesh: a large model *and* latency.

The model itself carries no parallelism logic — it exposes small no-op shard/gather seams, and
`cosmos_parallel.py` implements the entire path (collectives, GQA KV-head handling, ragged-length
padding, the dual-pathway attention, weight sharding) behind those two helpers. It is meant to be
read end to end and adapted.

The CLI imports these helpers, so you run **any modality** (text-to-image/video, image-to-video,
sound, action modes) multi-GPU by adding `--tp-degree` / `--cp-degree` and launching with
[torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html) — `--tp-degree * --cp-degree`
must equal `--nproc_per_node`:

```bash
# CP only (Nano): CP degree must divide the 32 query heads.
torchrun --nproc_per_node 4 examples/cosmos3/inference_cosmos3.py --model nano --cp-degree 4 --prompt "..."

# TP only (Super): TP degree must divide the 64 query heads and 8 KV heads.
torchrun --nproc_per_node 4 examples/cosmos3/inference_cosmos3.py --model super --tp-degree 4 --prompt "..."

# TP + CP (Super), 4 GPUs as 2 x 2, with sound:
torchrun --nproc_per_node 4 examples/cosmos3/inference_cosmos3.py \
    --model super --tp-degree 2 --cp-degree 2 --enable-sound --prompt "A waterfall in a forest."
```

### Modular pipeline setup

The CLI uses `Cosmos3OmniModularPipeline`, not the legacy task pipeline. Its distributed setup order is important:

1. Construct it with `Cosmos3OmniModularPipeline.from_pretrained(...)` and call
   `pipe.load_components(dtype=torch.bfloat16)` while it is still on CPU. Do not use
   `device_map`.
2. Initialize the NCCL process group and call `torch.cuda.set_device(local_rank)` before building
   the device mesh or applying TP. If needed, replace the scheduler with
   `pipe.update_components(scheduler=...)` and apply TP to `pipe.transformer` while it is still on
   CPU.
3. Move the pipeline to `cuda:${LOCAL_RANK}`, select the `native` attention backend, and then
   enable CP (or the GQA-safe dense attention helper for TP-only runs).
4. Before generation, call `pipe.enable_safety_checker()` to load and enable the default checker,
   or `pipe.disable_safety_checker()` to explicitly opt out. The task-pipeline
   `enable_safety_checker=` construction argument and `enable_safety_check=` call argument do not
   configure the modular pipeline.

Modular calls request outputs explicitly: `output="videos"` returns frames directly, while
`output=["videos", "sound", "sampling_rate", "action"]` returns a dictionary. Read values such as
`outputs["videos"]` from that dictionary rather than `result.video`, `result.sound`, or
`result.action`; sound and action are `None` when their respective workflows are not used. The
[pipeline documentation](../../docs/source/en/api/pipelines/cosmos3.md#context-parallelism) has a
complete direct-use example.

Notes:

- The helpers use the `native` attention backend (the only one that supports GQA's `enable_gqa`),
  and expand the KV heads to the query-head count so SDPA picks the flash kernel — passing
  `enable_gqa=True` forces the math kernel, which materializes the full `[S, S]` scores and OOMs
  on long sequences.
- Only Ulysses is supported (not ring attention).
- The CP/Ulysses degree must divide the query heads (32 for `Nano`, 64 for `Super`). For TP,
  `tp` must divide the KV heads (8), and `tp * cp` must divide the query heads.
- TP all-reduces on every block, so it's bandwidth-heavy — use the smallest TP degree that makes
  the weights fit and put the remaining GPUs into CP.
- Generation size is set with the usual CLI flags (`--num-frames` / `--height` / `--width`), and
  multi-GPU runs require a seed for reproducibility across ranks (the CLI sets one if you omit `--seed`).
- On some multi-GPU topologies the first NCCL all-to-all can hang; if a run stalls at the first
  denoising step, set `NCCL_P2P_DISABLE=1` before launching.

See the [pipeline docs](../../docs/source/en/api/pipelines/cosmos3.md#context-parallelism) for how
to enable CP and TP from your own pipeline code.