Installation

There are two options for installing Cosmos-Predict2: Using a Conda environment or using a Docker container.

Note

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

Option 1: Conda Environment Setup

1. Clone the Cosmos-Predict2 repository from GitHub.

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

2. Create and activate the conda environment.

   conda env create --file cosmos-predict2.yaml
   conda activate cosmos-predict2
   

3. Install Cosmos-Predict2 dependencies for the environment.

   pip install -r requirements-conda.txt
   pip install flash-attn==2.6.3 --no-build-isolation
   
   # Patch Transformer engine linking issues
   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
   pip install transformer-engine[pytorch]==1.13.0
   

4. Install the Apex library for training. This step is optional if you are using Cosmos-Predict2 for inference only.

   pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext --cuda_ext" git+https://github.com/NVIDIA/apex.git
   

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

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

6. Ensure the CUDA_HOME environment variable points to your Conda installation directory.

   export CUDA_HOME=$CONDA_PREFIX
   

Option 2: Docker Container Setup

1. Set up the Docker container by either pulling the pre-built container or building the container using the provided Dockerfile.

   1. Pull the pre-build container:

      docker pull nvcr.io/nvidia/cosmos/cosmos-predict2-container:1.0
      

   2. Build the container from the Dockerfile (ensure you run this command from the root directory of the repository):

      docker build -t cosmos-predict2-local -f Dockerfile .
      

2. Run the Docker container:

   Note

   Replace [CONTAINER_NAME] with either nvcr.io/nvidia/cosmos/cosmos-predict2-container:1.0 or cosmos-predict2-local.

   Note

   Replace /path/to/cosmos-predict2, /path/to/datasets, and /path/to/checkpoints with your actual local paths.

   docker run --gpus all -it --rm \
   -v /path/to/cosmos-predict2:/workspace \
   -v /path/to/datasets:/workspace/datasets \
   -v /path/to/checkpoints:/workspace/checkpoints \
   [CONTAINER_NAME]
   

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

   python /workspace/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.