Diffusion-DPO 项目技术文档

1. 安装指南

环境准备

在开始安装之前，请确保您的系统满足以下要求：

• Python 3.7 或更高版本
• CUDA 11.0 或更高版本（如果使用GPU）

安装步骤

1. 克隆项目仓库：

   git clone https://github.com/your-repo/diffusion-dpo.git
   cd diffusion-dpo
   

2. 安装依赖：

   pip install -r requirements.txt
   

3. 配置环境变量（可选）： 如果需要使用GPU进行训练，请确保已正确配置CUDA环境变量。

2. 项目的使用说明

项目概述

Diffusion-DPO 是一个基于扩散模型的文本到图像生成项目。该项目使用 Stable Diffusion 模型进行初始化，并通过 Direct Preference Optimization (DPO) 方法进行训练。

主要功能

• 文本到图像生成：根据输入的文本描述生成相应的图像。
• 模型训练：支持对 Stable Diffusion 1.5 和 Stable Diffusion XL 1.0 进行微调。
• 自动量化评估：使用 PickScore 等模型进行自动量化评估。

使用流程

1. 数据准备：准备训练数据集，并将其存储在本地缓存目录中。
2. 模型训练：使用 train.py 脚本进行模型训练。
3. 模型评估：使用 quick_samples.ipynb 进行模型生成结果的可视化和评估。

3. 项目API使用文档

主要API

• train.py：主训练脚本，支持多种参数配置。
• upload_model_to_hub.py：将训练好的模型上传到 Hugging Face Hub。

常用参数说明

• --pretrained_model_name_or_path：预训练模型的路径或名称。
• --output_dir：模型输出目录。
• --seed：训练种子。
• --sdxl：是否使用 SDXL 模型进行训练。
• --beta_dpo：DPO 方法中的 KL-divergence 参数 beta。
• --max_train_steps：最大训练步数。
• --learning_rate：学习率。
• --dataset_name：数据集名称。
• --cache_dir：数据集缓存目录。

4. 项目安装方式

安装步骤

1. 克隆项目仓库：

   git clone https://github.com/your-repo/diffusion-dpo.git
   cd diffusion-dpo
   

2. 安装依赖：

   pip install -r requirements.txt
   

3. 配置环境变量（可选）： 如果需要使用GPU进行训练，请确保已正确配置CUDA环境变量。

运行示例

以下是一个运行 Stable Diffusion 1.5 训练的示例：

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="/export/share/datasets/vision_language/pick_a_pic_v2/" \
  --checkpointing_steps 500 \
  --beta_dpo 5000 \
  --output_dir="tmp-sd15"

通过以上步骤，您可以成功安装并运行 Diffusion-DPO 项目。