ImageReward

📃 论文 • 🖼 数据集 • 🌐 中文博客 • 🤗 HF 仓库 • 🐦 Twitter

🔥🔥 新闻！ 2024/12/31: 我们发布了下一代模型，VisionReward，这是一种用于视觉生成（文本到图像 / 文本到视频）稳定 RLHF 的细粒度、多维度奖励模型！

🔥 新闻！ 2023/9/22: ImageReward 的论文已被 NeurIPS 2023 接收！

ImageReward：学习与评估人类对文本到图像生成的偏好

ImageReward 是首个通用的文本到图像人类偏好 RM，共基于13.7万组专家比较数据进行训练，在理解文本到图像合成中的人类偏好方面，优于现有的文本-图像评分方法，例如 CLIP（高出 38.6%）、Aesthetic（高出 39.6%）和 BLIP（高出 31.6%）。

此外，我们还引入了奖励反馈学习（ReFL），用于直接使用 ImageReward 优化文本到图像扩散模型。经过 ReFL 微调的 Stable Diffusion 在人工评估中比未微调版本高出 58.4%。

目前，ImageReward 和 ReFL 已全部打包为 Python image-reward 包！

只需三行代码，即可使用 image-reward 包进行 ImageReward 评分！

# pip install image-reward
import ImageReward as RM
model = RM.load("ImageReward-v1.0")

rewards = model.score("<prompt>", ["<img1_obj_or_path>", "<img2_obj_or_path>", ...])

只需四行代码，即可使用 image-reward 包进行 ReFL 微调！

# pip install image-reward
# pip install diffusers==0.16.0 accelerate==0.16.0 datasets==2.11.0
from ImageReward import ReFL
args = ReFL.parse_args()
trainer = ReFL.Trainer("CompVis/stable-diffusion-v1-4", "data/refl_data.json", args=args)
trainer.train(args=args)

如果您觉得 ImageReward 的开源工作很有帮助，请为我们点个星🌟，以鼓励我们继续开发！

• ImageReward
  • 快速入门
    • 安装依赖
    • 示例用法
  • ReFL
    • 安装依赖
    • 示例用法
  • ImageReward 和 ReFL 的演示
  • ImageReward 的训练代码
  • 集成到 Stable Diffusion Web UI
    • 功能
      • 为生成的图片打分并附加到图片信息中
        • 使用方法
        • 演示视频
      • 自动过滤低分图片
        • 使用方法
        • 演示视频
      • 查看已评分图片的分数
        • 使用方法
        • 示例
      • 其他功能
        • 内存管理
    • 常见问题解答
  • 重现表 1 中的实验
  • 重现表 3 中的实验
  • 引用

快速入门

安装依赖

我们将整个仓库整合为一个 Python 包 image-reward。请按照以下命令准备环境：

# 克隆 ImageReward 仓库（包含测试数据）
git clone https://github.com/THUDM/ImageReward.git
cd ImageReward

# 安装整合后的包 `image-reward`
pip install image-reward

示例用法

我们在该仓库的 assets/images 目录中提供了示例图片。示例提示词为：

一幅描绘海洋、云朵和鸟类的画作，白天场景，浅景深效果

使用以下代码从 ImageReward 获取人类偏好得分：

import os
import torch
import ImageReward as RM

if __name__ == "__main__":
    prompt = "一幅描绘海洋、云朵和鸟类的画作，白天场景，浅景深效果"
    img_prefix = "assets/images"
    generations = [f"{pic_id}.webp" for pic_id in range(1, 5)]
    img_list = [os.path.join(img_prefix, img) for img in generations]
    model = RM.load("ImageReward-v1.0")
    with torch.no_grad():
        ranking, rewards = model.inference_rank(prompt, img_list)
        # 打印结果
        print("\n偏好预测：\n")
        print(f"ranking = {ranking}")
        print(f"rewards = {rewards}")
        for index in range(len(img_list)):
            score = model.score(prompt, img_list[index])
            print(f"{generations[index]:>16s}: {score:.2f}")

输出应如下所示（具体数值可能因计算设备而略有不同）：

偏好预测：

ranking = [1, 2, 3, 4]
rewards = [[0.5811622738838196], [0.2745276093482971], [-1.4131819009780884], [-2.029569625854492]]
          1.webp: 0.58
          2.webp: 0.27
          3.webp: -1.41
          4.webp: -2.03

ReFL

安装依赖

pip install diffusers==0.16.0 accelerate==0.16.0 datasets==2.11.0

示例用法

我们在该仓库的 data/refl_data.json 中提供了 ReFL 的示例数据集。按如下方式运行 ReFL：

bash scripts/train_refl.sh

ImageReward 和 ReFL 的演示

ImageReward 的训练代码

1. 下载数据：🖼 数据集。

2. 制作数据集。

cd train
python src/make_dataset.py

3. 设置训练配置：train/src/config/config.yaml

4. 一条命令即可开始训练。

bash scripts/train_one_node.sh

集成到 Stable Diffusion Web UI

我们开发了一个自定义脚本，用于将 ImageReward 集成到 SD Web UI 中，以提供便捷的使用体验。

该脚本位于本仓库的 sdwebui/image_reward.py 文件中。

脚本的使用方法如下：

1. 安装：将自定义脚本放入 stable-diffusion-webui/scripts/ 目录。
2. 重新加载：重启服务，或在 SD Web UI 设置页底部点击“重新加载自定义脚本”按钮。（如果找不到该按钮，请尝试点击左侧边栏底部的“显示所有页面”按钮。）
3. 选择：返回“txt2img”或“img2img”选项卡，在左下角的“脚本”下拉菜单中选择“ImageReward - 生成人类偏好分数”。
4. 运行：具体用法根据功能需求有所不同，详见下方的“功能”部分。

功能

为生成的图像打分并附加到图像信息中

使用方法

1. 不要勾选“过滤低分图像”复选框。
2. 点击“生成”按钮以生成图像。
3. 在画廊下方的图像信息底部查看 ImageReward 分数。

演示视频

https://github.com/THUDM/ImageReward/assets/98524878/9d8a036d-1583-4978-aac7-4b758edf9b89

自动过滤低分图像

使用方法

1. 勾选“过滤低分图像”复选框。
2. 在“最低分数限制”中输入分数下限。（ImageReward 大致遵循标准正态分布，均值为 0，方差为 1。）
3. 点击“生成”按钮以生成图像。
4. 分数低于下限的图像将被自动过滤掉，并且不会出现在画廊中。
5. 在画廊下方的图像信息底部查看 ImageReward 分数。

演示视频

https://github.com/THUDM/ImageReward/assets/98524878/b9f01629-87d6-4c92-9990-fe065711b9c6

查看已评分图像的分数

使用方法

1. 在“PNG 信息”选项卡中上传已评分的图像文件。
2. 在右侧查看图像信息，图像分数位于底部。

示例

其他功能

内存管理

• ImageReward 模型仅在首次运行脚本时才会被加载。
• “重新加载 UI”不会重新加载或卸载模型，而是会复用当前已加载的模型（如果存在）。
• 提供了“卸载模型”按钮，可手动卸载当前加载的模型。

常见问题解答

如何调整 SD Web UI 使用的 Python 环境（例如重新安装某个包）？

请注意，SD Web UI 有两种设置 Python 环境的方式：

• 如果您使用 python launch.py 启动，Web UI 将使用您 PATH 中找到的 Python 环境（在 Linux 中，您可以使用 which python 查看其确切路径）。
• 如果您使用类似 webui-user.bat 的脚本启动，Web UI 会在 stable-diffusion-webui\venv 目录下创建一个新的 venv 环境。
  • 通常，您还需要执行一些操作来激活此环境。例如，在 Windows 中，您需要进入 stable-diffusion-webui\venv\Scripts 目录，运行 activate 或 activate.bat（如果您使用的是 cmd），或者运行 activate.ps1（如果您使用的是 PowerShell）。
  • 如果您看到命令行最左侧出现 (venv) 提示符，则表示您已成功激活 SD Web UI 创建的 venv 环境。

激活正确的 Python 环境后，您就可以按照常规方式执行所需操作了。

复现表 1 中的实验

注意：实验结果是在以下环境中产生的：

• (NVIDIA) 驱动版本：515.86.01
• CUDA 版本：11.7
• torch 版本：1.12.1+cu113

根据我们的复现实验经验，在其他环境中复现该实验可能会导致最后一位小数出现波动，通常在 ±0.1 的范围内。

运行以下脚本可自动下载数据、基线模型并运行实验：

bash ./scripts/test-benchmark.sh

然后您可以在 benchmark/results/ 或终端 中查看结果。

如果您想单独查看原始数据文件：

• 测试提示词及对应的图像人类评分位于 benchmark/benchmark-prompts.json。
• 每个提示词的生成结果（最初来自 DiffusionDB）可以从 Hugging Face 或 清华大学云盘 下载。
  • 每个 <model_name>.zip 文件包含一个同名目录，其中共有 1000 张图像，由 100 个提示词各生成 10 张图像。
  • 每个 <model_name>.zip 文件都应解压到 benchmark/generations/ 目录中，形成名为 <model_name> 的子目录，存放这些图像。

复现表 3 中的实验

运行以下脚本可自动下载数据、基线模型并运行实验：

bash ./scripts/test.sh

如果您想单独查看原始数据文件：

• 测试提示词及对应的图像人类评分位于 data/test.json。
• 每个提示词的生成结果（最初来自 DiffusionDB）可以从 Hugging Face 或 清华大学云盘 下载。下载后应解压至 data/test_images 目录。

引用

@inproceedings{xu2023imagereward,
  title={ImageReward: learning and evaluating human preferences for text-to-image generation},
  author={Xu, Jiazheng and Liu, Xiao and Wu, Yuchen and Tong, Yuxuan and Li, Qinkai and Ding, Ming and Tang, Jie and Dong, Yuxiao},
  booktitle={Proceedings of the 37th International Conference on Neural Information Processing Systems},
  pages={15903--15935},
  year={2023}
}

ImageReward 快速上手指南

ImageReward 是首个通用的文生图（Text-to-Image）人类偏好奖励模型。它基于 13.7 万对专家比较数据训练而成，能够准确评估生成图像是否符合人类审美偏好，并支持通过 ReFL 技术直接优化扩散模型。

环境准备

• 操作系统: Linux / macOS / Windows
• Python: 建议 Python 3.8+
• 依赖库: torch, transformers (安装主包时会自动处理)
• 可选依赖 (用于 ReFL 微调): diffusers==0.16.0, accelerate==0.16.0, datasets==2.11.0

提示：本指南优先使用官方 PyPI 源。若下载速度较慢，可添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源加速安装。

安装步骤

ImageReward 已封装为独立的 Python 包 image-reward，安装非常简单。

1. 克隆仓库（可选，仅用于获取测试图片示例）：

   git clone https://github.com/THUDM/ImageReward.git
   cd ImageReward
   

2. 安装核心包：

   pip install image-reward
   

   国内用户加速安装命令：

   pip install image-reward -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装 ReFL 微调所需依赖（如需进行模型微调）：

   pip install diffusers==0.16.0 accelerate==0.16.0 datasets==2.11.0
   

基本使用

以下示例展示如何加载预训练模型并对生成的图像进行打分和排序。

代码示例

假设你有一组由提示词生成的图片，想要知道哪张最符合人类偏好：

import os
import torch
import ImageReward as RM

if __name__ == "__main__":
    # 定义提示词
    prompt = "a painting of an ocean with clouds and birds, day time, low depth field effect"
    
    # 定义图片路径列表 (请替换为你自己的图片路径)
    # 这里假设当前目录下有 1.webp 到 4.webp
    img_list = ["1.webp", "2.webp", "3.webp", "4.webp"]
    
    # 加载预训练模型
    model = RM.load("ImageReward-v1.0")
    
    # 方法一：获取排序结果和具体奖励分数
    with torch.no_grad():
        ranking, rewards = model.inference_rank(prompt, img_list)
        
        print("\n偏好预测结果:\n")
        print(f"排序索引 (从优到劣): {ranking}")
        print(f"对应奖励分数: {rewards}")
        
        # 方法二：单独获取某张图片的分数
        for index, img_path in enumerate(img_list):
            score = model.score(prompt, img_path)
            print(f"{img_path}: {score:.2f}")

输出说明

运行上述代码后，你将得到类似以下的输出。分数越高代表图像越符合人类偏好（通常服从标准正态分布，均值为 0）：

偏好预测结果:

排序索引 (从优到劣): [1, 2, 3, 4]
对应奖励分数: [[0.58], [0.27], [-1.41], [-2.03]]
1.webp: 0.58
2.webp: 0.27
3.webp: -1.41
4.webp: -2.03

进阶：ReFL 微调 (可选)

如果你希望使用 ImageReward 作为奖励信号来微调 Stable Diffusion 模型，可以使用以下极简代码启动训练：

from ImageReward import ReFL

# 解析参数
args = ReFL.parse_args()

# 初始化训练器 (指定基础模型和数据集路径)
trainer = ReFL.Trainer("CompVis/stable-diffusion-v1-4", "data/refl_data.json", args=args)

# 开始训练
trainer.train(args=args)