Initialize the demo.

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: ImageReward Demo
+emoji: 👩‍🎨
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+sdk_version: 3.28.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,241 @@
+import os
+
+import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+
+import ImageReward as RM
+
+# initialize
+
+model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+)
+
+model = RM.load("ImageReward-v1.0")
+
+images_in_gallery = []
+rewards_in_gallery = []
+
+# event functions
+
+
+def generate_images(
+    prompt, magic_words, num, height, width, num_inference_steps, guidance_scale
+):
+    global images_in_gallery, rewards_in_gallery
+
+    if magic_words is not None:
+        prompt += ", ".join(magic_words)
+
+    images_in_gallery = pipe(
+        prompt,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num,
+    ).images
+    rewards_in_gallery = [None] * len(images_in_gallery)
+    return list(zip(images_in_gallery, rewards_in_gallery))
+
+
+def score_and_rank(prompt):
+    global rewards_in_gallery, images_in_gallery
+
+    num_not_scored = rewards_in_gallery.count(None)
+
+    if num_not_scored > 0:
+        images_to_score = images_in_gallery[-num_not_scored:]
+        with torch.no_grad():
+            ranks, rewards = model.inference_rank(prompt, images_to_score)
+
+        if not isinstance(rewards, list):
+            rewards = [rewards]
+        rewards_in_gallery = rewards_in_gallery[:-num_not_scored] + rewards
+
+        outputs = sorted(
+            zip(images_in_gallery, rewards_in_gallery), key=lambda x: x[1], reverse=True
+        )
+
+        images_in_gallery = [image for image, _ in outputs]
+        rewards_in_gallery = [reward for _, reward in outputs]
+
+        return outputs, [
+            [idx + 1, reward] for idx, reward in enumerate(rewards_in_gallery)
+        ]
+    else:
+        return list(zip(images_in_gallery, rewards_in_gallery)), [
+            [idx + 1, reward] for idx, reward in enumerate(rewards_in_gallery)
+        ]
+
+
+def upload_images_to_gallery(uploaded_image_files):
+    global images_in_gallery, rewards_in_gallery
+
+    uploaded_image_file_paths = [file.name for file in uploaded_image_files]
+    uploaded_images = [Image.open(path) for path in uploaded_image_file_paths]
+    for path in uploaded_image_file_paths:
+        os.remove(path)
+    images_in_gallery = images_in_gallery + uploaded_images
+    rewards_in_gallery = rewards_in_gallery + [None] * len(uploaded_images)
+
+    return list(zip(images_in_gallery, rewards_in_gallery))
+
+
+def clear_images():
+    global images_in_gallery, rewards_in_gallery
+    images_in_gallery = []
+    rewards_in_gallery = []
+    return None
+
+
+if __name__ == "__main__":
+    # UI
+    with gr.Blocks(
+        theme=gr.themes.Monochrome(),
+        css=r".caption-label { color: black; }",
+    ) as demo:
+        gr.HTML(
+            """
+            <h1 align="center">ImageReward Demo</h1>
+            <p align="center"><a href="https://github.com/THUDM/ImageReward">GitHub Repo</a> • 🤗 <a href="https://huggingface.co/THUDM/ImageReward" target="_blank">HF Repo</a> • 🐦 <a href="https://twitter.com/thukeg" target="_blank">Twitter</a> • 📃 <a href="https://arxiv.org/abs/2304.05977" target="_blank">Paper</a><br></p>
+            <br>
+            <p dir="auto">ImageReward is the first general-purpose text-to-image <strong>human preference RM</strong>, which is trained on in total <strong>137k pairs of expert comparisons</strong>!</p>
+            <p dir="auto">The calculation of ImageRewards is based on <strong>both the prompt and images</strong>.</p>
+            """
+        )
+        with gr.Row():
+            with gr.Column():
+                gr.HTML(
+                    """
+                    <p dir="auto">Try ImageReward with only 2 steps:</p>
+                    <ol dir="auto">
+                    <li>Click the <strong>"Generate"</strong> button <strong>in the middle of the bottom</strong>.</li>
+                    <li>Click the <strong>"Score&Rank"</strong> button <strong>below the gallery</strong>.</li>
+                    </ol>
+                    <p dir="auto">Finally, just check ImageRewards <strong>along with images or on the right of the gallery</strong>.</p>
+                    <br>
+                    <p dir="auto">This demo uses <code>runwayml/stable-diffusion-v1-5</code> as image generation model.</p>
+                    """
+                )
+            with gr.Column():
+                gr.HTML(
+                    """
+                    <p dir="auto">Besides generating images, you can also <strong>upload</strong> images to score:</p>
+                    <ol dir="auto">
+                    <li>Upload images <strong>in the bottom right corner</strong>.</li>
+                    <li>Change the <strong>"Prompt"</strong> to correspond to the images.</li>
+                    <li>Click the <strong>"Score&Rank"</strong> button <strong>below the gallery</strong>.</li>
+                    </ol>
+                    <br>
+                    <p dir="auto">For more details about using ImageReward in your own program, check <a href="https://github.com/THUDM/ImageReward">the README.md in our Github Repo</a>.</p>
+                    """
+                )
+
+        with gr.Row(elem_id="outputs_row"):
+            with gr.Column(elem_id="gallery_column", scale=4):
+                gallery = gr.Gallery(
+                    label="Images (scored ones sorted)",
+                    show_label=False,
+                    elem_id="gallery",
+                ).style(columns=4, object_fit="contain", full_width=True)
+            with gr.Column(elem_id="rewards_column"):
+                rewards = gr.Matrix(
+                    value=[[None, None]],
+                    headers=["Rank", "ImageReward"],
+                    datatype="number",
+                )
+        with gr.Row():
+            score_and_rank_button = gr.Button("Score&Rank")
+            clear_button = gr.Button("Clear Gallery")
+        with gr.Row().style(equal_height=True):
+            with gr.Column():
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    value="A painting of an ocean with clouds and birds, day time, low depth field effect, oil painting, impressionism",
+                )
+
+                examples = [
+                    "A painting of an ocean with clouds and birds, day time, low depth field effect, oil painting, impressionism",
+                    "A painting of a girl walking in a hallway and suddenly finds a giant sunflower on the floor blocking her way",
+                    "Coronation of the sun emperor, digital art, illustration,4k resolution,intricate extremely detailed, depth,vivid colors",
+                    "Symmetry!! Product render poster vivid colors divine proportion owl,glowing fog intricate,elegant, highly detailed",
+                    "A unicorn in a clearing.it has a single shining horn. volumetric light.by emmanuel shiu, harry potter, eragon",
+                    "Highly detailed portrait of a woman with long hairs,stephen bliss. unreal engine, fantasy art by greg rutkowski",
+                    "Sculpture made of flame,portrait, female,future, torch,fire,harper's bazaar,vogue, fashion magazine, intricate",
+                ]
+                prompt_examples = gr.Examples(
+                    examples=examples,
+                    label="Prompt Examples",
+                    inputs=[prompt],
+                    elem_id="prompt_examples",
+                )
+
+            with gr.Column():
+                choices = [
+                    "HDR, UHD, 4K, 8K, 64K",
+                    "highly detailed",
+                    "studio lighting",
+                    "professional",
+                    "trending on artstation",
+                    "unreal engine",
+                    "vivid colors",
+                ]
+                magic_words = gr.CheckboxGroup(
+                    choices=choices,
+                    value=choices,
+                    type="value",
+                    label="Magic Words to Append to Prompt",
+                )
+
+                num = gr.Slider(1, 16, step=1, label="Number of images", value=8)
+                height = gr.Slider(256, 2048, step=256, label="Height", value=512)
+                width = gr.Slider(256, 2048, step=256, label="Width", value=512)
+                num_inference_steps = gr.Slider(
+                    0, 200, step=10, label="Number of inference steps", value=50
+                )
+                guidance_scale = gr.Slider(
+                    0, 25, step=0.1, label="Guidance scale", value=7.5
+                )
+
+                generate_button = gr.Button("Generate")
+            with gr.Column():
+                gr.Markdown(
+                    """
+                    - To clear all uploaded images, click the **"Clear Gallery"** button above.
+                    - To clear the upload list and add additional images, click the **`x` in the upper right corner of the uploading window**.
+                    - Additional images will be appended to the gallery, instead of replacing the existing ones.
+                    """
+                )
+                uploaded_image_files = gr.File(
+                    file_count="multiple",
+                    file_types=["image"],
+                    type="file",
+                    label="Upload Images",
+                    show_label=True,
+                )
+
+        generate_button.click(
+            generate_images,
+            [
+                prompt,
+                magic_words,
+                num,
+                height,
+                width,
+                num_inference_steps,
+                guidance_scale,
+            ],
+            [gallery],
+        )
+        score_and_rank_button.click(score_and_rank, [prompt], [gallery, rewards])
+        uploaded_image_files.upload(
+            upload_images_to_gallery, [uploaded_image_files], [gallery]
+        )
+        clear_button.click(clear_images, None, [gallery])
+
+    demo.launch()

requirements.txt

@@ -0,0 +1,13 @@
+image-reward
+timm==0.6.13
+transformers==4.27.4
+fairscale==0.4.13
+huggingface_hub==0.13.4
+clip @ git+https://github.com/openai/CLIP.git
+
+torch
+diffusers
+accelerate
+
+Pillow
+gradio