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

batches, example workflows

Browse Source
This commit is contained in:
Kijai
2024-04-12 16:50:54 +03:00
parent 1eb9315a14
commit a1f1bcb0be
4 changed files with 1279 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
brushnet/pipeline_brushnet.py
+3 -1
View File
@@ -27,7 +27,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusio
from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from comfy.utils import ProgressBar as comfy_pbar
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
@@ -1164,6 +1164,7 @@ class StableDiffusionBrushNetPipeline(
is_unet_compiled = is_compiled_module(self.unet)
is_brushnet_compiled = is_compiled_module(self.brushnet)
is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")
comfy_pbar(num_inference_steps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
# Relevant thread:
@@ -1245,6 +1246,7 @@ class StableDiffusionBrushNetPipeline(
# call the callback, if provided
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
comfy_pbar(1)
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
examples/brushnet_example_outpaint.json
+481
View File
@@ -0,0 +1,481 @@
{
"last_node_id": 54,
"last_link_id": 109,
"nodes": [
{
"id": 3,
"type": "CheckpointLoaderSimple",
"pos": [
202,
402
],
"size": {
"0": 351.8843078613281,
"1": 98
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
1
],
"shape": 3
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
2
],
"shape": 3,
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
3
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"1_5/darkSushi25D25D_v40.safetensors"
]
},
{
"id": 7,
"type": "LoadImage",
"pos": [
215,
582
],
"size": {
"0": 316,
"1": 405
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
39
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": [],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"clipspace/clipspace-mask-1261422.png [input]",
"image"
]
},
{
"id": 1,
"type": "brushnet_model_loader",
"pos": [
688,
402
],
"size": {
"0": 337,
"1": 98
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 1,
"slot_index": 0
},
{
"name": "clip",
"type": "CLIP",
"link": 2
},
{
"name": "vae",
"type": "VAE",
"link": 3
}
],
"outputs": [
{
"name": "brushnet",
"type": "BRUSHNET",
"links": [
4
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "brushnet_model_loader"
},
"widgets_values": [
"brushnet_segmentation_mask"
]
},
{
"id": 24,
"type": "ImageResize+",
"pos": [
578,
591
],
"size": {
"0": 315,
"1": 218
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 39
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
99
],
"shape": 3,
"slot_index": 0
},
{
"name": "width",
"type": "INT",
"links": null,
"shape": 3
},
{
"name": "height",
"type": "INT",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "ImageResize+"
},
"widgets_values": [
512,
512,
"lanczos",
true,
"always",
2
]
},
{
"id": 52,
"type": "ImagePadForOutpaint",
"pos": [
918,
598
],
"size": {
"0": 315,
"1": 174
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 99
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
100,
108
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": [
101,
109
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "ImagePadForOutpaint"
},
"widgets_values": [
256,
256,
256,
256,
1
]
},
{
"id": 53,
"type": "PreviewImage",
"pos": [
928,
824
],
"size": [
210,
246
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 108
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 54,
"type": "MaskPreview+",
"pos": [
1148,
826
],
"size": [
210,
246
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 109
}
],
"properties": {
"Node name for S&R": "MaskPreview+"
}
},
{
"id": 5,
"type": "brushnet_sampler",
"pos": [
1283,
404
],
"size": {
"0": 399,
"1": 299
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "brushnet",
"type": "BRUSHNET",
"link": 4,
"slot_index": 0
},
{
"name": "image",
"type": "IMAGE",
"link": 100,
"slot_index": 1
},
{
"name": "mask",
"type": "MASK",
"link": 101
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": [
86
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "brushnet_sampler"
},
"widgets_values": [
30,
1,
29,
"fixed",
"UniPCMultistepScheduler",
"1girl, blue dress, forest, best quality, masterpiece"
]
},
{
"id": 47,
"type": "PreviewImage",
"pos": [
1728,
406
],
"size": [
555.6796875,
582.3743743896484
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 86,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
}
],
"links": [
[
1,
3,
0,
1,
0,
"MODEL"
],
[
2,
3,
1,
1,
1,
"CLIP"
],
[
3,
3,
2,
1,
2,
"VAE"
],
[
4,
1,
0,
5,
0,
"BRUSHNET"
],
[
39,
7,
0,
24,
0,
"IMAGE"
],
[
86,
5,
0,
47,
0,
"IMAGE"
],
[
99,
24,
0,
52,
0,
"IMAGE"
],
[
100,
52,
0,
5,
1,
"IMAGE"
],
[
101,
52,
1,
5,
2,
"MASK"
],
[
108,
52,
0,
53,
0,
"IMAGE"
],
[
109,
52,
1,
54,
0,
"MASK"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
examples/brushnet_example_workflow_blend.json
+767
View File
@@ -0,0 +1,767 @@
{
"last_node_id": 51,
"last_link_id": 98,
"nodes": [
{
"id": 3,
"type": "CheckpointLoaderSimple",
"pos": [
202,
402
],
"size": {
"0": 351.8843078613281,
"1": 98
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
1
],
"shape": 3
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
2
],
"shape": 3,
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
3
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"1_5/darkSushi25D25D_v40.safetensors"
]
},
{
"id": 24,
"type": "ImageResize+",
"pos": [
578,
591
],
"size": {
"0": 315,
"1": 218
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 39
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
81,
87,
90
],
"shape": 3,
"slot_index": 0
},
{
"name": "width",
"type": "INT",
"links": null,
"shape": 3
},
{
"name": "height",
"type": "INT",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "ImageResize+"
},
"widgets_values": [
512,
512,
"lanczos",
true,
"always",
2
]
},
{
"id": 51,
"type": "RemapMaskRange",
"pos": [
1240,
885
],
"size": {
"0": 315,
"1": 82
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 94
}
],
"outputs": [
{
"name": "mask",
"type": "MASK",
"links": [
95,
96
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "RemapMaskRange"
},
"widgets_values": [
0,
0.8
]
},
{
"id": 50,
"type": "MaskToImage",
"pos": [
1251,
1020
],
"size": {
"0": 210,
"1": 26
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 96
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
92
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "MaskToImage"
}
},
{
"id": 49,
"type": "ImageCompositeMasked",
"pos": [
1251,
1093
],
"size": {
"0": 315,
"1": 146
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "destination",
"type": "IMAGE",
"link": 90
},
{
"name": "source",
"type": "IMAGE",
"link": 92
},
{
"name": "mask",
"type": "MASK",
"link": 89,
"slot_index": 2
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
93
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageCompositeMasked"
},
"widgets_values": [
0,
0,
true
]
},
{
"id": 46,
"type": "PreviewImage",
"pos": [
2124,
794
],
"size": {
"0": 491.6383056640625,
"1": 480.1783142089844
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 85,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 6,
"type": "PreviewImage",
"pos": [
1593,
795
],
"size": {
"0": 491.6383056640625,
"1": 480.1783142089844
},
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 93,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 44,
"type": "ImageCompositeMasked",
"pos": [
1722,
568
],
"size": {
"0": 315,
"1": 146
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "destination",
"type": "IMAGE",
"link": 81
},
{
"name": "source",
"type": "IMAGE",
"link": 82
},
{
"name": "mask",
"type": "MASK",
"link": 95,
"slot_index": 2
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
85
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageCompositeMasked"
},
"widgets_values": [
0,
0,
true
]
},
{
"id": 48,
"type": "MaskPreview+",
"pos": [
987,
991
],
"size": [
210,
246
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 88
}
],
"properties": {
"Node name for S&R": "MaskPreview+"
}
},
{
"id": 45,
"type": "GrowMaskWithBlur",
"pos": [
638,
962
],
"size": {
"0": 315,
"1": 246
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 84
}
],
"outputs": [
{
"name": "mask",
"type": "MASK",
"links": [
88,
89,
94
],
"shape": 3,
"slot_index": 0
},
{
"name": "mask_inverted",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "GrowMaskWithBlur"
},
"widgets_values": [
2,
0,
true,
false,
25.5,
1,
1,
false
]
},
{
"id": 7,
"type": "LoadImage",
"pos": [
215,
582
],
"size": {
"0": 316,
"1": 405
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
39
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": [
84,
98
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"clipspace/clipspace-mask-381446.png [input]",
"image"
]
},
{
"id": 47,
"type": "PreviewImage",
"pos": [
1711,
147
],
"size": [
361.57662353515616,
361.1871826171875
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 86,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 1,
"type": "brushnet_model_loader",
"pos": [
688,
402
],
"size": {
"0": 337,
"1": 98
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 1,
"slot_index": 0
},
{
"name": "clip",
"type": "CLIP",
"link": 2
},
{
"name": "vae",
"type": "VAE",
"link": 3
}
],
"outputs": [
{
"name": "brushnet",
"type": "BRUSHNET",
"links": [
4
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "brushnet_model_loader"
},
"widgets_values": [
"brushnet_segmentation_mask"
]
},
{
"id": 5,
"type": "brushnet_sampler",
"pos": [
1250,
401
],
"size": {
"0": 399,
"1": 299
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "brushnet",
"type": "BRUSHNET",
"link": 4,
"slot_index": 0
},
{
"name": "image",
"type": "IMAGE",
"link": 87,
"slot_index": 1
},
{
"name": "mask",
"type": "MASK",
"link": 98
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": [
82,
86
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "brushnet_sampler"
},
"widgets_values": [
30,
1,
27,
"fixed",
"UniPCMultistepScheduler",
"purple eye"
]
}
],
"links": [
[
1,
3,
0,
1,
0,
"MODEL"
],
[
2,
3,
1,
1,
1,
"CLIP"
],
[
3,
3,
2,
1,
2,
"VAE"
],
[
4,
1,
0,
5,
0,
"BRUSHNET"
],
[
39,
7,
0,
24,
0,
"IMAGE"
],
[
81,
24,
0,
44,
0,
"IMAGE"
],
[
82,
5,
0,
44,
1,
"IMAGE"
],
[
84,
7,
1,
45,
0,
"MASK"
],
[
85,
44,
0,
46,
0,
"IMAGE"
],
[
86,
5,
0,
47,
0,
"IMAGE"
],
[
87,
24,
0,
5,
1,
"IMAGE"
],
[
88,
45,
0,
48,
0,
"MASK"
],
[
89,
45,
0,
49,
2,
"MASK"
],
[
90,
24,
0,
49,
0,
"IMAGE"
],
[
92,
50,
0,
49,
1,
"IMAGE"
],
[
93,
49,
0,
6,
0,
"IMAGE"
],
[
94,
45,
0,
51,
0,
"MASK"
],
[
95,
51,
0,
44,
2,
"MASK"
],
[
96,
51,
0,
50,
0,
"MASK"
],
[
98,
7,
1,
5,
2,
"MASK"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
nodes.py
+28 -12
View File
@@ -1,6 +1,8 @@
import os
from contextlib import nullcontext
import torch
import torch.nn.functional as F
try:
from diffusers import (
DPMSolverMultistepScheduler,
@@ -83,13 +85,19 @@ class brushnet_model_loader:
if not os.path.exists(checkpoint_path):
print(f"Selected model: {checkpoint_path} not found, downloading...")
from huggingface_hub import snapshot_download
snapshot_download(repo_id="Kijai/BrushNet-fp16", allow_patterns=[f"*{brushnet_model}*"], local_dir=brushnet_model_folder, local_dir_use_symlinks=False)
snapshot_download(repo_id="Kijai/BrushNet-fp16",
allow_patterns=[f"*{brushnet_model}*"],
local_dir=brushnet_model_folder,
local_dir_use_symlinks=False
)
brushnet = BrushNetModel(**brushnet_config)
brushnet_sd = comfy.utils.load_torch_file(checkpoint_path)
brushnet.load_state_dict(brushnet_sd)
brushnet.to(dtype)
pbar.update(1)
clip_sd = None
load_models = [model]
load_models.append(clip.load_model())
@@ -131,10 +139,6 @@ class brushnet_model_loader:
scheduler=DPMSolverMultistepScheduler(**scheduler_config)
pbar.update(1)
del sd
pbar.update(1)
print("creating pipeline")
self.pipe = StableDiffusionBrushNetPipeline(
unet=unet,
@@ -147,9 +151,8 @@ class brushnet_model_loader:
safety_checker=None,
feature_extractor=None
)
print("pipeline created")
pbar.update(1)
brushnet = {
"pipe": self.pipe,
}
@@ -229,18 +232,31 @@ class brushnet_sampler:
noise_scheduler = UniPCMultistepScheduler(**scheduler_config)
pipe.scheduler = noise_scheduler
B, H, W, C = image.shape
image = image.permute(0, 3, 1, 2).to(device)
mask = mask.unsqueeze(0).to(device)
image = image * (1-mask)
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
mask = mask.to(device)
if mask.shape[0] < B:
repeat_times = B // mask.shape[0]
mask = mask.repeat(repeat_times, 1, 1, 1)
resized_mask = F.interpolate(mask.unsqueeze(1), size=[H, W], mode='nearest').squeeze(1)
image = image * (1-resized_mask)
prompt_list = []
prompt_list.append(prompt)
if len(prompt_list) < B:
prompt_list += [prompt_list[-1]] * (B - len(prompt_list))
autocast_condition = (dtype != torch.float32) and not mm.is_device_mps(device)
with torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext():
generator = torch.Generator(device).manual_seed(seed)
images = pipe(
prompt,
prompt_list,
image=image,
mask=mask,
mask=resized_mask,
num_inference_steps=steps,
generator=generator,
brushnet_conditioning_scale=guidance_scale,