scripts/dreambooth.sh
#!/bin/bash
set -ex

source .venv/bin/activate
which python
python -V

nvidia-smi

export CONCEPT="wolf_plushie"
export CLASS="stuffed animal"
export EXP="00031-667a-wolf_plushie"
export NPROMPTS=1000

export WANDB_MODE="offline"

echo "Started: $(date)"
accelerate launch diffusers/examples/dreambooth/train_dreambooth.py \
    --pretrained_model_name_or_path stabilityai/stable-diffusion-2-base \
    --instance_data_dir dreambooth/dataset/$CONCEPT \
    --output_dir experiments/dreambooth/$EXP \
    --instance_prompt="a photo of sks $CLASS" \
    --resolution=512 \
    --train_text_encoder \
    --max_train_steps=600 \
    --train_batch_size=1 \
    --learning_rate=2e-6 \
    --gradient_accumulation_steps=1 \
    --lr_scheduler="constant" \
    --lr_warmup_steps=0 \
    --checkpointing_steps=100 \
    --seed=0 \
    --report_to wandb
echo "Finished (train): $(date)"

python clean_exp_ckpts.py experiments/dreambooth/$EXP

python generate.py \
    --pretrained_model_or_path stabilityai/stable-diffusion-2-base \
    --exp_dir experiments/dreambooth/$EXP \
    --concept "sks $CLASS" \
    --prompt_source eval_set \
    --eval_set medium \
    --samples_per_prompt 10 \
    --prompts_per_batch 8
echo "Finished (genval): $(date)"

python generate.py \
    --pretrained_model_or_path stabilityai/stable-diffusion-2-base \
    --ckpt_path experiments/dreambooth/$EXP/checkpoint-400 \
    --concept "sks $CLASS" \
    --prompt_source json \
    --prompts_json_path data/$CLASS/$NPROMPTS-2-parts/0.json \
    --samples_per_prompt 10 \
    --prompts_per_batch 8
echo "Finished (genprompt): $(date)"
