diff --git a/README.md b/README.md index f0d09d1..85eb61f 100644 --- a/README.md +++ b/README.md @@ -1 +1,40 @@ -# ComfyUI-BrushNet \ No newline at end of file +## ComfyUI-BrushNet + +Custom nodes for ComfyUI allow to inpaint using Brushnet: ["BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"](https://arxiv.org/abs/2403.06976). + +My contribution is limited to the ComfyUI adaptation, and all credit goes to the authors of the paper. + + +## Installation + +Clone the repo into the custom_nodes directory and install the requirements: + +``` +git clone https://github.com/nullquant/ComfyUI-BrushNet.git +pip install -r requirements.txt +``` + +Checkpoints of BrushNet can be downloaded from [here](https://drive.google.com/drive/folders/1fqmS1CEOvXCxNWFrsSYd_jHYXxrydh1n?usp=drive_link). +Both `segmentation_mask_brushnet_ckpt` and `random_mask_brushnet_ckpt` should be placed inside models/inpaint/brushnet folder. +Base model should be SD1.5 type (realisticVisionV60B1_v51VAE, for example). + +## Usage + +Below is an example for the intended workflow. The [json file](example/BrushNet example.json) for the example can be found inside the 'workflow' directory. + +![example workflow](example/BrushNet example.png?raw=true) + + +## Bugs + +1. Nodes use custom `diffusers` library. It can potentially conflicts with other custom nodes. +2. Sometimes `BrushNetPipeline` freezes on loading. Looking into the problem. + + +## Credits + +The code is based on + +- [BrushNet](https://github.com/TencentARC/BrushNet) +- [PowerPaint](https://github.com/zhuang2002/PowerPaint) +- [diffusers](https://github.com/huggingface/diffusers) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e3a6c83 --- /dev/null +++ b/__init__.py @@ -0,0 +1,24 @@ +from .brushnet_nodes import BrushNetLoader, BrushNetPipeline, BrushNetInpaint, BlendInpaint +""" +@author: nullquant +@title: BrushNet +@nickname: BrushName nodes +@description: This repository contains an custom nodes for inpaint using BrushNet and PowerPaint models +""" + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "BrushNetLoader": BrushNetLoader, + "BrushNetPipeline": BrushNetPipeline, + "BrushNetInpaint": BrushNetInpaint, + "BlendInpaint": BlendInpaint, +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "BrushNetLoader": "BrushNet Loader", + "BrushNetPipeline": "BrushNet Pipeline", + "BrushNetInpaint": "BrushNet Inpaint", + "BlendInpaint": "Blend Inpaint", +} \ No newline at end of file diff --git a/brushnet_nodes.py b/brushnet_nodes.py new file mode 100644 index 0000000..ad6c4b6 --- /dev/null +++ b/brushnet_nodes.py @@ -0,0 +1,176 @@ +from typing import Tuple +import torch +from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler +import os +import numpy as np +import cv2 +from PIL import Image + +import folder_paths +import warnings +warnings.filterwarnings("ignore", category=UserWarning, module="torch") +warnings.filterwarnings("ignore", category=UserWarning, module="safetensors") + + + +class BrushNetLoader: + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "brushnet": ([x[0][17:] for x in os.walk("./models/inpaint/")], ), + }, + } + + CATEGORY = "inpaint" + RETURN_TYPES = ("BRMODEL",) + RETURN_NAMES = ("brushnet",) + + FUNCTION = "brushnet_loading" + + def brushnet_loading(self, brushnet): + brushnet_path = "./models/inpaint/"+brushnet + print("Brushnet... ", end="") + brush_net = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch.float16) + print("loaded") + return (brush_net,) + + +class BrushNetPipeline: + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "model": (folder_paths.get_filename_list("checkpoints"), ), + "brushnet": ("BRMODEL", ), + }, + } + + CATEGORY = "inpaint" + RETURN_TYPES = ("BRPL",) + RETURN_NAMES = ("BRPL",) + + FUNCTION = "pipeline_loading" + + def pipeline_loading(self, model, brushnet): + model_path = "./models/checkpoints/"+model + print("Model... ", end="") + pipe = StableDiffusionBrushNetPipeline.from_single_file( + model_path, + brushnet=brushnet, + torch_dtype=torch.float16, + low_cpu_mem_usage=False + ) + pipe.to("cuda") + pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) + pipe.enable_model_cpu_offload() + print("loaded") + return (pipe,) + + + +class BrushNetInpaint: + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "BRPL": ("BRPL", ), + "image": ("IMAGE",), + "mask": ("MASK",), + "prompt": ("STRING", {"multiline": False, "default": ""}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), + "scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}), + }, + } + + CATEGORY = "inpaint" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "brushnet_inpaint" + + def brushnet_inpaint(self, BRPL, image: torch.Tensor, mask: torch.Tensor, prompt: str, seed: int, + steps: int, scale: float) -> Tuple[torch.Tensor]: + + + print("Working on image") + + init_image = image_from_tensor(image[0]) + mask_image = image_from_tensor(mask[0]) + + # resize image and mask + + mask_image = 1.*(mask_image>250)[:,:,np.newaxis] + init_image = init_image * (1-mask_image) + + init_image = Image.fromarray(init_image.astype(np.uint8)).convert("RGB") + mask_image = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB") + + print("Inference started") + + generator = torch.Generator("cuda").manual_seed(seed) + result = BRPL( + prompt, + init_image, + mask_image, + num_inference_steps=steps, + generator=generator, + brushnet_conditioning_scale=scale + ).images[0] + + return (torch.stack([numpy_to_tensor(np.array(result)).squeeze(0)]),) + + +class BlendInpaint: + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "mask": ("MASK",), + "blur": ("INT", {"default": 21, "min": 0, "max": 1000}), + }, + } + + CATEGORY = "inpaint" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "blend_inpaint" + + def blend_inpaint(self, image1: torch.Tensor, image2: torch.Tensor, mask: torch.Tensor, blur: int) -> Tuple[torch.Tensor]: + + init_image1 = image_from_tensor(image1[0]) + init_image2 = image_from_tensor(image2[0]) + mask_image = image_from_tensor(mask[0]) + + mask_image = 1.*(mask_image>250)[:,:,np.newaxis] + + # blur, you can adjust the parameters for better performance + mask_blurred = cv2.GaussianBlur(mask_image*255, (blur, blur), 0)/255 + mask_blurred = mask_blurred[:,:,np.newaxis] + mask_image = 1-(1-mask_image) * (1-mask_blurred) + + image_pasted=init_image1 * (1-mask_image) + init_image2*mask_image + result=image_pasted.astype(init_image1.dtype) + + return (torch.stack([numpy_to_tensor(np.array(result)).squeeze(0)]),) + + + +def image_from_tensor(t: torch.Tensor) -> np.ndarray: + image_np = t.numpy() + # Convert the numpy array back to the original range (0-255) and data type (uint8) + image_np = (image_np * 255).astype(np.uint8) + return image_np + +def numpy_to_tensor(array: np.ndarray) -> torch.Tensor: + """Convert a numpy array to a tensor and scale its values from 0-255 to 0-1.""" + array = array.astype(np.float32) / 255.0 + return torch.from_numpy(array)[None,] diff --git a/example/BrushNet example.json b/example/BrushNet example.json new file mode 100644 index 0000000..fd8d024 --- /dev/null +++ b/example/BrushNet example.json @@ -0,0 +1 @@ +{"last_node_id": 44, "last_link_id": 68, "nodes": [{"id": 29, "type": "GroundingDinoSAMSegment (segment anything)", "pos": [749, 640], "size": {"0": 352.79998779296875, "1": 124.0001220703125}, "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "sam_model", "type": "SAM_MODEL", "link": 55}, {"name": "grounding_dino_model", "type": "GROUNDING_DINO_MODEL", "link": 56}, {"name": "image", "type": "IMAGE", "link": 57}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [59, 61], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "GroundingDinoSAMSegment (segment anything)"}, "widgets_values": ["cake", 0.3]}, {"id": 27, "type": "GroundingDinoModelLoader (segment anything)", "pos": [45, 1126], "size": {"0": 532.800048828125, "1": 104.4000244140625}, "flags": {}, "order": 2, "mode": 0, "outputs": [{"name": "GROUNDING_DINO_MODEL", "type": "GROUNDING_DINO_MODEL", "links": [56], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "GroundingDinoModelLoader (segment anything)"}, "widgets_values": ["GroundingDINO_SwinT_OGC (694MB)"]}, {"id": 28, "type": "SAMModelLoader (segment anything)", "pos": [254, 1014], "size": {"0": 315, "1": 58}, "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "SAM_MODEL", "type": "SAM_MODEL", "links": [55], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "SAMModelLoader (segment anything)"}, "widgets_values": ["sam_vit_h (2.56GB)"]}, {"id": 1, "type": "LoadImage", "pos": [101, 386], "size": [470.19439697265625, 578.6853942871094], "flags": {}, "order": 3, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [52, 57], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": null, "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "LoadImage"}, "widgets_values": ["test_image.jpg", "image"]}, {"id": 24, "type": "BrushNetLoader", "pos": [222, 260], "size": {"0": 354.79998779296875, "1": 61.60003662109375}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "BrushNet", "type": "BRMODEL", "links": [67], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetLoader"}, "widgets_values": ["brushnet/random_mask_brushnet_ckpt"]}, {"id": 44, "type": "BrushNetPipeline", "pos": [704, 263], "size": [371.60003662109375, 103.20001220703125], "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "brushnet", "type": "BRMODEL", "link": 67}], "outputs": [{"name": "BRPL", "type": "BRPL", "links": [66], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetPipeline"}, "widgets_values": ["realisticVisionV60B1_v51VAE.safetensors"]}, {"id": 26, "type": "BrushNetInpaint", "pos": [1254, 411], "size": {"0": 315, "1": 194}, "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "BRPL", "type": "BRPL", "link": 66}, {"name": "image", "type": "IMAGE", "link": 52}, {"name": "mask", "type": "MASK", "link": 61}], "outputs": [{"name": "Image", "type": "IMAGE", "links": [53], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetInpaint"}, "widgets_values": ["burger", 0, "fixed", 30, 1]}, {"id": 12, "type": "PreviewImage", "pos": [1634, 411], "size": [523.5944213867188, 547.4853210449219], "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 53}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 32, "type": "PreviewImage", "pos": [1386, 711], "size": {"0": 210, "1": 246}, "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 60}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 31, "type": "MaskToImage", "pos": [1140, 711], "size": {"0": 210, "1": 26}, "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "mask", "type": "MASK", "link": 59}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [60], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "MaskToImage"}}], "links": [[52, 1, 0, 26, 1, "IMAGE"], [53, 26, 0, 12, 0, "IMAGE"], [55, 28, 0, 29, 0, "SAM_MODEL"], [56, 27, 0, 29, 1, "GROUNDING_DINO_MODEL"], [57, 1, 0, 29, 2, "IMAGE"], [59, 29, 1, 31, 0, "MASK"], [60, 31, 0, 32, 0, "IMAGE"], [61, 29, 1, 26, 2, "MASK"], [66, 44, 0, 26, 0, "BRPL"], [67, 24, 0, 44, 0, "BRMODEL"]], "groups": [], "config": {}, "extra": {}, "version": 0.4} \ No newline at end of file diff --git a/example/BrushNet example.png b/example/BrushNet example.png new file mode 100644 index 0000000..753053b Binary files /dev/null and b/example/BrushNet example.png differ diff --git a/example/test_image.jpg b/example/test_image.jpg new file mode 100644 index 0000000..dba55ac Binary files /dev/null and b/example/test_image.jpg differ diff --git a/example/test_mask.jpg b/example/test_mask.jpg new file mode 100644 index 0000000..686aa01 Binary files /dev/null and b/example/test_mask.jpg differ diff --git a/requirements.txt b/requirements.txt new file mode 100755 index 0000000..4aede95 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +torch +numpy +opencv-python +Pillow +accelerate +git+https://github.com/nullquant/diffusers.git