#!/bin/bash

set -e
set -o nounset
set -o pipefail

apt update
apt install -y less vim unzip zip

cd /workspace

echo "Installing dependencies"

pip install --upgrade \
    torch \
    torchaudio \
    torchvision \
    --index-url https://download.pytorch.org/whl/cu118

echo "Installing xformers"

pip install --no-deps xformers==0.0.17

echo "Installing TensorRT"

pip install tensorrt

echo "Cloning the project..."

git clone https://github.com/bmaltais/kohya_ss/

echo "Installing the project..."

DIR="/workspace/kohya_ss"
TEMP_REQUIREMENTS_FILE="${DIR}/requirements_tmp_for_setup.txt"
awk -v dir="$DIR" '/#.*kohya_ss.*library/{print; getline; sub(/^\.$/, dir)}1' "$DIR/requirements_unix.txt" >"$TEMP_REQUIREMENTS_FILE"
python -m pip install --use-pep517 --upgrade -r "$TEMP_REQUIREMENTS_FILE"

echo "Installing accelerate config"

mkdir -p "$HOME/.cache/huggingface/accelerate"
cp "${DIR}/config_files/accelerate/default_config.yaml"\
    "$HOME/.cache/huggingface/accelerate/default_config.yaml"

VENV_DIR=$(python -c "import site; print(site.getsitepackages()[0])")
VENV_DIR="${VENV_DIR%/lib/python3.10/site-packages}"

echo "Downloading base model"
pushd /workspace
curl -L -o model.safetensors -J https://civitai.com/api/download/models/28562
popd

echo "Creating directories"

mkdir -p /workspace/images/train/img \
    /workspace/images/train/model \
    /workspace/images/train/log

echo "Cleanup"
rm -r ~/.cache/huggingface ~/.cache/pip*

print "Setup complete."
exit 0