Installation

Follow these steps to clone the Cosmos Predict2 repository from GitHub and run the Docker container.

Note

Ensure you have the necessary hardware and software prerequisites before installation.

Cosmos-Predict2 Setup

1. Clone the Cosmos-Predict2 repository from GitHub.

   git clone git@github.com:nvidia-cosmos/cosmos-predict2.git
   cd cosmos-predict2
   

2. Set up the conda environment and install Cosmos-Predict2 dependencies for the environment.

   # Create the cosmos-predict2 conda environment.
   conda env create --file cosmos-predict2.yaml
   # Activate the cosmos-predict2 conda environment.
   conda activate cosmos-predict2
   # Install the dependencies.
   pip install -r requirements.txt
   pip install flash-attn==2.7.4.post1 --no-build-isolation
   # Patch Transformer engine linking issues in conda environments.
   ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/
   ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10
   # Install Transformer engine.
   pip install transformer-engine[pytorch]==1.12.0
   

3. Test the environment using the test_environment.py script.

   CUDA_HOME=$CONDA_PREFIX PYTHONPATH=$(pwd) python scripts/test_environment.py
   

Next Steps

If you haven’t done so already, determine which model you want to work with–refer to the Model Matrix page for more information.

You can then move on to the Predict2 Quickstart Guide.