diff --git a/README.md b/README.md index a2d37d1..2a22791 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ My contribution is limited to the ComfyUI adaptation, and all credit goes to the ## Updates -May 9, 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 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. @@ -158,6 +158,15 @@ To upscale you should use base model, not BrushNet. The same is true for conditi +
+ Image batch + +![example workflow](example/BrushNet_image_batch.png?raw=true) + +[workflow](example/BrushNet_image_batch.json) + +
+
PowerPaint outpaint diff --git a/brushnet_nodes.py b/brushnet_nodes.py index d0c2f31..65dfa55 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -1,5 +1,6 @@ import torch import torchvision.transforms as T +import math import os import folder_paths @@ -10,8 +11,6 @@ from sys import platform comfy_parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')) sys.path.append(comfy_parent_dir) import comfy -import nodes -#import latent_preview from accelerate import init_empty_weights, load_checkpoint_and_dispatch @@ -34,6 +33,11 @@ powerpaint_config_file = os.path.join(current_directory,'brushnet', 'powerpaint. sd15_scaling_factor = 0.18215 sdxl_scaling_factor = 0.13025 +ComfySDLayers = [comfy.ops.disable_weight_init.Conv2d, + comfy.ldm.modules.attention.SpatialTransformer, + comfy.ldm.modules.diffusionmodules.openaimodel.Downsample, + comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock + ] class BrushNetLoader: @@ -220,9 +224,9 @@ class PowerPaint: torch_dtype = powerpaint['dtype'] # prepare conditioning latents - conditioning_latents, conditioning_latents2 = get_image_latents(masked_image, mask, vae, scaling_factor) - conditioning_latents = conditioning_latents.to(dtype=torch_dtype).to(powerpaint['brushnet'].device) - conditioning_latents2 = conditioning_latents2.to(dtype=torch_dtype).to(powerpaint['brushnet'].device) + conditioning_latents = get_image_latents(masked_image, mask, vae, scaling_factor) + conditioning_latents[0] = conditioning_latents[0].to(dtype=torch_dtype).to(powerpaint['brushnet'].device) + conditioning_latents[1] = conditioning_latents[1].to(dtype=torch_dtype).to(powerpaint['brushnet'].device) # prepare embeddings @@ -271,9 +275,7 @@ class PowerPaint: negative_prompt_embedsB = clip.encode_from_tokens(tokens, return_pooled=False) prompt_embeds_pp = (prompt_embedsA * fitting + (1.0 - fitting) * prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device) - negative_prompt_embeds = (negative_prompt_embedsA * fitting + (1.0 - fitting) * negative_prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device) - - prompt_embeds_pp2 = torch.cat([negative_prompt_embeds, prompt_embeds_pp]) + negative_prompt_embeds_pp = (negative_prompt_embedsA * fitting + (1.0 - fitting) * negative_prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device) # apply patch to code @@ -291,13 +293,11 @@ class PowerPaint: powerpaint['brushnet'], torch_dtype, conditioning_latents, - conditioning_latents2, (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), - prompt_embeds_pp, None, None, - prompt_embeds_pp2, None, None) + prompt_embeds_pp, negative_prompt_embeds_pp, + None, None, None) - latent = torch.zeros([batch, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=powerpaint['brushnet'].device) - #latent = torch.randn([1, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=brushnet['brushnet'].device) + latent = torch.zeros([batch, 4, conditioning_latents[0].shape[2], conditioning_latents[0].shape[3]], device=powerpaint['brushnet'].device) return (model, positive, negative, {"samples":latent},) @@ -355,22 +355,22 @@ class BrushNet: torch_dtype = brushnet['dtype'] # prepare conditioning latents - conditioning_latents, conditioning_latents2 = get_image_latents(masked_image, mask, vae, scaling_factor) - conditioning_latents = conditioning_latents.to(dtype=torch_dtype).to(brushnet['brushnet'].device) - conditioning_latents2 = conditioning_latents2.to(dtype=torch_dtype).to(brushnet['brushnet'].device) + conditioning_latents = get_image_latents(masked_image, mask, vae, scaling_factor) + conditioning_latents[0] = conditioning_latents[0].to(dtype=torch_dtype).to(brushnet['brushnet'].device) + conditioning_latents[1] = conditioning_latents[1].to(dtype=torch_dtype).to(brushnet['brushnet'].device) # prepare embeddings - prompt_embeds = positive[0][0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device) - negative_prompt_embeds = negative[0][0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device) + prompt_embeds = positive[0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device) + negative_prompt_embeds = negative[0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device) if len(positive[0]) > 1 and 'pooled_output' in positive[0][1] and positive[0][1]['pooled_output'] is not None: - add_text_embeds = positive[0][1]['pooled_output'].to(dtype=torch_dtype).to(brushnet['brushnet'].device) + pooled_prompt_embeds = positive[0][1]['pooled_output'].to(dtype=torch_dtype).to(brushnet['brushnet'].device) else: print('BrushNet: positive conditioning has not pooled_output') if is_SDXL: print('BrushNet will not produce correct results') - add_text_embeds = torch.empty([2, 1280], device=brushnet['brushnet'].device).to(dtype=torch_dtype) + pooled_prompt_embeds = torch.empty([2, 1280], device=brushnet['brushnet'].device).to(dtype=torch_dtype) if len(negative[0]) > 1 and 'pooled_output' in negative[0][1] and negative[0][1]['pooled_output'] is not None: negative_pooled_prompt_embeds = negative[0][1]['pooled_output'].to(dtype=torch_dtype).to(brushnet['brushnet'].device) @@ -378,14 +378,14 @@ class BrushNet: print('BrushNet: negative conditioning has not pooled_output') if is_SDXL: print('BrushNet will not produce correct results') - negative_pooled_prompt_embeds = torch.empty([1, add_text_embeds.shape[1]], device=brushnet['brushnet'].device).to(dtype=torch_dtype) + negative_pooled_prompt_embeds = torch.empty([1, pooled_prompt_embeds.shape[1]], device=brushnet['brushnet'].device).to(dtype=torch_dtype) - add_time_ids = torch.FloatTensor([[height, width, 0., 0., height, width]]).to(dtype=torch_dtype).to(brushnet['brushnet'].device) - negative_add_time_ids = add_time_ids + time_ids = torch.FloatTensor([[height, width, 0., 0., height, width]]).to(dtype=torch_dtype).to(brushnet['brushnet'].device) - prompt_embeds2 = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - add_text_embeds2 = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) - add_time_ids2 = torch.cat([negative_add_time_ids, add_time_ids], dim=0) + if not is_SDXL: + pooled_prompt_embeds = None + negative_pooled_prompt_embeds = None + time_ids = None # apply patch to code @@ -403,13 +403,11 @@ class BrushNet: brushnet['brushnet'], torch_dtype, conditioning_latents, - conditioning_latents2, (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), - prompt_embeds, add_text_embeds, add_time_ids, - prompt_embeds2, add_text_embeds2, add_time_ids2) + prompt_embeds, negative_prompt_embeds, + pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids) - latent = torch.zeros([batch, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=brushnet['brushnet'].device) - #latent = torch.randn([1, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=brushnet['brushnet'].device) + latent = torch.zeros([batch, 4, conditioning_latents[0].shape[2], conditioning_latents[0].shape[3]], device=brushnet['brushnet'].device) return (model, positive, negative, {"samples":latent},) @@ -527,10 +525,6 @@ def prepare_image(image, mask): if len(image.shape) < 4: # image tensor shape should be [B, H, W, C], but batch somehow is missing image = image[None,:,:,:] - #elif image.shape[0] > 1: - # print("") - # # batch > 1, take first image - # image = image[0][None,:,:,:] if len(mask.shape) > 3: # mask tensor shape should be [B, H, W] but we get [B, H, W, C], image may be? @@ -539,9 +533,6 @@ def prepare_image(image, mask): elif len(mask.shape) < 3: # mask tensor shape should be [B, H, W] but batch somehow is missing mask = mask[None,:,:] - #elif mask.shape[0] > 1: - # # batch > 1, take first mask - # mask = mask[0][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])) @@ -573,11 +564,7 @@ def prepare_image(image, mask): @torch.inference_mode() def get_image_latents(masked_image, mask, vae, scaling_factor): processed_image = masked_image.to(vae.device) - processed_image2 = torch.cat([masked_image] * 2).to(vae.device) - image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor - image_latents2 = vae.encode(processed_image2[:,:,:,:3]) * scaling_factor - processed_mask = 1. - mask[:,None,:,:] interpolated_mask = torch.nn.functional.interpolate( processed_mask, @@ -588,63 +575,12 @@ def get_image_latents(masked_image, mask, vae, scaling_factor): ) interpolated_mask = interpolated_mask.to(image_latents.device) - processed_mask2 = torch.cat([1. - mask[:,None,:,:]] * 2) - interpolated_mask2 = torch.nn.functional.interpolate( - processed_mask2, - size=( - image_latents2.shape[-2], - image_latents2.shape[-1] - ) - ) - interpolated_mask2 = interpolated_mask2.to(image_latents2.device) + conditioning_latents = [image_latents, interpolated_mask] - conditioning_latents = torch.concat([image_latents, interpolated_mask], 1) - conditioning_latents2 = torch.concat([image_latents2, interpolated_mask2], 1) + print('BrushNet CL: image_latents shape =', image_latents.shape, 'interpolated_mask shape =', interpolated_mask.shape) - print('BrushNet CL: image_latents shape =', image_latents2.shape, 'interpolated_mask shape =', interpolated_mask2.shape) + return conditioning_latents - return (conditioning_latents, conditioning_latents2) - - -#def modified_common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, -# disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): -# ''' -# #Modified by BrushNet nodes -# ''' -# latent_image = latent["samples"] -# if disable_noise: -# noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") -# else: -# batch_inds = latent["batch_index"] if "batch_index" in latent else None -# noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds) -# -# noise_mask = None -# if "noise_mask" in latent: -# noise_mask = latent["noise_mask"] -# -# ####################################################################################### -# # -# latent_preview_callback = latent_preview.prepare_callback(model, steps) -# -# to = add_model_patch_option(model) -# to['model_patch']['step'] = 0 -# to['model_patch']['total_steps'] = steps -# to['model_patch']['cfg'] = cfg -# -# def callback(step, x0, x, total_steps): -# to['model_patch']['step'] = step + 1 -# latent_preview_callback(steps, x0, x, total_steps) -# # -# ####################################################################################### -# -# disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED -# samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, -# denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step, -# force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, -# disable_pbar=disable_pbar, seed=seed) -# out = latent.copy() -# out["samples"] = samples -# return (out, ) # Model needs current step number and cfg at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one. # The first versions had modified_common_ksampler, but it broke custom KSampler nodes @@ -661,12 +597,17 @@ def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigm # to = add_model_patch_option(model) to['model_patch']['all_sigmas'] = sigmas - to['model_patch']['cfg'] = cfg + + if math.isclose(cfg, 1.0) and model_options.get("disable_cfg1_optimization", False) == False: + to['model_patch']['free_guidance'] = False + else: + to['model_patch']['free_guidance'] = True # ####################################################################################### return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) + # Main function where magic happens @torch.inference_mode() def brushnet_inference(x, timesteps, transformer_options): @@ -680,18 +621,13 @@ def brushnet_inference(x, timesteps, transformer_options): brushnet = mp['brushnet_model'] if isinstance(brushnet, BrushNetModel) or isinstance(brushnet, PowerPaintModel): torch_dtype = mp['brushnet_dtype'] - cl = mp['brushnet_latents'] - cl2 = mp['brushnet_latents2'] + cl_list = mp['brushnet_latents'] brushnet_conditioning_scale, control_guidance_start, control_guidance_end = mp['brushnet_controls'] pe = mp['brushnet_prompt'] - pe2 = mp['brushnet_prompt2'] - ack = mp['brushnet_add_embeds'] - ack2 = mp['brushnet_add_embeds2'] - cfg = mp['cfg'] + npe = mp['brushnet_negative_prompt'] + ppe, nppe, time_ids = mp['brushnet_add_embeds'] - sub_idx = None - if 'ad_params' in transformer_options and 'sub_idxs' in transformer_options['ad_params']: - sub_idx = transformer_options['ad_params']['sub_idxs'] + do_classifier_free_guidance = mp['free_guidance'] x = x.detach().clone() x = x.to(torch_dtype).to(brushnet.device) @@ -706,66 +642,112 @@ def brushnet_inference(x, timesteps, transformer_options): added_cond_kwargs = {} - do_classifier_free_guidance = cfg > 1 - if do_classifier_free_guidance and x.shape[0] % 2 != 0: - if step == 0: - print('Do_classifier_free_guidance and x.shape', x.shape, ', set DCFG = False') - do_classifier_free_guidance = False - if do_classifier_free_guidance and step == 0: print('BrushNet inference: do_classifier_free_guidance is True') - if do_classifier_free_guidance: - prompt_embeds = pe2 - if sub_idx: - idx = [] - for i in sub_idx: - idx.extend([2 * i, 2 * i + 1]) - cl2apply = torch.index_select(cl2, 0, torch.IntTensor(idx).to(cl2.device)) - else: - cl2apply = cl2 - else: - prompt_embeds = pe - if sub_idx: - cl2apply = torch.index_select(cl, 0, torch.IntTensor(sub_idx).to(cl2.device)) - else: - cl2apply = cl + sub_idx = None + if 'ad_params' in transformer_options and 'sub_idxs' in transformer_options['ad_params']: + sub_idx = transformer_options['ad_params']['sub_idxs'] - if step == 0: - print('BrushNet inference: sample', x.shape, ', CL', cl2apply.shape) - if sub_idx: + # we have batch input images + batch = cl_list[0].shape[0] + # we have incoming latents + latents_incoming = x.shape[0] + # and we already got some + latents_got = mp['brushnet_latent_id'] + if step == 0 or batch > 1: + print('BrushNet inference, step = %d: image batch = %d, got %d latents, starting from %d' \ + % (step, batch, latents_incoming, latents_got)) + + image_latents = [] + masks = [] + prompt_embeds = [] + negative_prompt_embeds = [] + pooled_prompt_embeds = [] + negative_pooled_prompt_embeds = [] + if sub_idx: + if step == 0: print('BrushNet inference: AnimateDiff indexes detected and applied') - check = False - if (x.shape[0], x.shape[2], x.shape[3]) != (cl2apply.shape[0], cl2apply.shape[2], cl2apply.shape[3]): - check = True + batch = len(sub_idx) - if x.shape[0] > cl2apply.shape[0]: - batch = x.shape[0] // cl2apply.shape[0] - if step == 0: - print('BrushNet inference: latent batch', batch) - conditioning_latents = torch.cat([cl2apply] * batch, dim=0).to(torch_dtype).to(brushnet.device) - elif x.shape[0] < cl2apply.shape[0]: - if step == 0: - print('BrushNet inference: sample batch is less, than CL!') - conditioning_latents = cl2apply[:x.shape[0],:,:,:].to(torch_dtype).to(brushnet.device) - else: - conditioning_latents = cl2apply.to(torch_dtype).to(brushnet.device) - - if ack2['text_embeds'] is not None and ack2['time_ids'] is not None: if do_classifier_free_guidance: - if x.shape[0] > 2: - added_cond_kwargs['text_embeds'] = torch.cat([ack2['text_embeds']] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device) - added_cond_kwargs['time_ids'] = torch.cat([ack2['time_ids']] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device) - else: - added_cond_kwargs = ack2 - elif x.shape[0] > 1: - added_cond_kwargs['text_embeds'] = torch.cat([ack['text_embeds']] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device) - added_cond_kwargs['time_ids'] = torch.cat([ack['time_ids']] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device) + for i in sub_idx: + image_latents.append(cl_list[0][i][None,:,:,:]) + masks.append(cl_list[1][i][None,:,:,:]) + prompt_embeds.append(pe) + negative_prompt_embeds.append(npe) + pooled_prompt_embeds.append(ppe) + negative_pooled_prompt_embeds.append(nppe) + for i in sub_idx: + image_latents.append(cl_list[0][i][None,:,:,:]) + masks.append(cl_list[1][i][None,:,:,:]) else: - added_cond_kwargs = ack + for i in sub_idx: + image_latents.append(cl_list[0][i][None,:,:,:]) + masks.append(cl_list[1][i][None,:,:,:]) + prompt_embeds.append(pe) + pooled_prompt_embeds.append(ppe) else: - added_cond_kwargs = ack + # do_classifier_free_guidance = 2 passes, 1st pass is cond, 2nd is uncond + continue_batch = True + for i in range(latents_incoming): + number = latents_got + i + if number < batch: + # 1st pass, cond + image_latents.append(cl_list[0][number][None,:,:,:]) + masks.append(cl_list[1][number][None,:,:,:]) + prompt_embeds.append(pe) + pooled_prompt_embeds.append(ppe) + elif do_classifier_free_guidance and number < batch * 2: + # 2nd pass, uncond + image_latents.append(cl_list[0][number-batch][None,:,:,:]) + masks.append(cl_list[1][number-batch][None,:,:,:]) + negative_prompt_embeds.append(npe) + negative_pooled_prompt_embeds.append(nppe) + else: + # latent batch + image_latents.append(cl_list[0][0][None,:,:,:]) + masks.append(cl_list[1][0][None,:,:,:]) + prompt_embeds.append(pe) + pooled_prompt_embeds.append(ppe) + latents_got = -i + continue_batch = False + + if continue_batch: + # we don't have full batch yet + if do_classifier_free_guidance: + if number < batch * 2 - 1: + mp['brushnet_latent_id'] = number + 1 + else: + mp['brushnet_latent_id'] = 0 + else: + if number < batch - 1: + mp['brushnet_latent_id'] = number + 1 + else: + mp['brushnet_latent_id'] = 0 + else: + mp['brushnet_latent_id'] = 0 + + cl = [] + for il, m in zip(image_latents, masks): + cl.append(torch.concat([il, m], dim=1)) + cl2apply = torch.concat(cl, dim=0) + + conditioning_latents = cl2apply.to(torch_dtype).to(brushnet.device) + + prompt_embeds.extend(negative_prompt_embeds) + prompt_embeds = torch.concat(prompt_embeds, dim=0).to(torch_dtype).to(brushnet.device) + + if ppe is not None: + added_cond_kwargs = {} + added_cond_kwargs['time_ids'] = torch.concat([time_ids] * latents_incoming, dim = 0).to(torch_dtype).to(brushnet.device) + + pooled_prompt_embeds.extend(negative_pooled_prompt_embeds) + pooled_prompt_embeds = torch.concat(pooled_prompt_embeds, dim=0).to(torch_dtype).to(brushnet.device) + added_cond_kwargs['text_embeds'] = pooled_prompt_embeds + else: + added_cond_kwargs = None if x.shape[2] != conditioning_latents.shape[2] or x.shape[3] != conditioning_latents.shape[3]: if step == 0: @@ -777,7 +759,7 @@ def brushnet_inference(x, timesteps, transformer_options): ), mode='bicubic', ).to(torch_dtype).to(brushnet.device) - if check and step == 0: + if step == 0: print('BrushNet inference: sample', x.shape, ', CL', conditioning_latents.shape) if step < control_guidance_start or step > control_guidance_end: @@ -785,29 +767,20 @@ def brushnet_inference(x, timesteps, transformer_options): else: cond_scale = brushnet_conditioning_scale - if isinstance(brushnet, PowerPaintModel): - return brushnet(x, - encoder_hidden_states=prompt_embeds, - brushnet_cond=conditioning_latents, - timestep = timesteps, - conditioning_scale=cond_scale, - guess_mode=False, - return_dict=False, - ) - else: - return brushnet(x, - encoder_hidden_states=prompt_embeds, - brushnet_cond=conditioning_latents, - timestep = timesteps, - conditioning_scale=cond_scale, - guess_mode=False, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - ) + return brushnet(x, + encoder_hidden_states=prompt_embeds, + brushnet_cond=conditioning_latents, + timestep = timesteps, + conditioning_scale=cond_scale, + guess_mode=False, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + ) else: print('BrushNet model is not a BrushNetModel class') return ([], 0, []) + # Check and add 'model_patch' to model.model_options['transformer_options'] def add_model_patch_option(model): if 'transformer_options' not in model.model_options: @@ -817,21 +790,21 @@ def add_model_patch_option(model): to["model_patch"] = {} return to + # This is main patch function -def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, conditioning_latents2, +def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, controls, - prompt_embeds, add_text_embeds, add_time_ids, - prompt_embeds2, add_text_embeds2, add_time_ids2): + prompt_embeds, negative_prompt_embeds, + pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids): to = add_model_patch_option(model) to['model_patch']['brushnet_model'] = brushnet to['model_patch']['brushnet_dtype'] = torch_dtype to['model_patch']['brushnet_latents'] = conditioning_latents - to['model_patch']['brushnet_latents2'] = conditioning_latents2 to['model_patch']['brushnet_controls'] = controls to['model_patch']['brushnet_prompt'] = prompt_embeds - to['model_patch']['brushnet_prompt2'] = prompt_embeds2 - to['model_patch']['brushnet_add_embeds'] = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - to['model_patch']['brushnet_add_embeds2'] = {"text_embeds": add_text_embeds2, "time_ids": add_time_ids2} + to['model_patch']['brushnet_negative_prompt'] = negative_prompt_embeds + to['model_patch']['brushnet_add_embeds'] = (pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids) + to['model_patch']['brushnet_latent_id'] = 0 is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL) @@ -870,7 +843,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, condi # patch layers `forward` so we can apply brushnet def forward_patched_by_brushnet(self, x, *args, **kwargs): h = self.original_forward(x, *args, **kwargs) - if hasattr(self, 'brushnet_sample'): + if hasattr(self, 'brushnet_sample') and type(self) in ComfySDLayers: if torch.is_tensor(self.brushnet_sample): #print(type(self), h.shape, self.brushnet_sample.shape) h += self.brushnet_sample.to(h.dtype).to(h.device) diff --git a/example/BrushNet_image_batch.json b/example/BrushNet_image_batch.json new file mode 100644 index 0000000..61d9ab4 --- /dev/null +++ b/example/BrushNet_image_batch.json @@ -0,0 +1 @@ +{"last_node_id": 18, "last_link_id": 24, "nodes": [{"id": 6, "type": "SAMModelLoader (segment anything)", "pos": [329, 68], "size": [347.87583007145474, 58], "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "SAM_MODEL", "type": "SAM_MODEL", "links": [2], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "SAMModelLoader (segment anything)"}, "widgets_values": ["sam_vit_h (2.56GB)"]}, {"id": 4, "type": "GroundingDinoModelLoader (segment anything)", "pos": [324, 175], "size": [361.20001220703125, 58], "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "GROUNDING_DINO_MODEL", "type": "GROUNDING_DINO_MODEL", "links": [3], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "GroundingDinoModelLoader (segment anything)"}, "widgets_values": ["GroundingDINO_SwinT_OGC (694MB)"]}, {"id": 10, "type": "BrushNetLoader", "pos": [338, 744], "size": {"0": 315, "1": 82}, "flags": {}, "order": 2, "mode": 0, "outputs": [{"name": "brushnet", "type": "BRMODEL", "links": [7], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetLoader"}, "widgets_values": ["brushnet/random_mask.safetensors", "float16"]}, {"id": 12, "type": "CLIPTextEncode", "pos": [805, 922], "size": [393.06744384765625, 101.02725219726562], "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 9}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [11], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": ["burger"]}, {"id": 11, "type": "CLIPTextEncode", "pos": [810, 786], "size": [388.26751708984375, 88.82723999023438], "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 8}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [10], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": [""]}, {"id": 9, "type": "BrushNet", "pos": [1279, 577], "size": {"0": 315, "1": 226}, "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 6}, {"name": "vae", "type": "VAE", "link": 12}, {"name": "image", "type": "IMAGE", "link": 13}, {"name": "mask", "type": "MASK", "link": 14}, {"name": "brushnet", "type": "BRMODEL", "link": 7}, {"name": "positive", "type": "CONDITIONING", "link": 10}, {"name": "negative", "type": "CONDITIONING", "link": 11}], "outputs": [{"name": "model", "type": "MODEL", "links": [15], "shape": 3, "slot_index": 0}, {"name": "positive", "type": "CONDITIONING", "links": [16], "shape": 3, "slot_index": 1}, {"name": "negative", "type": "CONDITIONING", "links": [17], "shape": 3, "slot_index": 2}, {"name": "latent", "type": "LATENT", "links": [18], "shape": 3, "slot_index": 3}], "properties": {"Node name for S&R": "BrushNet"}, "widgets_values": [1, 0, 10000]}, {"id": 8, "type": "CheckpointLoaderSimple", "pos": [333, 574], "size": [404.79998779296875, 98], "flags": {}, "order": 3, "mode": 0, "outputs": [{"name": "MODEL", "type": "MODEL", "links": [6], "shape": 3, "slot_index": 0}, {"name": "CLIP", "type": "CLIP", "links": [8, 9], "shape": 3, "slot_index": 1}, {"name": "VAE", "type": "VAE", "links": [12, 21], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, "widgets_values": ["realisticVisionV60B1_v51VAE.safetensors"]}, {"id": 14, "type": "VAEDecode", "pos": [2049, 464], "size": {"0": 210, "1": 46}, "flags": {}, "order": 13, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 19}, {"name": "vae", "type": "VAE", "link": 21}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [20], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 15, "type": "PreviewImage", "pos": [2273, 575], "size": [394.86773681640625, 360.6271057128906], "flags": {}, "order": 14, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 20}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 13, "type": "KSampler", "pos": [1709, 576], "size": {"0": 315, "1": 262}, "flags": {}, "order": 11, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 15}, {"name": "positive", "type": "CONDITIONING", "link": 16}, {"name": "negative", "type": "CONDITIONING", "link": 17}, {"name": "latent_image", "type": "LATENT", "link": 18}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [19], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSampler"}, "widgets_values": [0, "fixed", 20, 8, "euler", "normal", 1]}, {"id": 2, "type": "VHS_LoadImagesPath", "pos": [334, 306], "size": [226.8000030517578, 194], "flags": {}, "order": 4, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [4, 13, 22], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": null, "shape": 3}, {"name": "INT", "type": "INT", "links": null, "shape": 3}], "properties": {"Node name for S&R": "VHS_LoadImagesPath"}, "widgets_values": {"directory": "./output/", "image_load_cap": 6, "skip_first_images": 0, "select_every_nth": 1, "choose folder to upload": "image", "videopreview": {"hidden": false, "paused": false, "params": {"frame_load_cap": 6, "skip_first_images": 0, "filename": "./output/", "type": "path", "format": "folder", "select_every_nth": 1}}}}, {"id": 5, "type": "GroundingDinoSAMSegment (segment anything)", "pos": [997, 71], "size": [352.79998779296875, 122], "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "sam_model", "type": "SAM_MODEL", "link": 2}, {"name": "grounding_dino_model", "type": "GROUNDING_DINO_MODEL", "link": 3}, {"name": "image", "type": "IMAGE", "link": 4}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [14, 23], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "GroundingDinoSAMSegment (segment anything)"}, "widgets_values": ["burger", 0.3]}, {"id": 17, "type": "MaskToImage", "pos": [1424, 90], "size": {"0": 210, "1": 26}, "flags": {}, "order": 10, "mode": 0, "inputs": [{"name": "mask", "type": "MASK", "link": 23}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [24], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "MaskToImage"}}, {"id": 16, "type": "PreviewImage", "pos": [782, 255], "size": [379.2693328857422, 273.5831527709961], "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 22}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 18, "type": "PreviewImage", "pos": [1734, 90], "size": [353.90962829589853, 245.8631393432617], "flags": {}, "order": 12, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 24}], "properties": {"Node name for S&R": "PreviewImage"}}], "links": [[2, 6, 0, 5, 0, "SAM_MODEL"], [3, 4, 0, 5, 1, "GROUNDING_DINO_MODEL"], [4, 2, 0, 5, 2, "IMAGE"], [6, 8, 0, 9, 0, "MODEL"], [7, 10, 0, 9, 4, "BRMODEL"], [8, 8, 1, 11, 0, "CLIP"], [9, 8, 1, 12, 0, "CLIP"], [10, 11, 0, 9, 5, "CONDITIONING"], [11, 12, 0, 9, 6, "CONDITIONING"], [12, 8, 2, 9, 1, "VAE"], [13, 2, 0, 9, 2, "IMAGE"], [14, 5, 1, 9, 3, "MASK"], [15, 9, 0, 13, 0, "MODEL"], [16, 9, 1, 13, 1, "CONDITIONING"], [17, 9, 2, 13, 2, "CONDITIONING"], [18, 9, 3, 13, 3, "LATENT"], [19, 13, 0, 14, 0, "LATENT"], [20, 14, 0, 15, 0, "IMAGE"], [21, 8, 2, 14, 1, "VAE"], [22, 2, 0, 16, 0, "IMAGE"], [23, 5, 1, 17, 0, "MASK"], [24, 17, 0, 18, 0, "IMAGE"]], "groups": [], "config": {}, "extra": {}, "version": 0.4} \ No newline at end of file diff --git a/example/BrushNet_image_batch.png b/example/BrushNet_image_batch.png new file mode 100644 index 0000000..32a3034 Binary files /dev/null and b/example/BrushNet_image_batch.png differ