# Note: This will only work on Navi21 GPUs (6800/6900+).
# See: https://github.com/RadeonOpenCompute/ROCm/issues/1668#issuecomment-1043994570

# Install Conda (latest from https://docs.conda.io/en/latest/miniconda.html#linux-installers)
wget https://repo.anaconda.com/miniconda/Miniconda3-py39_4.12.0-Linux-x86_64.sh
bash Miniconda3-py39_4.12.0-Linux-x86_64.sh
# follow the prompts to install it, and run `conda` to make sure it's working.

# Install git and curl, and clone the stable-diffusion repo
sudo apt install -y git curl
cd Downloads
git clone https://github.com/CompVis/stable-diffusion.git

# Install dependencies and activate environment
cd stable-diffusion
conda env create -f environment.yaml
conda activate ldm

# Download Stable Diffusion weights
curl https://www.googleapis.com/storage/v1/b/aai-blog-files/o/sd-v1-4.ckpt?alt=media > sd-v1-4.ckpt

# Symlink the weights into place
mkdir -p models/ldm/stable-diffusion-v1/
ln -s -r sd-v1-4.ckpt models/ldm/stable-diffusion-v1/model.ckpt

# Install AMD ROCm support
wget https://repo.radeon.com/amdgpu-install/22.10/ubuntu/focal/amdgpu-install_22.10.50100-1_all.deb
sudo apt-get install ./amdgpu-install_22.10.50100-1_all.deb
sudo amdgpu-install --usecase=dkms,graphics,rocm,lrt,hip,hiplibsdk
# make sure you see your GPU by running rocm-smi

# Make AMD GPU work with ROCm
cd stable-diffusion/
conda remove cudatoolkit -y
pip3 uninstall torch torchvision -y 
# Install PyTorch ROCm
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/rocm5.1.0
pip3 install transformers==4.19.2 scann kornia==0.6.4 torchmetrics==0.6.0

# Generate an image
python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms