| # Setup Guide |
|
|
| ## System Requirements |
|
|
| * NVIDIA GPUs with Ampere architecture (RTX 30 Series, A100) or newer |
| * NVIDIA driver compatible with CUDA 12.6 |
| * Linux x86-64 |
| * glibc>=2.31 (e.g Ubuntu >=22.04) |
| * Python 3.10 |
|
|
| ## Installation |
|
|
| ### Clone the repository |
|
|
| ```bash |
| git clone git@github.com:nvidia-cosmos/cosmos-predict2.git |
| cd cosmos-predict2 |
| ``` |
|
|
| ### Option 1: Virtual environment |
|
|
| Install system dependencies: |
|
|
| [uv](https://docs.astral.sh/uv/getting-started/installation/) |
|
|
| ```shell |
| curl -LsSf https://astral.sh/uv/install.sh | sh |
| source $HOME/.local/bin/env |
| ``` |
|
|
| Install the package into a new environment: |
|
|
| ```shell |
| uv sync --extra cu126 |
| source .venv/bin/activate |
| ``` |
|
|
| Or, install the package into the active environment (e.g. conda): |
|
|
| ```shell |
| uv sync --extra cu126 --active --inexact |
| ``` |
|
|
| ### Option 2: Docker container |
|
|
| Please make sure you have access to Docker on your machine and the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) is installed. |
|
|
| Build and run the container: |
|
|
| ```bash |
| docker run --gpus all --rm -v .:/workspace -v /workspace/.venv -it $(docker build -q .) |
| ``` |
|
|
| ## Downloading Checkpoints |
|
|
| 1. Get a [Hugging Face Access Token](https://huggingface.co/settings/tokens) with `Read` permission |
| 2. Install [Hugging Face CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli): `uv tool install -U "huggingface_hub[cli]"` |
| 3. Login: `hf auth login` |
| 4. Accept the [Llama-Guard-3-8B terms](https://huggingface.co/meta-llama/Llama-Guard-3-8B). |
|
|
| To download a specific model: |
|
|
| ```shell |
| ./scripts/download_checkpoints.py --model_types <model_type> --model_sizes <model_size> |
| ``` |
|
|
| | Models | Link | Download Arguments | Notes | |
| |--------|------|--------------------|-------| |
| | Cosmos-Predict2-0.6B-Text2Image | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-0.6B-Text2Image) | `--model_types text2image --model_sizes 0.6B` | N/A | |
| | Cosmos-Predict2-2B-Text2Image | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image) | `--model_types text2image --model_sizes 2B` | N/A | |
| | Cosmos-Predict2-14B-Text2Image | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Text2Image) | `--model_types text2image --model_sizes 14B` | N/A | |
| | Cosmos-Predict2-2B-Video2World | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) | `--model_types video2world --model_sizes 2B` | Download 720P, 16FPS by default. Supports 480P and 720P resolution. Supports 10FPS and 16FPS | |
| | Cosmos-Predict2-14B-Video2World | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Video2World) | `--model_types video2world --model_sizes 14B` | Download 720P, 16FPS by default. Supports 480P and 720P resolution. Supports 10FPS and 16FPS | |
| | Cosmos-Predict2-2B-Sample-Action-Conditioned | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Sample-Action-Conditioned) | `--model_types sample_action_conditioned` | Supports 480P and 4FPS. | |
| | Cosmos-Predict2-14B-Sample-GR00T-Dreams-GR1 | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Sample-GR00T-Dreams-GR1) | `--model_types sample_gr00t_dreams_gr1` | Supports 480P and 16FPS. | |
| | Cosmos-Predict2-14B-Sample-GR00T-Dreams-DROID | [🤗 Huggingface](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Sample-GR00T-Dreams-DROID) | `--model_types sample_gr00t_dreams_droid` | Supports 480P and 16FPS. | |
|
|
| To download all the checkpoints (requires ~250GB of disk space), run: |
|
|
| ```shell |
| ./scripts/download_checkpoints.py |
| ``` |
|
|
| To see the full list of options, run: |
|
|
| ```shell |
| ./scripts/download_checkpoints.py --help |
| ``` |
|
|
| ## Troubleshooting |
|
|
| ### CUDA/GPU Issues |
|
|
| * **CUDA driver version insufficient**: Update NVIDIA drivers to latest version compatible with CUDA 12.6+ |
|
|
| * **Out of Memory (OOM) errors**: Use 2B models instead of 14B, or reduce batch size/resolution |
| * **Missing CUDA libraries**: Set paths with `export CUDA_HOME=$CONDA_PREFIX` |
|
|
| ### Installation Issues |
|
|
| * **Conda environment conflicts**: Create fresh environment with `conda create -n cosmos-predict2-clean python=3.10 -y` |
|
|
| * **Flash-attention build failures**: Install build tools with `apt-get install build-essential` |
| * **Transformer engine linking errors**: Reinstall with `pip install --force-reinstall transformer-engine==1.12.0` |
|
|
| For other issues, check [GitHub Issues](https://github.com/nvidia-cosmos/cosmos-predict2/issues). |
|
|
