DOCSEARCHbyAlgolia
Submit Docs

ImageReward Documentation

repository·main·Indexed 23 days ago

https://github.com/zai-org/imagereward

ImageReward is a human preference reward model (RM) for text-to-image generation. It provides scoring for image-text alignment and includes Reward Feedback Learning (ReFL) for fine-tuning diffusion models to better match human preferences. The library allows users to score images, rank image lists, and integrate with Stable Diffusion Web UI. It includes a training pipeline using a pre-trained BLIP model and MLP head, supporting distributed training and fine-grained layer freezing.

Tokens
2.5K
Snippets
4
Records
12
Agent score
83%

What's inside ImageReward

  1. Fine-tune diffusion models with ReFL

    main

    Reward Feedback Learning (ReFL) allows for direct optimization of a text-to-image diffusion model using ImageReward.

    Dependencies:

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

    Usage Example: Use ReFL.parse_args() to get configuration arguments and ReFL.Trainer to initialize the training process with a base model (e.g., CompVis/stable-diffusion-v1-4) and a dataset path.

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

  2. Install the image-reward package

    main

    To use ImageReward for scoring or ReFL fine-tuning, install the integrated Python package via pip. It is recommended to clone the repository first if you need the testing data or scripts.

    # Clone the repository for testing data and scripts
git clone https://github.com/THUDM/ImageReward.git
cd ImageReward

# Install the package
pip install image-reward
    # 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

  3. Reproduce Experiments in Table 3

    main

    To automatically download necessary data and baseline models and run the experiments described in Table 3 of the paper, execute the provided test script.

    Data Locations

    If you need to inspect or manually manage the raw data files:

    • Test prompts and human rankings: Located in data/test.json.
    • Generated outputs: These are images originally from DiffusionDB. You can download them from Hugging Face or Tsinghua Cloud.
    • Important: Downloaded images must be decompressed into the data/test_images directory.
    bash ./scripts/test.sh

  4. Integrate ImageReward into Stable Diffusion Web UI

    main

    You can use a custom script to add ImageReward features to the Stable Diffusion Web UI.

    Setup Steps:

    1. Install: Copy the sdwebui/image_reward.py script from this repository into your stable-diffusion-webui/scripts/ directory.
    2. Reload: Restart the Web UI service or click "Reload custom script" in the settings tab.
    3. Select: In the "txt2img" or "img2img" tabs, select "ImageReward - generate human preference scores" from the "Script" dropdown menu.

    Features:

    • Score and Append: Generates images and appends the score to the image information below the gallery.
    • Automatic Filtering: Check "Filter out images with low scores" and set a "Lower score limit" to prevent low-quality images from appearing in the gallery.
    • View Scores: Upload scored images to the "PNG Info" tab to view their scores.
    • Memory Management: The model only loads on the first run. Use the "Unload Model" button to free up VRAM.

  5. Image preprocessing in ImageReward

    main

    Images are processed using a standard transformation pipeline before being passed to the visual encoder. This ensures consistency with the pre-trained BLIP model.

    Pipeline steps:

    1. Resize: Resizes to the configured BLIP['image_size'] using BICUBIC interpolation.
    2. CenterCrop: Crops the center of the image to the target size.
    3. RGB Conversion: Ensures the image is in RGB format.
    4. ToTensor: Converts the PIL image to a PyTorch tensor.
    5. Normalize: Applies mean (0.48145466, 0.4578275, 0.40821073) and std (0.26862954, 0.26130258, 0.27577711).

  6. Understand the training loss and accuracy metrics

    main

    The training script uses a specific loss_func to evaluate the reward model.

    • Loss: Calculated using F.cross_entropy against a target of zeros. This effectively treats the reward training as a classification task where the model is encouraged to output a specific target value.
    • Accuracy (Acc): Calculated based on the difference between paired rewards. For a pair of rewards [reward_0, reward_1], accuracy is the mean of instances where reward_0 - reward_1 > 0.

  7. Configure ImageReward layer freezing

    main

    The ImageReward model allows for fine-grained control over which layers of the underlying BLIP model are frozen during training using opts (configuration options):

    • opts.fix_base: If set to True, the entire BLIP base model is frozen (requires_grad_(False)).
    • opts.fix_rate: A float (e.g., 0.5) that determines the ratio of layers to freeze. It calculates a cutoff point for both the text encoder and the visual encoder layers. Layers up to the calculated index are frozen.
    • loose_layer(fix_rate): A method that can be called to 'unfreeze' specific layers. It targets the top layers of the text encoder and visual encoder based on the provided fix_rate, setting their requires_grad to True.

  8. Train the ImageReward model

    main
    The train.py script serves as the entrypoint for fine-tuning the ImageReward model. It supports both single-GPU and distributed training modes. The training process involves loading datasets (either standard rank_dataset or rank_pair_dataset depending on configuration), optimizing the model using the Adam optimizer, and performing periodic validation to save the best performing model based on validation loss.

  9. Score images using ImageReward

    main

    You can use the ImageReward module to score how well images match a text prompt based on human preference.

    Use RM.load("ImageReward-v1.0") to load the model. You can then use model.score(prompt, image_path) for a single image or model.inference_rank(prompt, img_list) to rank a list of images and get their rewards simultaneously.

    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}")

  10. Reference: ImageReward input data formats

    main
    The ImageReward model processes data in two primary formats depending on the encoding mode selected via opts.rank_pair.

  11. Use the ImageReward class for scoring

    main

    The ImageReward class is the core model used to score the quality of text-to-image generation. It utilizes a pre-trained BLIP model and an MLP head to compare image-text pairs.

    Depending on the configuration (opts.rank_pair), the model can operate in two modes:

    1. Pairwise Encoding (encode_pair): Takes a batch containing explicit img_better and img_worse tensors and their corresponding text IDs/masks to compute rewards.
    2. Data Encoding (encode_data): Takes a batch of raw data (prompts, image paths, and rankings) and performs cross-attention between text and images to generate embeddings for comparison.

    The forward method returns a concatenated tensor of rewards for the 'better' and 'worse' samples in the batch.

  12. Configure training via command-line options

    main

    The training script relies on an opts object (imported from config.options) to control execution. While the exact CLI parser is defined in config/options.py, the following parameters are directly utilized in the training logic:

    • opts.distributed: Boolean flag to enable distributed training using torch.distributed with the nccl backend.
    • opts.rank_pair: Boolean flag to switch between rank_dataset and rank_pair_dataset.
    • opts.batch_size: The per-device batch size.
    • opts.accumulation_steps: Number of steps for gradient accumulation.
    • opts.gpu_num: Number of GPUs used.
    • opts.seed: Random seed for reproducibility.
    • opts.epochs: Total number of training epochs.
    • opts.valid_per_epoch: Frequency of validation steps per epoch.
    • opts.lr: Learning rate for the Adam optimizer.
    • opts.preload_path: Path to a pre-trained model to load before training.
    • opts.std_log: Boolean flag to enable standard logging to a text file.
    • opts.gpu_id: Environment variable setting for CUDA_VISIBLE_DEVICES.
    • opts.adam_beta1, opts.adam_beta2, opts.adam_eps: Adam optimizer hyperparameters.