export RUN_NAME="gift_immunized"
export WANDB_NAME="${RUN_NAME}-Attack"

# export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR='results/immunization'
export CKPT="${OUTPUT_DIR}/${RUN_NAME}.pt"

export INSTANCE_PROMPT="A photo of a penguin plushie"
export VALIDATION_PROMPT="A penguin plushie riding a bicycle in front of Eiffel Tower"
export INSTANCE_DIR="data/malicious/plushie_penguin/test"
export OUTPUT_DIR_FT="results/${RUN_NAME}-Attack"

nohup accelerate launch --num_processes=1 --gpu_ids='0' train/train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR_FT \
  --instance_prompt="${INSTANCE_PROMPT}" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=2500 \
  --validation_prompt="${VALIDATION_PROMPT}" \
  --validation_steps=1 \
  --seed="0" \
  --num_validation_images=4 \
  --use_8bit_adam \
  --gradient_checkpointing \
  --gift_ckpt=$CKPT \
#   --saved_pipeline_name=$RUN_NAME \