tongyx361
mageofthesands
Permalink
Name already in use
A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?
ImageReward/sdwebui/image_reward.py /
Go to fileThis commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
121 lines (105 sloc)
4.14 KB
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import os | |
| import subprocess | |
| import sys | |
| from pathlib import Path | |
| import gradio as gr | |
| import modules.images as images | |
| import modules.scripts as scripts | |
| import torch | |
| from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE | |
| from modules import sd_samplers, shared | |
| from modules.processing import ( | |
| Processed, | |
| StableDiffusionProcessing, | |
| create_infotext, | |
| process_images, | |
| ) | |
| from modules.shared import cmd_opts, opts, state | |
| try: | |
| import ImageReward as reward | |
| except ModuleNotFoundError as error: | |
| print('"image-reward" package has not been properly installed. Installing...') | |
| if subprocess.check_call(["pip", "install", "image-reward"]) == 0: | |
| print(f'"image-reward" package is successfully installed!') | |
| import ImageReward as reward | |
| else: | |
| print('"image-reward" package installation failed!') | |
| print( | |
| "Please open an issue with full error message at https://github.com/THUDM/ImageReward/issues" | |
| ) | |
| def unload_image_reward_model(): | |
| del shared.image_reward_model | |
| 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): | |
| with gr.Blocks(): | |
| with gr.Row(): | |
| gr.Markdown( | |
| value="**Tip**: It will take a little time to **load** the ImageReward model before **the first generation**." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| filter_out_low_scores = gr.Checkbox( | |
| value=False, label="Filter out images with low scores" | |
| ) | |
| with gr.Column(): | |
| lower_score_limit = gr.Textbox(value=0, label="Lower score limit") | |
| with gr.Row(): | |
| gr.Markdown( | |
| value="ImageReward model takes about **1,600 MB** of memory." | |
| ) | |
| with gr.Row(): | |
| unload_button = gr.Button(value="Unload Model From Memory") | |
| unload_button.click(unload_image_reward_model) | |
| return [filter_out_low_scores, lower_score_limit] | |
| def run(self, p, filter_out_low_scores, lower_score_limit): | |
| try: | |
| shared.image_reward_model # if loaded, do nothing | |
| except AttributeError: | |
| # load the model | |
| if sys.platform == "win32": | |
| download_root = HUGGINGFACE_HUB_CACHE | |
| else: | |
| download_root = None | |
| print(f"Loading ImageReward model from {download_root}...") | |
| shared.image_reward_model = reward.load( | |
| "ImageReward-v1.0", download_root=download_root | |
| ) # using shared to make the model object global among modules | |
| # preprocess parameters | |
| if lower_score_limit != "": | |
| lower_score_limit = float(lower_score_limit) | |
| # generate images | |
| proc = process_images(p) | |
| # score | |
| gens = proc.images | |
| with torch.no_grad(): | |
| for img in gens: | |
| score = shared.image_reward_model.score(p.prompt, img) | |
| img.info["score"] = score | |
| if img.info.get("parameters") is None: | |
| img.info["parameters"] = f"ImageReward Score: {score:.4f}" | |
| else: | |
| 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, | |
| ) |
ImageReward does not work. I am using stable-diffusion-webui-directml. I followed instructions from "Quick Start" as well as "Integration into Stable Diffusion Web UI". Attached screenshots to show that ImageReward is installed, and script is in then place instructions say.
