nodes/nullquant-ComfyUI-BrushNet
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cut for inpaint node

Browse Source
This commit is contained in:
nullquant
2024-05-12 22:27:43 +03:00
parent 2e37308cee
commit f8a08c4a84
6 changed files with 182 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
BIG_IMAGE.md
+6
View File
@@ -0,0 +1,6 @@
![example workflow](example/BrushNet_cut_for_inpaint.png?raw=true)
[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.
README.md
+16
View File
@@ -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
</details>
<details>
<summary>Big image inpaint</summary>
![example workflow](example/BrushNet_cut_for_inpaint.png?raw=true)
[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.
</details>
<details>
<summary>PowerPaint outpaint</summary>
__init__.py
+3 -1
View File
@@ -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",
}
brushnet_nodes.py
+156 -24
View File
@@ -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):
example/BrushNet_cut_for_inpaint.json
+1
View File
File diff suppressed because one or more lines are too long
example/BrushNet_cut_for_inpaint.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.3 MiB