Diffusion-DPO

To train a Stable Diffusion 1.5 model using Direct Preference Optimization (DPO), use the accelerate launch command pointing to train.py.

Note on Batch Size: The effective batch size is calculated as (N_GPU * train_batch_size * gradient_accumulation_steps). The original paper used an effective batch size of 2048.

Example Launch Script:

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATASET_NAME="yuvalkirstain/pickapic_v2"

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="/path/to/your/cache/" \
  --checkpointing_steps 500 \
  --beta_dpo 5000 \
   --output_dir="tmp-sd15"

The train.py script accepts several arguments to control the training process.

General Arguments

• --pretrained_model_name_or_path: The model to initialize from.
• --output_dir: Directory to save logs and checkpoints.
• --seed: Training seed.
• --sdxl: Flag to enable Stable Diffusion XL training.
• --sft: Flag to run Supervised Fine-Tuning (SFT) instead of DPO.

DPO Arguments

• --beta_dpo: The KL-divergence parameter beta for DPO.
• --choice_model: The model used for AI feedback (options include Aesthetics, CLIP, PickScore, or HPS).

Optimizer and Learning Rate Arguments

• --max_train_steps: Total number of training steps.
• --gradient_accumulation_steps: Number of steps for gradient accumulation.
• --train_batch_size: Batch size per device.
• --checkpointing_steps: Frequency of saving model checkpoints.
• --gradient_checkpointing: Automatically enabled for SDXL.
• --learning_rate: The learning rate.
• --scale_lr: Scales the learning rate (recommended for stability).
• --lr_scheduler: Type of scheduler (e.g., linear_warmup_to_constant).
• --lr_warmup_steps: Number of warmup steps.
• --use_adafactor: Uses Adafactor instead of Adam (lower memory, default for SDXL).

Data Arguments

• --dataset_name: Name of the dataset (e.g., Pick-a-Pic).
• --cache_dir: Local directory for caching the dataset.
• --resolution: Image resolution (defaults to 512 for SD1.5, 1024 for SDXL).
• --random_crop: Enables random cropping data augmentation.
• --no_hflip: Disables horizontal flipping data augmentation.
• --dataloader_num_workers: Number of dataloader workers.

To compare a standard Stable Diffusion model with a DPO-tuned version, you can swap the unet attribute of an existing pipeline. This allows you to use the same pipeline infrastructure (scheduler, text encoder, etc.) while switching between the original and the DPO-tuned weights.

# Assuming 'pipe' is your loaded pipeline and 'unets' is a list [original_unet, dpo_unet]
unets = [pipe.unet, dpo_unet]
names = ["Orig. SDXL", "DPO SDXL"]

def gen(prompt, seed=0, run_baseline=True):
    ims = []
    generator = torch.Generator(device='cuda')
    # Loop through original and DPO unets if run_baseline is True
    for unet_i in ([0, 1] if run_baseline else [1]):
        print(f"Prompt: {prompt}\nSeed: {seed}\n{names[unet_i]}")
        pipe.unet = unets[unet_i]
        generator = generator.manual_seed(seed)
        
        im = pipe(prompt=prompt, generator=generator, guidance_scale=7.5).images[0]
        ims.append(im)
    return ims

Depending on whether you are using Stable Diffusion v1.5 or SDXL, use StableDiffusionPipeline or StableDiffusionXLPipeline. For SDXL, it is recommended to use variant="fp16" and use_safetensors=True.

Note: You may want to set pipe.safety_checker = None to prevent the safety checker from aggressively filtering generated images.

from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
import torch

pretrained_model_name = "stabilityai/stable-diffusion-xl-base-1.0"

if 'stable-diffusion-xl' in pretrained_model_name:
    pipe = StableDiffusionXLPipeline.from_pretrained(
        pretrained_model_name, 
        torch_dtype=torch.float16,
        variant="fp16", 
        use_safetensors=True
    ).to("cuda")
else:
    pipe = StableDiffusionPipeline.from_pretrained(
        pretrained_model_name, 
        torch_dtype=torch.float16
    ).to("cuda")

pipe.safety_checker = None