Add integration into SD Web UI

AniAggarwal · Apr 24, 2023 · c553037 · c553037
1 parent c0b9080
commit c553037
Show file tree

Hide file tree

Showing 2 changed files with 131 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -6,7 +6,7 @@
 
 **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**. 
+ImageReward is the first general-purpose text-to-image human preference RM, which is trained on in total **137k pairs of expert comparisons**.
 
 It outperforms 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. Try `image-reward` package in only 3 lines of code!
 
@@ -85,6 +85,61 @@ rewards = [[0.5811622738838196], [0.2745276093482971], [-1.4131819009780884], [-
           4.webp: -2.03
 ```
 
+## Integration into [Stable Diffusion Web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
+
+We have developed a **custom script** to integrate ImageReward into SD Web UI for a convenient experience.
+
+The script is located at [`demo/sdwebui/image_reward.py`](demo/sdwebui/image_reward.py) in this repository.
+
+The **usage** of the script is described as follows:
+
+1. **Dependency**: [`pip install image-reward`](https://pypi.org/project/image-reward/)
+2. **Install**: put the custom script into the [`stable-diffusion-webui/scripts/`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/tree/master/scripts) directory
+3. **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.)
+4. **Select**: go back to the **"txt2img"/"img2img"** tab, and select **"ImageReward - generate human preference scores"** from the "**Script"** dropdown menu in the lower left corner.
+5. **Run**: the specific usage varies depending on the functional requirements, as described in the **"Features"** section below.
+
+### Features
+
+#### Score generated images and append to image information
+
+##### Usage:
+
+1. **Do not** check the "Filter out images with low scores" checkbox.
+2. Click the **"Generate"** button to generate images.
+3. Check the ImageReward at the **bottom** of the image information **below the gallery**.
+
+##### Demo video:
+
+https://user-images.githubusercontent.com/98524878/233889441-d593675a-dff4-43aa-ad6b-48cc68326fb0.mp4
+
+#### Automatically filter out images with low scores
+
+##### Usage:
+
+1. Check the **"Filter out images with low scores"** checkbox.
+2. Enter the score lower limit in **"Lower score limit"**. (ImageReward roughly follows the standard normal distribution, with a mean of 0 and a variance of 1.)
+3. Click the **"Generate"** button to generate images.
+4. Images with scores below the lower limit will be automatically filtered out and **will not appear in the gallery**.
+5. Check the ImageReward at the **bottom** of the image information **below the gallery**.
+
+##### Demo video:
+
+https://user-images.githubusercontent.com/98524878/233889490-5c4a062f-bb5e-4179-ba98-b336cda4d290.mp4
+
+#### View the scores of images that have been scored
+
+##### Usage:
+
+1. Upload the scored image file in the **"PNG Info"** tab
+2. Check the image information on the right with the score of the image at the **bottom**.
+
+##### Example:
+
+<p align="center">
+    <img src="https://user-images.githubusercontent.com/98524878/233829640-12190bff-f62b-4160-b05d-29624fa83677.jpg" width="700px">
+</p>
+
 ## Reproduce Experiments in Table 2
 
 <p align="center">
@@ -99,10 +154,11 @@ bash ./scripts/test.sh
 
 If you want to check the raw data files individually:
 
-* Test prompts and corresponding human rankings for images are located in [`data/test.json`](data/test.json). 
+* Test prompts and corresponding human rankings for images are located in [`data/test.json`](data/test.json).
 * Generated outputs for each prompt (originally from [DiffusionDB](https://github.com/poloclub/diffusiondb)) can be downloaded from [Huggingface](https://huggingface.co/THUDM/ImageReward/blob/main/test_images.zip) or [Tsinghua Cloud](https://cloud.tsinghua.edu.cn/f/9bd245027652422499f4/?dl=1). It should be decompressed to `data/test_images`.
 
 ## Reproduce Experiments in Table 4
+
 TODO
 
 ## Citation

diff --git a/demo/sdwebui/image_reward.py b/demo/sdwebui/image_reward.py
@@ -0,0 +1,73 @@
+import modules.scripts as scripts
+import gradio as gr
+
+from modules import sd_samplers, shared
+from modules.processing import Processed, process_images, StableDiffusionProcessing, create_infotext
+import modules.images as images
+from modules.shared import opts, cmd_opts, state
+
+import torch
+import os
+import sys
+from pathlib import Path
+import ImageReward as reward
+
+# load the model
+# by default, it will:
+# 1. set the device to cuda if available
+# 2. download the model and cache it in `~/.cache/` if model is not found
+# you can alse configure the device and cache dir by passing in the arguments
+model = reward.load("ImageReward-v1.0")
+
+class Script(scripts.Script):
+
+    def title(self):  
+        return "ImageReward - generate human preference scores"
+
+    def show(self, is_txt2img):
+        return True
+
+    def ui(self, is_txt2img):
+        filter_out_low_scores = gr.Checkbox(value=False, label="Filter out images with low scores")
+        lower_score_limit = gr.Textbox(value=0, label="Lower score limit")
+        return [filter_out_low_scores, lower_score_limit]
+
+    def run(self, p, filter_out_low_scores, lower_score_limit):
+        # preprocess parameters
+        if lower_score_limit != '':
+            lower_score_limit = float(lower_score_limit)
+
+        # generate images
+        proc = process_images(p)
+
+        # score
+        gens = proc.images
+        for img in gens:
+            with torch.no_grad():
+                score = model.score(p.prompt, img)
+            img.info['score'] = score
+            img.info['parameters'] += f"\n ImageReward Score: {score:.4f}"
+
+        # filter out images with scores lower than the lower limit
+        if filter_out_low_scores:    
+            imgs = list(filter(lambda x: x.info['score'] > lower_score_limit, gens))
+        else:
+            imgs = gens
+
+        # append score to info
+        infotexts = [img.info['parameters'] for img in imgs]
+
+        # sort to score
+        img_info_list = list(zip(imgs, infotexts))
+        img_info_list.sort(key=lambda x: x[0].info['score'], reverse=True)
+        imgs, infotexts = list(zip(*img_info_list))
+
+        # return Processed object
+        return Processed(
+            p=p,
+            images_list=imgs,
+            info=proc.info,
+            seed=proc.seed,
+            infotexts=infotexts,
+            index_of_first_image=proc.index_of_first_image
+        )