Cosmos3 Generation Mode

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

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.

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.

Architecture

The path follows PhyAI’s usual engine + plugin contract. Inside the plugin the work splits into a handful of pieces:

phyai/src/phyai/models/cosmos3

main_cosmos3.py

scheduler_cosmos3.py

model_runner_cosmos3.py

model_runner_vae_cosmos3.py

modeling_cosmos3.py

vae_wan.py

avae_sound.py

sampler_unipc.py

configuration_cosmos3.py

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

1. Prepare weights

   Download Cosmos3-Nano. The examples expect this layout:

   /path/to/Cosmos3-Nano/
     transformer/
     vae/
     text_tokenizer/
     sound_tokenizer/   # required for T2AV
     scheduler/
   
2. 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.

   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.

3. 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.

   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.

4. 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

   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.

5. Run generation

   output = engine.step(request)
   

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

   {
       "video": pixels,
       "sound": waveform,
       "sample_rate": sample_rate,
   }
   
6. 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.

   from phyai_utils_tools.models.cosmos3 import Cosmos3GenerationPostProcessor
   
   postprocessor = Cosmos3GenerationPostProcessor(fps=fps)
   media = postprocessor.postprocess(output)
   postprocessor.save_mp4(media, ".cache/cosmos3_t2v.mp4")
   

End-to-end examples

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

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:

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:

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:

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.

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.

Full example

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()