#!/bin/bash

# Configuration
install_dir=~/  # Define the installation directory
do_aws_setup=true  # Flag to control AWS setup

# Create necessary directories
mkdir -p "$install_dir/models"
mkdir -p "$install_dir/datasets"

# Clone and set up naifu-diffusion
cd "$install_dir" || exit
git clone https://github.com/Mikubill/naifu-diffusion nd && cd nd
git checkout sgm
pip install -r requirements.txt xformers prodigyopt

# Download parameter calculator
param_calc="https://huggingface.co/trojblue/nd-param-calculator/resolve/main/nd_param_calculator_latest.ipynb"
wget "$param_calc"

# AWS setup
if [ "$do_aws_setup" = true ]; then
    echo "Setting up AWS..."
    # Check if aws command exists
    if ! command -v aws &> /dev/null; then
        echo "Installing AWS CLI..."
        curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
        sudo apt install unzip && unzip awscliv2.zip
        sudo ./aws/install
        rm awscliv2.zip
    else
        echo "AWS CLI is already installed."
    fi

    # AWS S3 copy command
    cd "$install_dir/models" || exit
    aws s3 cp s3://bucket-external/trainer_opt/sdxl-fulldan-run-continue/checkpoint/nd-epoch-09.ckpt fd5me9.ckpt
else
    echo "Skipping AWS setup..."
fi

# Add any additional setup or installation commands below