diff --git a/brushnet_nodes.py b/brushnet_nodes.py index 47803d0..4e83b7f 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -1,5 +1,4 @@ import os -import types from typing import Tuple import torch @@ -274,11 +273,11 @@ class PowerPaint: # unload vae and CLIPs del vae del clip - for loaded_model in comfy.model_management.current_loaded_models: - if type(loaded_model.model.model) in ModelsToUnload: - comfy.model_management.current_loaded_models.remove(loaded_model) - loaded_model.model_unload() - del loaded_model + # for loaded_model in comfy.model_management.current_loaded_models: + # if type(loaded_model.model) in ModelsToUnload: + # comfy.model_management.current_loaded_models.remove(loaded_model) + # loaded_model.model_unload() + # del loaded_model # apply patch to model @@ -362,11 +361,11 @@ class BrushNet: # unload vae del vae - for loaded_model in comfy.model_management.current_loaded_models: - if type(loaded_model.model.model) in ModelsToUnload: - comfy.model_management.current_loaded_models.remove(loaded_model) - loaded_model.model_unload() - del loaded_model + # for loaded_model in comfy.model_management.current_loaded_models: + # if type(loaded_model.model) in ModelsToUnload: + # comfy.model_management.current_loaded_models.remove(loaded_model) + # loaded_model.model_unload() + # del loaded_model # prepare embeddings @@ -752,14 +751,14 @@ def get_image_latents(masked_image, mask, vae, scaling_factor): # Main function where magic happens @torch.inference_mode() def brushnet_inference(x, timesteps, transformer_options, debug): - if 'model_patch' not in transformer_options: + if 'brush_model_patch' not in transformer_options: print('BrushNet inference: there is no model_patch key in transformer_options') return ([], 0, []) - mp = transformer_options['model_patch'] - if 'brushnet' not in mp: + brush_model_patch = transformer_options['brush_model_patch'] + if 'brushnet' not in brush_model_patch: print('BrushNet inference: there is no brushnet key in mdel_patch') return ([], 0, []) - bo = mp['brushnet'] + bo = brush_model_patch['brushnet'] if 'model' not in bo: print('BrushNet inference: there is no model key in brushnet') return ([], 0, []) @@ -784,8 +783,8 @@ def brushnet_inference(x, timesteps, transformer_options, debug): timesteps = timesteps.detach().clone() timesteps = timesteps.to(torch_dtype).to(brushnet.device) - total_steps = mp['total_steps'] - step = mp['step'] + total_steps = brush_model_patch['total_steps'] + step = brush_model_patch['step'] added_cond_kwargs = {} @@ -938,8 +937,16 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL) + if model.model.model_config.custom_operations is None: + fp8 = model.model.model_config.optimizations.get("fp8", model.model.model_config.scaled_fp8 is not None) + operations = comfy.ops.pick_operations(model.model.model_config.unet_config.get("dtype", None), model.model.manual_cast_dtype, + fp8_optimizations=fp8, scaled_fp8=model.model.model_config.scaled_fp8) + else: + # such as gguf + operations = model.model.model_config.custom_operations + if is_SDXL: - input_blocks = [[0, comfy.ops.disable_weight_init.Conv2d], + input_blocks = [[0, operations.Conv2d], [1, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock], [2, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock], [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample], @@ -961,7 +968,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, [7, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock], [8, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock]] else: - input_blocks = [[0, comfy.ops.disable_weight_init.Conv2d], + input_blocks = [[0, operations.Conv2d], [1, comfy.ldm.modules.attention.SpatialTransformer], [2, comfy.ldm.modules.attention.SpatialTransformer], [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample], @@ -999,8 +1006,10 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, return -1, layer_list.reverse() return len(layer_list) - 1 - layer_list.index(tp), layer_list - def brushnet_forward(model, x, timesteps, transformer_options, control): - if 'brushnet' not in transformer_options['model_patch']: + def brushnet_forward(unetModel, x, timesteps, transformer_options, control): + # 这里的model是unetModel + brush_model_patch = transformer_options['brush_model_patch'] + if 'brushnet' not in brush_model_patch: input_samples = [] mid_sample = 0 output_samples = [] @@ -1010,31 +1019,32 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, # give additional samples to blocks for i, tp in input_blocks: - idx, layer_list = last_layer_index(model.input_blocks[i], tp) + idx, layer_list = last_layer_index(unetModel.input_blocks[i], tp) if idx < 0: print("BrushNet can't find", tp, "layer in", i,"input block:", layer_list) continue - model.input_blocks[i][idx].add_sample_after = input_samples.pop(0) if input_samples else 0 + unetModel.input_blocks[i][idx].add_sample_after = input_samples.pop(0) if input_samples else 0 - idx, layer_list = last_layer_index(model.middle_block, middle_block[1]) + idx, layer_list = last_layer_index(unetModel.middle_block, middle_block[1]) if idx < 0: print("BrushNet can't find", middle_block[1], "layer in middle block", layer_list) - model.middle_block[idx].add_sample_after = mid_sample + unetModel.middle_block[idx].add_sample_after = mid_sample for i, tp in output_blocks: - idx, layer_list = last_layer_index(model.output_blocks[i], tp) + idx, layer_list = last_layer_index(unetModel.output_blocks[i], tp) if idx < 0: print("BrushNet can't find", tp, "layer in", i,"outnput block:", layer_list) continue - model.output_blocks[i][idx].add_sample_after = output_samples.pop(0) if output_samples else 0 + unetModel.output_blocks[i][idx].add_sample_after = output_samples.pop(0) if output_samples else 0 patch_model_function_wrapper(model, brushnet_forward) to = add_model_patch_option(model) - mp = to['model_patch'] - if 'brushnet' not in mp: - mp['brushnet'] = {} - bo = mp['brushnet'] + brush_model_patch = to['brush_model_patch'] + + if 'brushnet' not in brush_model_patch: + brush_model_patch['brushnet'] = {} + bo = brush_model_patch['brushnet'] bo['model'] = brushnet bo['dtype'] = torch_dtype @@ -1045,37 +1055,4 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, bo['add_embeds'] = (pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids) bo['latent_id'] = 0 - # 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, 'add_sample_after') and type(self): - to_add = self.add_sample_after - if torch.is_tensor(to_add): - # interpolate due to RAUNet - if h.shape[2] != to_add.shape[2] or h.shape[3] != to_add.shape[3]: - to_add = torch.nn.functional.interpolate(to_add, size=(h.shape[2], h.shape[3]), mode='bicubic') - h += to_add.to(h.dtype).to(h.device) - else: - h += self.add_sample_after - self.add_sample_after = 0 - return h - for i, block in enumerate(model.model.diffusion_model.input_blocks): - for j, layer in enumerate(block): - if not hasattr(layer, 'original_forward'): - layer.original_forward = layer.forward - layer.forward = types.MethodType(forward_patched_by_brushnet, layer) - layer.add_sample_after = 0 - - for j, layer in enumerate(model.model.diffusion_model.middle_block): - if not hasattr(layer, 'original_forward'): - layer.original_forward = layer.forward - layer.forward = types.MethodType(forward_patched_by_brushnet, layer) - layer.add_sample_after = 0 - - for i, block in enumerate(model.model.diffusion_model.output_blocks): - for j, layer in enumerate(block): - if not hasattr(layer, 'original_forward'): - layer.original_forward = layer.forward - layer.forward = types.MethodType(forward_patched_by_brushnet, layer) - layer.add_sample_after = 0 diff --git a/model_patch.py b/model_patch.py index bbe55e1..ca25cb4 100644 --- a/model_patch.py +++ b/model_patch.py @@ -1,130 +1,163 @@ import torch import comfy +import types -# Check and add 'model_patch' to model.model_options['transformer_options'] +# Check and add 'brush_model_patch' to model.model_options['transformer_options'] def add_model_patch_option(model): if 'transformer_options' not in model.model_options: model.model_options['transformer_options'] = {} to = model.model_options['transformer_options'] - if "model_patch" not in to: - to["model_patch"] = {} + if "brush_model_patch" not in to: + to["brush_model_patch"] = {} return to # Patch model with model_function_wrapper def patch_model_function_wrapper(model, forward_patch, remove=False): - def brushnet_model_function_wrapper(apply_model_method, options_dict): - to = options_dict['c']['transformer_options'] - - control = None - if 'control' in options_dict['c']: - control = options_dict['c']['control'] - - x = options_dict['input'] - timestep = options_dict['timestep'] - - # check if there are patches to execute - if 'model_patch' not in to or 'forward' not in to['model_patch']: - return apply_model_method(x, timestep, **options_dict['c']) - - mp = to['model_patch'] - unet = mp['unet'] - - - - #print(model.get_model_object("model_sampling").sigmas, len(model.get_model_object("model_sampling").sigmas)) - #print(mp['all_sigmas'], len(mp['all_sigmas'])) - - - all_sigmas = mp['all_sigmas'] - sigma = to['sigmas'][0].item() - total_steps = all_sigmas.shape[0] - 1 - step = torch.argmin((all_sigmas - sigma).abs()).item() - - mp['step'] = step - mp['total_steps'] = total_steps - - # comfy.model_base.apply_model - xc = model.model.model_sampling.calculate_input(timestep, x) - if 'c_concat' in options_dict['c'] and options_dict['c']['c_concat'] is not None: - xc = torch.cat([xc] + [options_dict['c']['c_concat']], dim=1) - t = model.model.model_sampling.timestep(timestep).float() - # execute all patches - for method in mp['forward']: - method(unet, xc, t, to, control) - - return apply_model_method(x, timestep, **options_dict['c']) - - if "model_function_wrapper" in model.model_options and model.model_options["model_function_wrapper"]: - print('BrushNet is going to replace existing model_function_wrapper:', model.model_options["model_function_wrapper"]) - model.set_model_unet_function_wrapper(brushnet_model_function_wrapper) - + patch_wrapper = {} + comfy.patcher_extension.add_wrapper(comfy.patcher_extension.WrappersMP.APPLY_MODEL, + apply_model_function_wrapper, + patch_wrapper, + is_model_options=False) to = add_model_patch_option(model) - mp = to['model_patch'] + brush_model_patch = to['brush_model_patch'] if isinstance(model.model.model_config, comfy.supported_models.SD15): - mp['SDXL'] = False + brush_model_patch['SDXL'] = False elif isinstance(model.model.model_config, comfy.supported_models.SDXL): - mp['SDXL'] = True + brush_model_patch['SDXL'] = True else: print('Base model type: ', type(model.model.model_config)) raise Exception("Unsupported model type: ", type(model.model.model_config)) - if 'forward' not in mp: - mp['forward'] = [] + if 'forward' not in brush_model_patch: + brush_model_patch['forward'] = [] if remove: - if forward_patch in mp['forward']: - mp['forward'].remove(forward_patch) + if forward_patch in brush_model_patch['forward']: + brush_model_patch['forward'].remove(forward_patch) else: - mp['forward'].append(forward_patch) + brush_model_patch['forward'].append(forward_patch) - mp['unet'] = model.model.diffusion_model - mp['step'] = 0 - mp['total_steps'] = 1 + brush_model_patch['step'] = 0 + brush_model_patch['total_steps'] = 1 # apply patches to code - if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__: - comfy.samplers.original_sample = comfy.samplers.sample - comfy.samplers.sample = modified_sample + comfy.patcher_extension.add_wrapper(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, + brushNet_out_sample_wrapper, + patch_wrapper, + is_model_options=False) + model.wrappers.update(patch_wrapper['wrappers']) - if comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__ is None or \ - 'BrushNet' not in comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__: - comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control = comfy.ldm.modules.diffusionmodules.openaimodel.apply_control - comfy.ldm.modules.diffusionmodules.openaimodel.apply_control = modified_apply_control +def apply_model_function_wrapper(apply_model_executor, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs): + baseModel = apply_model_executor.class_obj + to = transformer_options + + timestep = t + + # check if there are patches to execute + if 'brush_model_patch' not in to or 'forward' not in to['brush_model_patch']: + return apply_model_executor(x, timestep, c_concat, c_crossattn, control, transformer_options, **kwargs) + + brush_model_patch = to['brush_model_patch'] + unet = baseModel.diffusion_model + + all_sigmas = brush_model_patch['all_sigmas'] + sigma = to['sigmas'][0].item() + total_steps = all_sigmas.shape[0] - 1 + step = torch.argmin((all_sigmas - sigma).abs()).item() + + brush_model_patch['step'] = step + brush_model_patch['total_steps'] = total_steps + + # comfy.model_base.apply_model + xc = baseModel.model_sampling.calculate_input(timestep, x) + if c_concat is not None: + xc = torch.cat([xc] + [c_concat], dim=1) + brush_t = baseModel.model_sampling.timestep(timestep).float() + # execute all patches + for method in brush_model_patch['forward']: + method(unet, xc, brush_t, to, control) + + return apply_model_executor(x, timestep, c_concat, c_crossattn, control, transformer_options, **kwargs) + +def brushNet_out_sample_wrapper(wrapper_executor, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + cfg_guider = wrapper_executor.class_obj + # set hook + diffusion_model = cfg_guider.model_patcher.model.diffusion_model + set_brushNet_hook(diffusion_model) + # Use cfd_guider.model_options, which is copied from modelPatcher.model_options and will be restored after execution without any unexpected contamination + to = add_model_patch_option(cfg_guider) + to['brush_model_patch']['all_sigmas'] = sigmas + try: + out = wrapper_executor(noise, latent_image, sampler, sigmas, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) + finally: + # cleanup hook + clean_brushNet_hook(diffusion_model) + 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 -def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, - latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - ''' - Modified by BrushNet nodes - ''' - cfg_guider = comfy.samplers.CFGGuider(model) - cfg_guider.set_conds(positive, negative) - cfg_guider.set_cfg(cfg) +def set_brushNet_hook(diffusion_model): + for i, block in enumerate(diffusion_model.input_blocks): + for j, layer in enumerate(block): + if not hasattr(layer, 'original_forward'): + layer.original_forward = layer.forward + layer.forward = types.MethodType(forward_patched_by_brushnet, layer) + layer.add_sample_after = 0 + for j, layer in enumerate(diffusion_model.middle_block): + if not hasattr(layer, 'original_forward'): + layer.original_forward = layer.forward + layer.forward = types.MethodType(forward_patched_by_brushnet, layer) + layer.add_sample_after = 0 + for i, block in enumerate(diffusion_model.output_blocks): + for j, layer in enumerate(block): + if not hasattr(layer, 'original_forward'): + layer.original_forward = layer.forward + layer.forward = types.MethodType(forward_patched_by_brushnet, layer) + layer.add_sample_after = 0 - ### Modified part ###################################################################### - # - to = add_model_patch_option(model) - to['model_patch']['all_sigmas'] = sigmas - # - #sigma_start = model.get_model_object("model_sampling").percent_to_sigma(start_at) - #sigma_end = model.get_model_object("model_sampling").percent_to_sigma(end_at) - # - # - #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) + comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control = comfy.ldm.modules.diffusionmodules.openaimodel.apply_control + comfy.ldm.modules.diffusionmodules.openaimodel.apply_control = modified_apply_control +def clean_brushNet_hook(diffusion_model): + for i, block in enumerate(diffusion_model.input_blocks): + for j, layer in enumerate(block): + if hasattr(layer, 'original_forward'): + layer.forward = layer.original_forward + del layer.original_forward + del layer.add_sample_after + for j, layer in enumerate(diffusion_model.middle_block): + if hasattr(layer, 'original_forward'): + layer.forward = layer.original_forward + del layer.original_forward + del layer.add_sample_after + for i, block in enumerate(diffusion_model.output_blocks): + for j, layer in enumerate(block): + if hasattr(layer, 'original_forward'): + layer.forward = layer.original_forward + del layer.original_forward + del layer.add_sample_after + + if hasattr(comfy.ldm.modules.diffusionmodules.openaimodel, 'original_apply_control'): + comfy.ldm.modules.diffusionmodules.openaimodel.apply_control = comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control + del comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control + +# 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, 'add_sample_after') and type(self): + to_add = self.add_sample_after + if torch.is_tensor(to_add): + # interpolate due to RAUNet + if h.shape[2] != to_add.shape[2] or h.shape[3] != to_add.shape[3]: + to_add = torch.nn.functional.interpolate(to_add, size=(h.shape[2], h.shape[3]), mode='bicubic') + h += to_add.to(h.dtype).to(h.device) + else: + h += self.add_sample_after + self.add_sample_after = 0 + return h # To use Controlnet with RAUNet it is much easier to modify apply_control a little def modified_apply_control(h, control, name): @@ -142,3 +175,4 @@ def modified_apply_control(h, control, name): print.warning("warning control could not be applied {} {}".format(h.shape, ctrl.shape)) return h + diff --git a/raunet_nodes.py b/raunet_nodes.py index 8c1b0d6..4ff01d9 100644 --- a/raunet_nodes.py +++ b/raunet_nodes.py @@ -42,11 +42,11 @@ class RAUNet: def add_raunet_patch(model, du_start, du_end, xa_start, xa_end): def raunet_forward(model, x, timesteps, transformer_options, control): - if 'model_patch' not in transformer_options: - print("RAUNet: 'model_patch' not in transformer_options, skip") + if 'brush_model_patch' not in transformer_options: + print("RAUNet: 'brush_model_patch' not in transformer_options, skip") return - mp = transformer_options['model_patch'] + mp = transformer_options['brush_model_patch'] is_SDXL = mp['SDXL'] if is_SDXL and type(model.input_blocks[6][0]) != comfy.ldm.modules.diffusionmodules.openaimodel.Downsample: @@ -87,7 +87,7 @@ def add_raunet_patch(model, du_start, du_end, xa_start, xa_end): model.set_model_output_block_patch(out_xattn_patch) to = add_model_patch_option(model) - mp = to['model_patch'] + mp = to['brush_model_patch'] if 'raunet' not in mp: mp['raunet'] = {} ro = mp['raunet'] @@ -103,10 +103,10 @@ def in_xattn_patch(h, transformer_options): if transformer_options["block"] != ("input", 4): # wrong block return h - if 'model_patch' not in transformer_options: - print("RAUNet (i-x-p): 'model_patch' not in transformer_options") + if 'brush_model_patch' not in transformer_options: + print("RAUNet (i-x-p): 'brush_model_patch' not in transformer_options") return h - mp = transformer_options['model_patch'] + mp = transformer_options['brush_model_patch'] if 'raunet' not in mp: print("RAUNet (i-x-p): 'raunet' not in model_patch options") return h @@ -123,10 +123,10 @@ def in_xattn_patch(h, transformer_options): def out_xattn_patch(h, hsp, transformer_options): - if 'model_patch' not in transformer_options: - print("RAUNet (o-x-p): 'model_patch' not in transformer_options") + if 'brush_model_patch' not in transformer_options: + print("RAUNet (o-x-p): 'brush_model_patch' not in transformer_options") return h, hsp - mp = transformer_options['model_patch'] + mp = transformer_options['brush_model_patch'] if 'raunet' not in mp: print("RAUNet (o-x-p): 'raunet' not in model_patch options") return h