Awesome

Intro

This is the training code for Diffusion-DPO. The script is adapted from the diffusers library.

Model Checkpoints

The below are initialized with StableDiffusion models and trained as described in the paper (replicable with launchers/ scripts assuming 16 GPUs, scale gradient accumulation accordingly).

StableDiffusion1.5

StableDiffusion-XL-1.0

Use this notebook to compare generations. It also has a sample of automatic quantative evaluation using PickScore.

Setup

pip install -r requirements.txt

Structure

launchers/ is examples of running SD1.5 or SDXL training
utils/ has the scoring models for evaluation or AI feedback (PickScore, HPS, Aesthetics, CLIP)
quick_samples.ipynb is visualizations from a pretrained model vs baseline
requirements.txt Basic pip requirements
train.py Main script, this is pretty bulky at >1000 lines, training loop starts at ~L1000 at this commit (ctrl-F "for epoch").
upload_model_to_hub.py Uploads a model checkpoint to HF (simple utility, current values are placeholder)

Running the training

Example SD1.5 launch

# from launchers/sd15.sh
export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATASET_NAME="yuvalkirstain/pickapic_v2"

# Effective BS will be (N_GPU * train_batch_size * gradient_accumulation_steps)
# Paper used 2048. Training takes ~24 hours / 2000 steps

accelerate launch --mixed_precision="fp16"  train.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --train_batch_size=1 \
  --dataloader_num_workers=16 \
  --gradient_accumulation_steps=1 \
  --max_train_steps=2000 \
  --lr_scheduler="constant_with_warmup" --lr_warmup_steps=500 \
  --learning_rate=1e-8 --scale_lr \
  --cache_dir="/export/share/datasets/vision_language/pick_a_pic_v2/" \
  --checkpointing_steps 500 \
  --beta_dpo 5000 \
   --output_dir="tmp-sd15"

Important Args

General

--pretrained_model_name_or_path what model to train/initalize from
--output_dir where to save/log to
--seed training seed (not set by default)
--sdxl run SDXL training
--sft run SFT instead of DPO

DPO

--beta_dpo KL-divergence parameter beta for DPO
--choice_model Model for AI feedback (Aesthetics, CLIP, PickScore, HPS)

Optimizers/learning rates

--max_train_steps How many train steps to take
--gradient_accumulation_steps
--train_batch_size see above notes in script for actual BS
--checkpointing_steps how often to save model
--gradient_checkpointing turned on automatically for SDXL
--learning_rate
--scale_lr Found this to be very helpful but isn't default in code
--lr_scheduler Type of LR warmup/decay. Default is linear warmup to constant
--lr_warmup_steps number of scheduler warmup steps
--use_adafactor Adafactor over Adam (lower memory, default for SDXL)

Data

--dataset_name if you want to switch from Pick-a-Pic
--cache_dir where dataset is cached locally (users will want to change this to fit their file system)
--resolution defaults to 512 for non-SDXL, 1024 for SDXL.
--random_crop and --no_hflip changes data aug
--dataloader_num_workers number of total dataloader workers

Citation

@misc{wallace2023diffusion,
      title={Diffusion Model Alignment Using Direct Preference Optimization}, 
      author={Bram Wallace and Meihua Dang and Rafael Rafailov and Linqi Zhou and Aaron Lou and Senthil Purushwalkam and Stefano Ermon and Caiming Xiong and Shafiq Joty and Nikhil Naik},
      year={2023},
      eprint={2311.12908},
      archivePrefix={arXiv},
      primaryClass={cs.CV}