Merge pull request #166 from lldacing/fix
Use the hook method provided by ComfyUI to inject the patch, solve some compatibility issues of ComfyUI
This commit is contained in:
nullquant
+41
-64
@@ -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
|
||||
|
||||
+128
-94
@@ -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
|
||||
|
||||
|
||||
|
||||
+10
-10
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user