#!/usr/bin/env sh
#license: MIT

# amd warning: optimizer AdamW8bit doesnt work on amd
# more info: https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md



# more info:
# - https://rentry.org/lora_train
# - https://github.com/darkstorm2150/sd-scripts#links-to-usage-documentation
# - https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_README-en.md
# if you have amd, remove xformers from lora arguments

export SD_MODELS=/home/tekakutli/files/models/stable-diffusion/models/
export LORAS=/home/tekakutli/files/models/stable-diffusion/lora/
#git clone https://github.com/kohya-ss/sd-scripts
export LORA_REPO=/home/tekakutli/code/stable-diffusion/sd-scripts



model=dreamlikeDiffusion10_10.ckpt
config="/home/tekakutli/files/org/roam/neuralnomicon/gists/train.toml"


# takes one optional input: an output name
lora(){

     name=last
     if ! [[ -z "$1" ]]; then
           name="$1"
     fi
     cd $LORA_REPO
     # not using accelerate
     # accelerate launch --num_cpu_threads_per_process 1 train_network.py \
     python train_network.py \
     --pretrained_model_name_or_path="$SD_MODELS"$model \
     --dataset_config="$config"\
     --output_dir=$LORAS \
     --output_name="$name" \
     --save_model_as=safetensors \
     --prior_loss_weight=1.0 \
     --max_train_epochs=1 \
     --learning_rate=1e-4 \
     --optimizer_type="AdamW" \
     --mixed_precision="fp16" \
     --gradient_checkpointing \
     --save_every_n_steps=100 \
     --network_module=networks.lora
}
