> ## Documentation Index
> Fetch the complete documentation index at: https://phyai.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# Cosmos3 Generation Mode

> Run Cosmos3 T2V and T2AV through the same single-GPU or multi-GPU plugin

export const ModelCard = ({title, subtitle, icon, rows = {}}) => {
  const entries = Object.entries(rows);
  const renderValue = value => {
    if (value === null || value === undefined) {
      return <span className="phyai-model-card__empty">—</span>;
    }
    if (Array.isArray(value)) {
      return <div className="phyai-model-card__tags">
                    {value.map((tag, index) => <span key={index} className="phyai-model-card__tag">
                            {tag}
                        </span>)}
                </div>;
    }
    if (typeof value === "string" || typeof value === "number") {
      return <span className="phyai-model-card__text">{value}</span>;
    }
    return value;
  };
  const hasHeader = title || subtitle || icon;
  return <div className="phyai-model-card not-prose">
            {hasHeader && <div className="phyai-model-card__header">
                    {icon && <div className="phyai-model-card__icon">{icon}</div>}
                    <div className="phyai-model-card__heading">
                        {title && <div className="phyai-model-card__title">{title}</div>}
                        {subtitle && <div className="phyai-model-card__subtitle">{subtitle}</div>}
                    </div>
                </div>}

            <div className="phyai-model-card__rows">
                {entries.map(([key, value]) => <div key={key} className="phyai-model-card__row">
                        <div className="phyai-model-card__label">{key}</div>
                        <div className="phyai-model-card__value">{renderValue(value)}</div>
                    </div>)}
            </div>
        </div>;
};

<ModelCard
  title="Cosmos3-Nano"
  subtitle="Text-to-Video / Text-to-Audio-Video · Single or Multi-GPU"
  icon="C"
  rows={{
"Model Type": "World Foundation Model",
"Weights": <a href="https://huggingface.co/nvidia/Cosmos3-Nano" target="_blank" rel="noreferrer" className="text-sm text-[#003399] dark:text-[#60A5FA] underline underline-offset-2 hover:opacity-80 break-all">huggingface.co/nvidia/Cosmos3-Nano</a>,
"Paths": ["T2V", "T2AV"],
"Entry Point": <code className="px-2 py-0.5 rounded bg-[#003399]/10 dark:bg-[#60A5FA]/15 text-[#003399] dark:text-[#60A5FA] text-xs font-mono">Cosmos3T2VScheduler</code>,
"Plugin": <code className="px-2 py-0.5 rounded bg-[#003399]/10 dark:bg-[#60A5FA]/15 text-[#003399] dark:text-[#60A5FA] text-xs font-mono">cosmos3</code>,
"Sampler": "UniPC",
"Default Size": "720x1280 · 189 frames · 35 steps",
"Param Precision": "bf16",
}}
/>

# Overview

Cosmos3's generation path turns a sentence into a video. Switch on the sound stream and the same denoising run also writes an audio track that follows the frames. T2V produces video only; T2AV advances the video latent and the sound latent on one timeline, so the picture surfaces out of the noise while the waveform takes shape beside it.

One plugin, `cosmos3`, covers both paths and both machine sizes. With the default `cfg_size=tp_size=1` the engine runs inline in your process. Ask for `--cfg 2` or `--tp N` and it spawns one worker per rank on the first visible GPUs instead; the request and the outputs stay the same.

<Warning>
  This path is built for correctness and reference alignment, not speed. The denoising loop is a Python-level UniPC loop with CUDA graphs off, and there are no dedicated kernels or batching yet. Read any timing you collect as a baseline.
</Warning>

# Architecture

The path follows PhyAI's usual <Tooltip headline="Engine + plugin" tip="Engine resolves an Entry by plugin name. Entry.setup() builds models, loads weights, and prepares the scheduler. Entry.step() accepts a canonical request and returns model output.">engine + plugin contract</Tooltip>. Inside the plugin the work splits into a handful of pieces:

<Tree>
  <Tree.Folder name="phyai/src/phyai/models/cosmos3" defaultOpen>
    <Tree.File name="main_cosmos3.py" />

    <Tree.File name="scheduler_cosmos3.py" />

    <Tree.File name="model_runner_cosmos3.py" />

    <Tree.File name="model_runner_vae_cosmos3.py" />

    <Tree.File name="modeling_cosmos3.py" />

    <Tree.File name="vae_wan.py" />

    <Tree.File name="avae_sound.py" />

    <Tree.File name="sampler_unipc.py" />

    <Tree.File name="configuration_cosmos3.py" />
  </Tree.Folder>
</Tree>

| Component                        | Responsibility                                                                            |
| -------------------------------- | ----------------------------------------------------------------------------------------- |
| `Cosmos3Entry`                   | Reads `Cosmos3Args`; loads the transformer, the VAE, and the AVAE when sound is on        |
| `Cosmos3T2VScheduler`            | Drives the denoising loop and the UniPC sampler; decodes video and sound at the end       |
| `Cosmos3T2VRunner`               | Calls the transformer and caches the text conditions, which do not depend on the timestep |
| `Cosmos3VAERunner`               | Turns video latent into pixels in `[0, 1]`                                                |
| `Cosmos3SoundVAERunner`          | Turns sound latent into a waveform in `[-1, 1]`                                           |
| `Cosmos3Processor`               | Tokenizes prompts and appends prompt metadata, outside the engine                         |
| `Cosmos3GenerationPostProcessor` | Moves results to the CPU and writes the mp4, outside the engine                           |

# Run path

<Steps>
  <Step title="Prepare weights">
    Download <a href="https://huggingface.co/nvidia/Cosmos3-Nano" target="_blank" rel="noreferrer">Cosmos3-Nano</a>. The examples expect this layout:

    ```text theme={null}
    /path/to/Cosmos3-Nano/
      transformer/
      vae/
      text_tokenizer/
      sound_tokenizer/   # required for T2AV
      scheduler/
    ```
  </Step>

  <Step title="Construct the engine">
    T2V needs the transformer and the VAE. T2AV also loads the AVAE from `sound_tokenizer`, which is what `load_sound=True` asks for.

    ```python theme={null}
    import torch

    from phyai.engine import Engine, EngineArgs
    from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
    from phyai.models.cosmos3.main_cosmos3 import Cosmos3Args

    checkpoint_dir = "/path/to/Cosmos3-Nano"
    with_sound = False

    engine = Engine(
        EngineArgs(
            plugin="cosmos3",
            plugin_args=Cosmos3Args(
                checkpoint_dir=checkpoint_dir,
                flow_shift=10.0,
                use_karras_sigmas=False,
                load_sound=(True if with_sound else None),
            ),
            config=EngineConfig(
                device=DeviceConfig(target="cuda", params_dtype=torch.bfloat16),
                runtime=RuntimeConfig(use_cuda_graph=False),
            ),
        )
    )
    ```

    `flow_shift=10.0` with `use_karras_sigmas=False` is the native linear-flow UniPC schedule; the example script uses the same values.
  </Step>

  <Step title="Tokenize the prompt">
    The scheduler never sees raw text. `Cosmos3Processor` applies the chat template, appends the `eos` and `<|vision_start|>` tokens, and returns token ids for the positive and negative prompts.

    ```python theme={null}
    from phyai_utils_tools.models.cosmos3 import Cosmos3Processor

    processor = Cosmos3Processor(
        f"{checkpoint_dir}/text_tokenizer",
        fps=24.0,
        num_frames=189,
        height=720,
        width=1280,
        append_metadata=True,
    )
    cond, uncond = processor.tokenize_pair(
        "A red sports car driving along a coastal road at sunset.",
        negative_prompt=None,
        device="cuda",
    )
    ```

    `negative_prompt=None` selects Cosmos3's built-in structured negative prompt. Pass `""` for an empty one.
  </Step>

  <Step title="Build the request">
    `Cosmos3T2VRequest` carries the tokenized conditions, the latent grid, and the sampler settings.

    | Field                            | Shape / Type            | Notes                                             |
    | -------------------------------- | ----------------------- | ------------------------------------------------- |
    | `text_ids` / `text_mask`         | `(1, S)` int64          | Positive prompt                                   |
    | `neg_text_ids` / `neg_text_mask` | `(1, S_neg)` int64      | Negative prompt                                   |
    | `video_shape`                    | `(t_lat, h_lat, w_lat)` | Latent grid, not pixels                           |
    | `fps`                            | `float`                 | Frame rate; also written into the prompt metadata |
    | `num_inference_steps`            | `int`                   | UniPC steps, `35` in the example                  |
    | `guidance_scale`                 | `float`                 | CFG scale, `6.0` in the example                   |
    | `seed`                           | `int`                   | Seed for the initial video and sound noise        |
    | `sound_frames`                   | `int` or `None`         | Anything other than `None` turns on T2AV          |

    ```python theme={null}
    import math

    from phyai.models.cosmos3 import Cosmos3T2VRequest, pixel_to_latent_shape

    num_frames = 189
    height = 720
    width = 1280
    fps = 24.0
    with_sound = False

    request = Cosmos3T2VRequest(
        text_ids=cond.text_ids,
        text_mask=cond.text_mask,
        neg_text_ids=uncond.text_ids,
        neg_text_mask=uncond.text_mask,
        video_shape=pixel_to_latent_shape(num_frames, height, width),
        fps=fps,
        num_inference_steps=35,
        guidance_scale=6.0,
        seed=42,
        sound_frames=(math.ceil(num_frames / fps * 25.0) if with_sound else None),
    )
    ```

    `pixel_to_latent_shape` divides the pixel dimensions by the VAE compression: `4` along time, `16` along each spatial axis.
  </Step>

  <Step title="Run generation">
    ```python theme={null}
    output = engine.step(request)
    ```

    T2V returns a pixel tensor shaped `(B, 3, T, H, W)` in `[0, 1]`. T2AV returns a dict:

    ```python theme={null}
    {
        "video": pixels,
        "sound": waveform,
        "sample_rate": sample_rate,
    }
    ```
  </Step>

  <Step title="Save media">
    The postprocessor moves everything to the CPU, converts the frames to uint8 RGB, and muxes the waveform into the same mp4 when there is one.

    ```python theme={null}
    from phyai_utils_tools.models.cosmos3 import Cosmos3GenerationPostProcessor

    postprocessor = Cosmos3GenerationPostProcessor(fps=fps)
    media = postprocessor.postprocess(output)
    postprocessor.save_mp4(media, ".cache/cosmos3_t2v.mp4")
    ```
  </Step>
</Steps>

# End-to-end examples

`examples/cosmos3/run_cosmos3.py` strings these steps together. A plain T2V run:

```bash theme={null}
uv run python examples/cosmos3/run_cosmos3.py \
    --checkpoint /path/to/Cosmos3-Nano \
    --prompt "A red sports car driving along a coastal road at sunset." \
    --out .cache/cosmos3_t2v
```

The same script spreads the run over eight GPUs once you add CFG and tensor parallelism. Eight ranks need eight visible devices, so pick them on the launching process:

```bash theme={null}
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 uv run python examples/cosmos3/run_cosmos3.py \
    --checkpoint /path/to/Cosmos3-Nano \
    --cfg 2 \
    --tp 4 \
    --out .cache/cosmos3_parallel
```

Add `--sound` for T2AV. The sound stream loads the AVAE and advances a second latent every step, so expect more memory and a longer run:

```bash theme={null}
uv run python examples/cosmos3/run_cosmos3.py \
    --checkpoint /path/to/Cosmos3-Nano \
    --prompt "ocean waves crashing on rocks" \
    --sound \
    --out .cache/cosmos3_t2av
```

The defaults are `720x1280`, `189` frames, and `35` steps, which takes a while. Shrink the run for a first smoke test:

```bash theme={null}
uv run python examples/cosmos3/run_cosmos3.py \
    --checkpoint /path/to/Cosmos3-Nano \
    --num-frames 49 \
    --height 480 \
    --width 832 \
    --steps 10 \
    --out .cache/cosmos3_smoke
```

The script prints timings for `model_load`, `preprocess`, `inference`, `to_cpu`, and `encode`. `inference` covers the denoising loop and the VAE decode; `encode` is PyAV writing the file.

# Multi-GPU execution

Two fields in `ParallelConfig` do the work. `cfg_size=2` runs the conditional and unconditional branches side by side on two rank groups instead of back to back on one GPU. `tp_size=N` shards the transformer inside each branch across `N` ranks. Every rank is one worker process, so a run needs `cfg_size * tp_size` GPUs, taken from the visible devices in order.

Compared with the inline snippet, little changes. `EngineConfig` gains a `parallel` block. Request tensors are built on the CPU, because the workers own the GPUs and the request has to cross a process boundary. And the engine moves behind an `if __name__ == "__main__":` guard: workers start with `spawn`, which re-imports the main module, and without the guard every worker would try to build an engine of its own. The `deployment` argument stays optional; here it only gives the workers longer to load the checkpoint.

```python theme={null}
import torch

from phyai import DeploymentConfig, Engine, EngineArgs
from phyai.engine_config import (
    AttentionParallelConfig,
    DenseParallelConfig,
    DeviceConfig,
    EngineConfig,
    OuterParallelConfig,
    ParallelConfig,
    RuntimeConfig,
)
from phyai.models.cosmos3 import Cosmos3T2VRequest, pixel_to_latent_shape
from phyai.models.cosmos3.main_cosmos3 import Cosmos3Args
from phyai.server import WorkerSupervisorConfig
from phyai_utils_tools.models.cosmos3 import (
    Cosmos3GenerationPostProcessor,
    Cosmos3Processor,
)

checkpoint_dir = "/path/to/Cosmos3-Nano"
cfg_size = 2  # cond and uncond branches on two rank groups
tp_size = 4  # transformer tensor parallelism inside each branch
num_frames = 189
height = 720
width = 1280
fps = 24.0


def main() -> None:
    engine = Engine(
        EngineArgs(
            plugin="cosmos3",
            plugin_args=Cosmos3Args(
                checkpoint_dir=checkpoint_dir,
                flow_shift=10.0,
                use_karras_sigmas=False,
            ),
            config=EngineConfig(
                device=DeviceConfig(target="cuda", params_dtype=torch.bfloat16),
                parallel=ParallelConfig(
                    outer=OuterParallelConfig(cfg_size=cfg_size),
                    dense=DenseParallelConfig(tp_size=tp_size),
                    attention=AttentionParallelConfig(tp_size=tp_size),
                ),
                runtime=RuntimeConfig(use_cuda_graph=False),
            ),
        ),
        # Optional. Without it the engine still spawns workers; this only
        # gives them longer to load the checkpoint.
        deployment=DeploymentConfig(
            process_config=WorkerSupervisorConfig(startup_timeout_s=1800.0),
        ),
    )
    assert engine.mode == "local"

    try:
        processor = Cosmos3Processor(
            f"{checkpoint_dir}/text_tokenizer",
            fps=fps,
            num_frames=num_frames,
            height=height,
            width=width,
            append_metadata=True,
        )
        # The workers own the GPUs, so request tensors stay on the CPU.
        cond, uncond = processor.tokenize_pair(
            "A red sports car driving along a coastal road at sunset.",
            negative_prompt=None,
            device="cpu",
        )
        request = Cosmos3T2VRequest(
            text_ids=cond.text_ids,
            text_mask=cond.text_mask,
            neg_text_ids=uncond.text_ids,
            neg_text_mask=uncond.text_mask,
            video_shape=pixel_to_latent_shape(num_frames, height, width),
            fps=fps,
            num_inference_steps=35,
            guidance_scale=6.0,
            seed=42,
        )

        output = engine.step(request)
        postprocessor = Cosmos3GenerationPostProcessor(fps=fps)
        media = postprocessor.postprocess(output)
        postprocessor.save_mp4(media, ".cache/cosmos3_t2v_parallel.mp4")
    finally:
        engine.close()


# Workers start with the "spawn" method and re-import this module, so the
# engine must sit behind the guard.
if __name__ == "__main__":
    main()
```

`engine.step()` hands back a CUDA-IPC view of the tensor on the output rank's GPU; the postprocessor copies it to the CPU, so saving works as before. Cosmos3 accepts `cfg_size` 1 or 2 and any `tp_size` that divides both the attention heads and the KV heads (1, 2, 4, or 8 for Cosmos3-Nano), with dense and attention TP kept equal. Other parallel axes are rejected before a worker starts. Serving replicas and `torchrun` launches are covered in [Parallel serving](/deployment/parallel-serving).

# Full example

```python theme={null}
import math

import torch

from phyai.engine import Engine, EngineArgs
from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
from phyai.models.cosmos3 import Cosmos3T2VRequest, pixel_to_latent_shape
from phyai.models.cosmos3.main_cosmos3 import Cosmos3Args
from phyai_utils_tools.models.cosmos3 import (
    Cosmos3GenerationPostProcessor,
    Cosmos3Processor,
)

checkpoint_dir = "/path/to/Cosmos3-Nano"
device = "cuda"
dtype = torch.bfloat16
num_frames = 189
height = 720
width = 1280
fps = 24.0
with_sound = False

engine = Engine(
    EngineArgs(
        plugin="cosmos3",
        plugin_args=Cosmos3Args(
            checkpoint_dir=checkpoint_dir,
            flow_shift=10.0,
            use_karras_sigmas=False,
            load_sound=(True if with_sound else None),
        ),
        config=EngineConfig(
            device=DeviceConfig(target=device, params_dtype=dtype),
            runtime=RuntimeConfig(use_cuda_graph=False),
        ),
    )
)

try:
    processor = Cosmos3Processor(
        f"{checkpoint_dir}/text_tokenizer",
        fps=fps,
        num_frames=num_frames,
        height=height,
        width=width,
        append_metadata=True,
    )
    cond, uncond = processor.tokenize_pair(
        "A red sports car driving along a coastal road at sunset.",
        negative_prompt=None,
        device=device,
    )

    request = Cosmos3T2VRequest(
        text_ids=cond.text_ids,
        text_mask=cond.text_mask,
        neg_text_ids=uncond.text_ids,
        neg_text_mask=uncond.text_mask,
        video_shape=pixel_to_latent_shape(num_frames, height, width),
        fps=fps,
        num_inference_steps=35,
        guidance_scale=6.0,
        seed=42,
        sound_frames=(math.ceil(num_frames / fps * 25.0) if with_sound else None),
    )

    output = engine.step(request)
    postprocessor = Cosmos3GenerationPostProcessor(fps=fps)
    media = postprocessor.postprocess(output)
    postprocessor.save_mp4(media, ".cache/cosmos3_t2v.mp4")
finally:
    engine.close()
```