#!/bin/bash
export HF_HOME="/mnt/data/cache" 
export MODEL_DIR="/mnt/data/StableDiffusion/stable-diffusion-v1-5"
export OUTPUT_DIR="/mnt/data/output_dir"

DATASET_NAME="lvis"
CONDITIONING_MODE="segmentation"

VALIDATION_IMAGE_1='./condition_val/epsilon_0_blackbg/000000286994.jpg'
VALIDATION_PROMPT_1='some elephants and one is by some water'

VALIDATION_IMAGE_2='./condition_val/dilate_20_within_bbox/000000286994.png'
VALIDATION_PROMPT_2='some elephants and one is by some water'

VALIDATION_IMAGE_3='./condition_val/dilate_40_within_bbox/000000286994.png'
VALIDATION_PROMPT_3='some elephants and one is by some water'


accelerate config default 
python train_controlnet.py \
  --pretrained_model_name_or_path=$MODEL_DIR \
  --output_dir=$OUTPUT_DIR \
  --dataset_name=$DATASET_NAME \
  --conditioning_mode=$CONDITIONING_MODE \
  --dilate_radius="random" \
  --resolution=512 \
  --learning_rate=1e-5 \
  --dataloader_num_workers=24 \
  --checkpointing_steps=2000 \
  --validation_steps=200 \
	--validation_image $VALIDATION_IMAGE_1 $VALIDATION_IMAGE_2 $VALIDATION_IMAGE_3 \
  --validation_prompt "${VALIDATION_PROMPT_1}" "${VALIDATION_PROMPT_2}" "${VALIDATION_PROMPT_3}" \
  --train_batch_size=24 \
	--num_train_epochs=10 \
  --proportion_empty_prompts=0.5 \
  --do_predict \
  --do_ratio_condition \
  --detach_feature \
  --low_cpu_mem_usage \
  --controlnet_model_name_or_path=$CONTROLNET_DIR
