# Nvidia Cosmos (Compatible with Blackwell)

### 1. Install wsl (this also install last version wsl and ubuntu 24.04.1)

wsl --install

### Update wsl (Optional)

wsl.exe --update

### (Optional). Install ubuntu 24.04 (install mannually)
### 2. Install CUDA

sudo apt-get remove --purge '^nvidia-.*'
sudo apt-get --purge remove nvidia*
sudo apt-get --purge remove "*cublas*" "cuda*"
sudo apt remove nvidia-*
sudo apt-get --purge remove "*nvidia*"
sudo apt autoremove -y
sudo apt-get update && sudo apt-get full-upgrade
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cudnn tensorrt
rm -rf cuda-keyring_1.1-1_all.deb 

sudo apt-key del 7fa2af80
sudo apt-get update && sudo apt-get full-upgrade
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-8
# sudo apt-get -y install nvidia-cuda-toolkit
rm -rf cuda-keyring_1.1-1_all.deb

#### Add in your path

nano ~/.bashrc

# CUDA Toolkit 12.8
export PATH=/usr/local/cuda-12.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:$LD_LIBRARY_PATH

source ~/.bashrc

### 3. Install Miniconda

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
chmod +x Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

source ~/.bashrc

conda update --all --yes

### 4. Clone Cosmos

mkdir Projects
cd Projects
git clone --recursive https://github.com/NVIDIA/Cosmos.git

cd ~/Projects/Cosmos

### 5. Install Cosmos requirements

sudo apt-get install build-essential gcc cmake ffmpeg ninja-build

### Create conda environent

conda env remove -n cosmos
conda create -n cosmos python=3.12
conda activate cosmos

### Install Requirements (Only ready for 3090 and 4090)

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

pip3 install -U einops attrs pynvml protobuf opencv-python

export MAX_JOBS=4
cd ~/Projects/
git clone --branch release_v2.0 --recursive https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine
export NVTE_FRAMEWORK=pytorch   # Optionally set framework
pip install . -v                  # Build and install

cd ~/Projects/Cosmos
pip3 install -U -r requirements.txt

### 6. Login Hugging face and accept all models

Generate a [Hugging Face](https://huggingface.co/settings/tokens) access token. Set the access token to 'Read' permission (default is 'Fine-grained').

huggingface-cli login

### 7. Download Models

## Download only 7B models

PYTHONPATH=$(pwd) python3 cosmos1/scripts/download_diffusion.py --model_sizes 7B --model_types Text2World Video2World

## If you have 2x24GB or 48GB download only all models

PYTHONPATH=$(pwd) python3 cosmos1/scripts/download_diffusion.py --model_sizes 7B 14B --model_types Text2World Video2World

### 8. Run Demo

PROMPT="A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves. \
The robot's metallic body gleams under the bright, even lighting, highlighting its futuristic design and intricate joints. \
A glowing blue light emanates from its chest, adding a touch of advanced technology. The background is dominated by rows of boxes, \
suggesting a highly organized storage system. The floor is lined with wooden pallets, enhancing the industrial setting. \
The camera remains static, capturing the robot's poised stance amidst the orderly environment, with a shallow depth of \
field that keeps the focus on the robot while subtly blurring the background for a cinematic effect."

PYTHONPATH=$(pwd) python3 cosmos1/models/diffusion/inference/text2world.py \
    --checkpoint_dir checkpoints \
    --diffusion_transformer_dir Cosmos-1.0-Diffusion-7B-Text2World \
    --prompt "$PROMPT" \
    --video_save_name Cosmos-1.0-Diffusion-7B-Text2World_memory_efficient \
    --offload_tokenizer \
    --offload_diffusion_transformer \
    --offload_text_encoder_model \
    --offload_prompt_upsampler \
    --offload_guardrail_models

### 8. See gpu performance

pip3 install nvitop
nvitop