zai-org/ImageReward

zai-org/ImageReward

Description: [NeurIPS 2023] ImageReward: Learning and Evaluating Human Preferences for Text-to-image Generation

Language: Python

License: Apache-2.0

Stars: 1681

Forks: 92

Open issues: 60

Created: 2023-04-01T09:04:17Z

Pushed: 2025-10-29T13:40:51Z

Default branch: main

Fork: no

Archived: no

README:

ImageReward

📃 Paper • 🖼 Dataset • 🌐 中文博客 • 🤗 HF Repo • 🐦 Twitter

🔥🔥 News! ``2024/12/31``: We released the next generation of model, [VisionReward](https://github.com/THUDM/VisionReward), which is a fine-grained and multi-dimensional reward model for stable RLHF for visual generation (text-to-image / text-to-video)!

🔥 News! ``2023/9/22``: The paper of ImageReward is accepted by NeurIPS 2023!

ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation

ImageReward is the first general-purpose text-to-image human preference RM, which is trained on in total 137k pairs of expert comparisons, outperforming existing text-image scoring methods, such as CLIP (by 38.6%), Aesthetic (by 39.6%), and BLIP (by 31.6%), in terms of understanding human preference in text-to-image synthesis.

Additionally, we introduce Reward Feedback Learning (ReFL) for direct optimizing a text-to-image diffusion model using ImageReward. ReFL-tuned Stable Diffusion wins against untuned version by 58.4% in human evaluation.

Both ImageReward and ReFL are all packed up to Python image-reward package now!

Try image-reward package in only 3 lines of code for ImageReward scoring!

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

rewards = model.score("", ["", "", ...])

Try image-reward package in only 4 lines of code for ReFL fine-tuning!

# 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)

If you find ImageReward's open-source effort useful, please 🌟 us to encourage our following developement!

• [ImageReward](#imagereward)
• [Quick Start](#quick-start)
• [Install Dependency](#install-dependency)
• [Example Use](#example-use)
• [ReFL](#refl)
• [Install Dependency](#install-dependency-1)
• [Example Use](#example-use-1)
• [Demos of ImageReward and ReFL](#demos-of-imagereward-and-refl)
• [Training code for ImageReward](#training-code-for-imagereward)
• [Integration into Stable Diffusion Web UI](#integration-into-stable-diffusion-web-ui)
• [Features](#features)
• [Score generated images and append to image information](#score-generated-images-and-append-to-image-information)
• [Usage](#usage)
• [Demo video](#demo-video)
• [Automatically filter out images with low scores](#automatically-filter-out-images-with-low-scores)
• [Usage](#usage-1)
• [Demo video](#demo-video-1)
• [View the scores of images that have been scored](#view-the-scores-of-images-that-have-been-scored)
• [Usage](#usage-2)
• [Example](#example)
• [Other Features](#other-features)
• [Memory Management](#memory-management)
• [FAQ](#faq)
• [Reproduce Experiments in Table 1](#reproduce-experiments-in-table-1)
• [Reproduce Experiments in Table 3](#reproduce-experiments-in-table-3)
• [Citation](#citation)

Quick Start

Install Dependency

We have integrated the whole repository to a single python package image-reward. Following the commands below to prepare the environment:

# Clone the ImageReward repository (containing data for testing)
git clone https://github.com/THUDM/ImageReward.git
cd ImageReward

# Install the integrated package `image-reward`
pip install image-reward

Example Use

We provide example images in the [assets/images](assets/images) directory of this repo. The example prompt is:

a painting of an ocean with clouds and birds, day time, low depth field effect

Use the following code to get the human preference scores from ImageReward:

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"
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 the result
print("\nPreference predictions:\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}")

The output should be like as follow (the exact numbers may be slightly different depending on the compute device):

Preference predictions:

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

Install Dependency

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

Example Use

We provide example dataset for ReFL in the [data/refl_data.json](data/refl_data.json) of this repo. Run ReFL as following:

bash scripts/train_refl.sh

Demos of ImageReward and ReFL

Training code for ImageReward

1. Download data: 🖼 Dataset.

2. Make dataset.

cd train
python src/make_dataset.py

3. Set training config: [train/src/config/config.yaml](train/src/config/config.yaml)

4. One command to train.

bash scripts/train_one_node.sh

Integration into Stable Diffusion Web UI

We have developed a custom script to integrate ImageReward into SD Web UI for a convenient experience.

The script is located at [sdwebui/image_reward.py](sdwebui/image_reward.py) in this repository.

The usage of the script is described as follows:

1. Install: put the custom script into the `stable-diffusion-webui/scripts/` directory 2. Reload: restart the service, or click the "Reload custom script" button at the bottom of the settings tab of SD Web UI. (If the button can't be found, try clicking the "Show all pages" button at the bottom of the left sidebar.) 3. Select:…