#!/bin/bash
export EXPERIMENT_NAME=103
export MODEL_PATH="stabilityai/stable-diffusion-2-1-base"
export CLASS_DIR="/home/anonymous/AIGC/diffusers/Style/data/CelebA-HQ/reference"
export INSTANCE_DIR="data/CelebA-HQ/$EXPERIMENT_NAME/all"
export DREAMBOOTH_OUTPUT_DIR="dreambooth-outputs/clean/CelebA-HQ/$EXPERIMENT_NAME"
export CUDA_VISIBLE_DEVICES=1,2,3,4

CUDA_VISIBLE_DEVICES=1,2,3,4 python train_dreambooth.py \
    --pretrained_model_name_or_path=$MODEL_PATH  \
    --enable_xformers_memory_efficient_attention \
    --train_text_encoder \
    --instance_data_dir=$INSTANCE_DIR \
    --class_data_dir=$CLASS_DIR \
    --output_dir=$DREAMBOOTH_OUTPUT_DIR \
    --with_prior_preservation \
    --prior_loss_weight=1.0 \
    --class_prompt="a photo of person" \
    --instance_prompt="a photo of sks person" \
    --class_prompt="a photo of person" \
    --inference_prompt="a photo of sks person" \
    --resolution=512 \
    --train_batch_size=2 \
    --gradient_accumulation_steps=1 \
    --learning_rate=5e-7 \
    --lr_scheduler="constant" \
    --lr_warmup_steps=0 \
    --num_class_images=16 \
    --max_train_steps=1000 \
    --checkpointing_steps=1000 \
    --center_crop \
    --mixed_precision=bf16 \
    --prior_generation_precision=bf16 \
    --sample_batch_size=1 \
    --seed=0
  
python infer.py \
      --model_path /home/anonymous/AIGC/diffusers/StyleGuard/dreambooth-outputs/clean/CelebA-HQ/103/checkpoint-1000 \
      --output_dir /home/anonymous/AIGC/diffusers/StyleGuard/dreambooth-outputs/clean/CelebA-HQ/103/checkpoint-1000-test-infer

