diff --git a/BIG_IMAGE.md b/BIG_IMAGE.md
new file mode 100644
index 0000000..84f735a
--- /dev/null
+++ b/BIG_IMAGE.md
@@ -0,0 +1,6 @@
+
+
+[workflow](example/BrushNet_cut_for_inpaint.json)
+
+When you work with big image and your inpaint mask is small it is better to cut part of the image, work with it and then blend it back.
+I created a node for such workflow, see example.
diff --git a/README.md b/README.md
index ea85c1c..aca71a0 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,8 @@ My contribution is limited to the ComfyUI adaptation, and all credit goes to the
## Updates
+May 12, 2024. CutForInpaint node, see [example](BIG_IMAGE.md).
+
May 11, 2024. Image batch is implemented. You can even add BrushNet to AnimateDiff vid2vid workflow, but they don't work together - they are different models and both try to patch UNet. Added some more examples.
May 6, 2024. PowerPaint v2 model is implemented. After update your workflow probably will not work. Don't panic! Check `end_at` parameter of BrushNode, if it equals 1, change it to some big number. Read about parameters in Usage section below.
@@ -169,6 +171,20 @@ To upscale you should use base model, not BrushNet. The same is true for conditi
+
+
+ Big image inpaint
+
+
+
+[workflow](example/BrushNet_cut_for_inpaint.json)
+
+When you work with big image and your inpaint mask is small it is better to cut part of the image, work with it and then blend it back.
+I created a node for such workflow, see example.
+
+
+
+
PowerPaint outpaint
diff --git a/__init__.py b/__init__.py
index b654901..9f6459d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,4 +1,4 @@
-from .brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, PowerPaintCLIPLoader, PowerPaint
+from .brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, PowerPaintCLIPLoader, PowerPaint, CutForInpaint
"""
@author: nullquant
@title: BrushNet
@@ -14,6 +14,7 @@ NODE_CLASS_MAPPINGS = {
"BlendInpaint": BlendInpaint,
"PowerPaintCLIPLoader": PowerPaintCLIPLoader,
"PowerPaint": PowerPaint,
+ "CutForInpaint": CutForInpaint,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -23,4 +24,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"BlendInpaint": "Blend Inpaint",
"PowerPaintCLIPLoader": "PowerPaint CLIP Loader",
"PowerPaint": "PowerPaint",
+ "CutForInpaint": "Cut For Inpaint",
}
diff --git a/brushnet_nodes.py b/brushnet_nodes.py
index 5b417fa..ca4163d 100644
--- a/brushnet_nodes.py
+++ b/brushnet_nodes.py
@@ -1,5 +1,6 @@
import torch
import torchvision.transforms as T
+import torch.nn.functional as F
import math
import os
@@ -450,6 +451,10 @@ class BlendInpaint:
"kernel": ("INT", {"default": 10, "min": 1, "max": 1000}),
"sigma": ("FLOAT", {"default": 10.0, "min": 0.01, "max": 1000}),
},
+ "optional":
+ {
+ "origin": ("VECTOR",),
+ },
}
CATEGORY = "inpaint"
@@ -458,32 +463,113 @@ class BlendInpaint:
FUNCTION = "blend_inpaint"
- def blend_inpaint(self, inpaint: torch.Tensor, original: torch.Tensor, mask: torch.Tensor, kernel: int, sigma:int) -> Tuple[torch.Tensor]:
+ def blend_inpaint(self, inpaint: torch.Tensor, original: torch.Tensor, mask, kernel: int, sigma:int, origin=None) -> Tuple[torch.Tensor]:
- # no batches over mask and original image
- if len(mask.shape) > 2:
- mask = mask[0]
- if len(original.shape) > 3:
- original = original[0]
+ original, mask = check_image_mask(original, mask, 'Blend Inpaint')
+
+ if len(inpaint.shape) < 4:
+ # image tensor shape should be [B, H, W, C], but batch somehow is missing
+ inpaint = inpaint[None,:,:,:]
+
+ if inpaint.shape[0] < original.shape[0]:
+ print("Blend Inpaint gets batch of original images (%d) but only (%d) inpaint images" % (original.shape[0], inpaint.shape[0]))
+ original= original[:inpaint.shape[0],:,:]
+ mask = mask[:inpaint.shape[0],:,:]
+
+ if inpaint.shape[0] > original.shape[0]:
+ # batch over inpaint
+ count = 0
+ original_list = []
+ mask_list = []
+ origin_list = []
+ while (count < inpaint.shape[0]):
+ for i in range(original.shape[0]):
+ original_list.append(original[i][None,:,:,:])
+ mask_list.append(mask[i][None,:,:])
+ if origin is not None:
+ origin_list.append(origin[i][None,:])
+ count += 1
+ if count >= inpaint.shape[0]:
+ break
+ original = torch.concat(original_list, dim=0)
+ mask = torch.concat(mask_list, dim=0)
+ if origin is not None:
+ origin = torch.concat(origin_list, dim=0)
if kernel % 2 == 0:
kernel += 1
transform = T.GaussianBlur(kernel_size=(kernel, kernel), sigma=(sigma, sigma))
- blurred_mask = transform(mask[None,None,:,:]).to(original.device).to(original.dtype)
-
- inpaint = torch.nn.functional.interpolate(
- inpaint.permute(0, 3, 1, 2),
- size=(
- original.shape[0],
- original.shape[1],
- )
- ).to(original.device).to(original.dtype)
ret = []
- for result in inpaint.permute(0, 2, 3, 1):
- ret.append(original * (1.0 - blurred_mask[0][0][:,:,None]) + result.to(original.device) * blurred_mask[0][0][:,:,None])
+ blurred = []
+ for i in range(inpaint.shape[0]):
+ if origin is None:
+ blurred_mask = transform(mask[i][None,None,:,:]).to(original.device).to(original.dtype)
+ blurred.append(blurred_mask[0])
- return (torch.stack(ret), blurred_mask[0],)
+ result = torch.nn.functional.interpolate(
+ inpaint[i][None,:,:,:].permute(0, 3, 1, 2),
+ size=(
+ original[i].shape[0],
+ original[i].shape[1],
+ )
+ ).permute(0, 2, 3, 1).to(original.device).to(original.dtype)
+ else:
+ # got mask from CutForInpaint
+ height, width, _ = original[i].shape
+ x0 = origin[i][0].item()
+ y0 = origin[i][1].item()
+
+ if mask[i].shape[0] < height or mask[i].shape[1] < width:
+ padded_mask = F.pad(input=mask[i], pad=(x0, width-x0-mask[i].shape[1],
+ y0, height-y0-mask[i].shape[0]), mode='constant', value=0)
+ else:
+ padded_mask = mask[i]
+ blurred_mask = transform(padded_mask[None,None,:,:]).to(original.device).to(original.dtype)
+ blurred.append(blurred_mask[0][0])
+
+ result = F.pad(input=inpaint[i], pad=(0, 0, x0, width-x0-inpaint[i].shape[1],
+ y0, height-y0-inpaint[i].shape[0]), mode='constant', value=0)
+ result = result[None,:,:,:].to(original.device).to(original.dtype)
+
+ ret.append(original[i] * (1.0 - blurred_mask[0][0][:,:,None]) + result[0] * blurred_mask[0][0][:,:,None])
+
+ return (torch.stack(ret), torch.stack(blurred), )
+
+
+class CutForInpaint:
+
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required":
+ {
+ "image": ("IMAGE",),
+ "mask": ("MASK",),
+ "width": ("INT", {"default": 512, "min": 64, "max": 2048}),
+ "height": ("INT", {"default": 512, "min": 64, "max": 2048}),
+ },
+ }
+
+ CATEGORY = "inpaint"
+ RETURN_TYPES = ("IMAGE","MASK","VECTOR",)
+ RETURN_NAMES = ("image","mask","origin",)
+
+ FUNCTION = "cut_for_inpaint"
+
+ def cut_for_inpaint(self, image: torch.Tensor, mask: torch.Tensor, width: int, height: int):
+
+ image, mask = check_image_mask(image, mask, 'BrushNet')
+
+ ret = []
+ msk = []
+ org = []
+ for i in range(image.shape[0]):
+ x0, y0, w, h = cut_with_mask(mask[i], width, height)
+ ret.append((image[i][y0:y0+h,x0:x0+w,:]))
+ msk.append((mask[i][y0:y0+h,x0:x0+w]))
+ org.append(torch.IntTensor([x0,y0]))
+
+ return (torch.stack(ret), torch.stack(msk), torch.stack(org), )
#### Utility function
@@ -556,8 +642,8 @@ def check_compatibilty(model, brushnet):
return (is_SDXL, is_PP)
-# Prepare image and mask
-def prepare_image(image, mask):
+
+def check_image_mask(image, mask, name):
if len(image.shape) < 4:
# image tensor shape should be [B, H, W, C], but batch somehow is missing
image = image[None,:,:,:]
@@ -571,18 +657,25 @@ def prepare_image(image, mask):
mask = mask[None,:,:]
if image.shape[0] > mask.shape[0]:
- print("BrushNet gets batch of images (%d) but only %d masks" % (image.shape[0], mask.shape[0]))
+ print(name, "gets batch of images (%d) but only %d masks" % (image.shape[0], mask.shape[0]))
if mask.shape[0] == 1:
- print("BrushNet will copy the mask to fill batch")
+ print(name, "will copy the mask to fill batch")
mask = torch.cat([mask] * image.shape[0], dim=0)
else:
- print("BrushNet will add empty masks to fill batch")
+ print(name, "will add empty masks to fill batch")
empty_mask = torch.zeros([image.shape[0] - mask.shape[0], mask.shape[1], mask.shape[2]])
mask = torch.cat([mask, empty_mask], dim=0)
elif image.shape[0] < mask.shape[0]:
- print("BrushNet gets batch of images (%d) but too many (%d) masks" % (image.shape[0], mask.shape[0]))
+ print(name, "gets batch of images (%d) but too many (%d) masks" % (image.shape[0], mask.shape[0]))
mask = mask[:image.shape[0],:,:]
+ return (image, mask)
+
+# Prepare image and mask
+def prepare_image(image, mask):
+
+ image, mask = check_image_mask(image, mask, 'BrushNet')
+
print("BrushNet image.shape =", image.shape, "mask.shape =", mask.shape)
if mask.shape[2] != image.shape[2] or mask.shape[1] != image.shape[1]:
@@ -596,6 +689,45 @@ def prepare_image(image, mask):
return (masked_image, mask)
+def cut_with_mask(mask, width, height):
+ iy, ix = (mask == 1).nonzero(as_tuple=True)
+ x_min = ix.min().item()
+ x_max = ix.max().item()
+ y_min = iy.min().item()
+ y_max = iy.max().item()
+
+ if x_max - x_min > width or y_max - y_min > height:
+ raise Exception("Mask is bigger than provided dimensions")
+
+ x_c = (x_min + x_max) / 2.0
+ y_c = (y_min + y_max) / 2.0
+
+ h0, w0 = mask.shape
+
+ width2 = width / 2.0
+ height2 = height / 2.0
+
+ if w0 <= width:
+ x0 = 0
+ w = w0
+ else:
+ x0 = max(0, x_c - width2)
+ w = width
+ if x0 + width > w0:
+ x0 = w0 - width
+
+ if h0 <= height:
+ y0 = 0
+ h = h0
+ else:
+ y0 = max(0, y_c - height2)
+ h = height
+ if y0 + height > h0:
+ y0 = h0 - height
+
+ return (int(x0), int(y0), int(w), int(h))
+
+
# Prepare conditioning_latents
@torch.inference_mode()
def get_image_latents(masked_image, mask, vae, scaling_factor):
diff --git a/example/BrushNet_cut_for_inpaint.json b/example/BrushNet_cut_for_inpaint.json
new file mode 100644
index 0000000..3f67f5f
--- /dev/null
+++ b/example/BrushNet_cut_for_inpaint.json
@@ -0,0 +1 @@
+{"last_node_id": 74, "last_link_id": 147, "nodes": [{"id": 45, "type": "BrushNetLoader", "pos": [8, 251], "size": {"0": 576.2000122070312, "1": 104}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "brushnet", "type": "BRMODEL", "links": [129], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetLoader"}, "widgets_values": ["brushnet/random_mask.safetensors", "float16"]}, {"id": 12, "type": "PreviewImage", "pos": [1963, 148], "size": [362.71480126953156, 313.34364410400406], "flags": {}, "order": 10, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 93}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 54, "type": "VAEDecode", "pos": [1964, 23], "size": {"0": 210, "1": 46}, "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 91}, {"name": "vae", "type": "VAE", "link": 92}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [93, 142], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 71, "type": "CutForInpaint", "pos": [756, 333], "size": {"0": 315, "1": 122}, "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "image", "type": "IMAGE", "link": 138}, {"name": "mask", "type": "MASK", "link": 139}], "outputs": [{"name": "image", "type": "IMAGE", "links": [140], "shape": 3, "slot_index": 0}, {"name": "mask", "type": "MASK", "links": [141, 144], "shape": 3, "slot_index": 1}, {"name": "origin", "type": "VECTOR", "links": [145], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "CutForInpaint"}, "widgets_values": [512, 512]}, {"id": 58, "type": "LoadImage", "pos": [10, 404], "size": [695.4412421875002, 781.0468775024417], "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [138, 143, 147], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [139], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "LoadImage"}, "widgets_values": ["clipspace/clipspace-mask-2517487.png [input]", "image"]}, {"id": 52, "type": "KSampler", "pos": [1617, 131], "size": {"0": 315, "1": 262}, "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 121}, {"name": "positive", "type": "CONDITIONING", "link": 122}, {"name": "negative", "type": "CONDITIONING", "link": 123}, {"name": "latent_image", "type": "LATENT", "link": 124, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [91], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSampler"}, "widgets_values": [0, "fixed", 20, 8, "euler_ancestral", "normal", 1]}, {"id": 74, "type": "Reroute", "pos": [736.4412231445312, 336.93898856946777], "size": [75, 26], "flags": {}, "order": 3, "mode": 0, "inputs": [{"name": "", "type": "*", "link": 147}], "outputs": [{"name": "", "type": "IMAGE", "links": null}], "properties": {"showOutputText": false, "horizontal": false}}, {"id": 62, "type": "BrushNet", "pos": [1254, 134], "size": {"0": 315, "1": 226}, "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 125}, {"name": "vae", "type": "VAE", "link": 126}, {"name": "image", "type": "IMAGE", "link": 140}, {"name": "mask", "type": "MASK", "link": 141}, {"name": "brushnet", "type": "BRMODEL", "link": 129}, {"name": "positive", "type": "CONDITIONING", "link": 127}, {"name": "negative", "type": "CONDITIONING", "link": 128}], "outputs": [{"name": "model", "type": "MODEL", "links": [121], "shape": 3, "slot_index": 0}, {"name": "positive", "type": "CONDITIONING", "links": [122], "shape": 3, "slot_index": 1}, {"name": "negative", "type": "CONDITIONING", "links": [123], "shape": 3, "slot_index": 2}, {"name": "latent", "type": "LATENT", "links": [124], "shape": 3, "slot_index": 3}], "properties": {"Node name for S&R": "BrushNet"}, "widgets_values": [0.8, 0, 10000]}, {"id": 50, "type": "CLIPTextEncode", "pos": [651, 168], "size": {"0": 339.20001220703125, "1": 96.39999389648438}, "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 80}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [128], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": [""], "color": "#322", "bgcolor": "#533"}, {"id": 47, "type": "CheckpointLoaderSimple", "pos": [3, 44], "size": {"0": 481, "1": 158}, "flags": {}, "order": 2, "mode": 0, "outputs": [{"name": "MODEL", "type": "MODEL", "links": [125], "shape": 3, "slot_index": 0}, {"name": "CLIP", "type": "CLIP", "links": [78, 80], "shape": 3, "slot_index": 1}, {"name": "VAE", "type": "VAE", "links": [92, 126], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, "widgets_values": ["SD15/toonyou_beta6.safetensors"]}, {"id": 49, "type": "CLIPTextEncode", "pos": [649, 21], "size": {"0": 339.20001220703125, "1": 96.39999389648438}, "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 78}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [127], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": ["a clear blue sky"], "color": "#232", "bgcolor": "#353"}, {"id": 72, "type": "BlendInpaint", "pos": [1385, 616], "size": {"0": 315, "1": 142}, "flags": {}, "order": 11, "mode": 0, "inputs": [{"name": "inpaint", "type": "IMAGE", "link": 142}, {"name": "original", "type": "IMAGE", "link": 143}, {"name": "mask", "type": "MASK", "link": 144}, {"name": "origin", "type": "VECTOR", "link": 145}], "outputs": [{"name": "image", "type": "IMAGE", "links": [146], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": null, "shape": 3}], "properties": {"Node name for S&R": "BlendInpaint"}, "widgets_values": [10, 10]}, {"id": 73, "type": "PreviewImage", "pos": [1784, 511], "size": [578.8481262207033, 616.0670013427734], "flags": {}, "order": 12, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 146}], "properties": {"Node name for S&R": "PreviewImage"}}], "links": [[78, 47, 1, 49, 0, "CLIP"], [80, 47, 1, 50, 0, "CLIP"], [91, 52, 0, 54, 0, "LATENT"], [92, 47, 2, 54, 1, "VAE"], [93, 54, 0, 12, 0, "IMAGE"], [121, 62, 0, 52, 0, "MODEL"], [122, 62, 1, 52, 1, "CONDITIONING"], [123, 62, 2, 52, 2, "CONDITIONING"], [124, 62, 3, 52, 3, "LATENT"], [125, 47, 0, 62, 0, "MODEL"], [126, 47, 2, 62, 1, "VAE"], [127, 49, 0, 62, 5, "CONDITIONING"], [128, 50, 0, 62, 6, "CONDITIONING"], [129, 45, 0, 62, 4, "BRMODEL"], [138, 58, 0, 71, 0, "IMAGE"], [139, 58, 1, 71, 1, "MASK"], [140, 71, 0, 62, 2, "IMAGE"], [141, 71, 1, 62, 3, "MASK"], [142, 54, 0, 72, 0, "IMAGE"], [143, 58, 0, 72, 1, "IMAGE"], [144, 71, 1, 72, 2, "MASK"], [145, 71, 2, 72, 3, "VECTOR"], [146, 72, 0, 73, 0, "IMAGE"], [147, 58, 0, 74, 0, "*"]], "groups": [], "config": {}, "extra": {}, "version": 0.4}
\ No newline at end of file
diff --git a/example/BrushNet_cut_for_inpaint.png b/example/BrushNet_cut_for_inpaint.png
new file mode 100644
index 0000000..85e36f4
Binary files /dev/null and b/example/BrushNet_cut_for_inpaint.png differ