├── .github └── workflows │ └── publish.yml ├── LICENSE ├── README.md ├── __init__.py ├── nodes.py ├── pyproject.toml ├── requirements.txt └── screenshot.png /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish to Comfy registry 2 | on: 3 | workflow_dispatch: 4 | push: 5 | branches: 6 | - main 7 | - master 8 | paths: 9 | - "pyproject.toml" 10 | 11 | permissions: 12 | issues: write 13 | 14 | jobs: 15 | publish-node: 16 | name: Publish Custom Node to registry 17 | runs-on: ubuntu-latest 18 | if: ${{ github.repository_owner == 'ZaneA' }} 19 | steps: 20 | - name: Check out code 21 | uses: actions/checkout@v4 22 | - name: Publish Custom Node 23 | uses: Comfy-Org/publish-node-action@v1 24 | with: 25 | ## Add your own personal access token to your Github Repository secrets and reference it here. 26 | personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} 27 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Zane Ashby 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ComfyUI-ImageReward 2 | 3 | Score images with [ImageReward](https://github.com/THUDM/ImageReward). 4 | 5 | ## Installation 6 | 7 | Within your ComfyUI directory, and within your Python venv: 8 | 9 | ```sh 10 | cd custom_nodes 11 | git clone https://github.com/ZaneA/ComfyUI-ImageReward 12 | pip install -r ComfyUI-ImageReward/requirements.txt 13 | ``` 14 | 15 | ## Usage 16 | 17 | Add an `ImageRewardLoader` node, this has the default model name prefilled, and is passed directly to the ImageReward loader. Note, the model will be downloaded on first run. 18 | 19 | Add an `ImageRewardScore` node, connect the model, your image, and your prompt (either enter this directly, or right click the node and convert prompt to an input first). Connect the `SCORE_FLOAT` or `SCORE_STRING` output to an appropriate node. 20 | 21 | ![Node screenshot](screenshot.png) 22 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from .nodes import NODE_CLASS_MAPPINGS 2 | 3 | __all__ = ["NODE_CLASS_MAPPINGS"] 4 | -------------------------------------------------------------------------------- /nodes.py: -------------------------------------------------------------------------------- 1 | import ImageReward as RM 2 | import numpy as np 3 | from PIL import Image 4 | 5 | class ImageRewardLoader: 6 | CATEGORY = "ImageReward" 7 | 8 | def __init__(self): 9 | pass 10 | 11 | @classmethod 12 | def INPUT_TYPES(s): 13 | return { 14 | "required": { 15 | "model": ("STRING", { 16 | "multiline": False, 17 | "default": "ImageReward-v1.0" 18 | }), 19 | }, 20 | } 21 | 22 | RETURN_TYPES = ("IMAGEREWARD_MODEL",) 23 | RETURN_NAMES = ("IMAGEREWARD_MODEL",) 24 | OUTPUT_NODE = False 25 | 26 | FUNCTION = "load_model" 27 | 28 | def load_model(self, model): 29 | model = RM.load(model) 30 | return (model,) 31 | 32 | class ImageRewardScore: 33 | CATEGORY = "ImageReward" 34 | 35 | def __init__(self): 36 | pass 37 | 38 | @classmethod 39 | def INPUT_TYPES(s): 40 | return { 41 | "required": { 42 | "model": ("IMAGEREWARD_MODEL",), 43 | "prompt": ("STRING", { 44 | "multiline": True 45 | }), 46 | "images": ("IMAGE",), 47 | }, 48 | } 49 | 50 | RETURN_TYPES = ("FLOAT","STRING") 51 | RETURN_NAMES = ("SCORE_FLOAT","SCORE_STRING") 52 | OUTPUT_NODE = False 53 | 54 | FUNCTION = "score_images" 55 | 56 | def score_images(self, model, prompt, images): 57 | score = 0.0 58 | for image in images: 59 | # convert to PIL image 60 | i = 255.0 * image.cpu().numpy() 61 | img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) 62 | score += model.score(prompt, [img]) 63 | score /= len(images) 64 | return (score,str(score)) 65 | 66 | 67 | NODE_CLASS_MAPPINGS = { 68 | "ImageRewardLoader": ImageRewardLoader, 69 | "ImageRewardScore": ImageRewardScore 70 | } 71 | 72 | NODE_DISPLAY_NAME_MAPPINGS = { 73 | "ImageRewardLoader": "ImageReward Loader", 74 | "ImageRewardScore": "ImageReward Score" 75 | } 76 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "comfyui-imagereward" 3 | description = "NODES: ImageRewardLoader, ImageRewardScore" 4 | version = "1.0.0" 5 | license = { file = "LICENSE" } 6 | dependencies = ["image-reward==1.4", "clip @ git+https://github.com/openai/CLIP.git"] 7 | 8 | [project.urls] 9 | Repository = "https://github.com/ZaneA/ComfyUI-ImageReward" 10 | # Used by Comfy Registry https://comfyregistry.org 11 | 12 | [tool.comfy] 13 | PublisherId = "zanea" 14 | DisplayName = "ComfyUI-ImageReward" 15 | Icon = "" 16 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | image-reward==1.4 2 | clip @ git+https://github.com/openai/CLIP.git 3 | -------------------------------------------------------------------------------- /screenshot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZaneA/ComfyUI-ImageReward/5e9936ded972d8447c26ea2670a92e999036821d/screenshot.png --------------------------------------------------------------------------------