diff --git a/brushnet/pipeline_brushnet.py b/brushnet/pipeline_brushnet.py index 41520d2..8e00e97 100644 --- a/brushnet/pipeline_brushnet.py +++ b/brushnet/pipeline_brushnet.py @@ -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) diff --git a/examples/brushnet_example_outpaint.json b/examples/brushnet_example_outpaint.json new file mode 100644 index 0000000..e8c762d --- /dev/null +++ b/examples/brushnet_example_outpaint.json @@ -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 +} \ No newline at end of file diff --git a/examples/brushnet_example_workflow_blend.json b/examples/brushnet_example_workflow_blend.json new file mode 100644 index 0000000..932f6a6 --- /dev/null +++ b/examples/brushnet_example_workflow_blend.json @@ -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 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py index 5b36ba2..d93f0df 100644 --- a/nodes.py +++ b/nodes.py @@ -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,