RAUNet

model_patch.py

@@ -13,31 +13,6 @@ def add_model_patch_option(model):
     return to
 
 
-# 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)
-
-    ### Modified part ######################################################################
-    #
-    to = add_model_patch_option(model)
-    to['model_patch']['all_sigmas'] = sigmas
-
-    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)
-
 # Patch model with model_function_wrapper
 def patch_model_function_wrapper(model, forward_patch):
 
@@ -73,12 +48,12 @@ def patch_model_function_wrapper(model, forward_patch):
         t = model.model.model_sampling.timestep(timestep).float()
         # execute all patches        
         for method in mp['forward']:
-            method(unet, xc, t, to)
+            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:', type(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)   
 
     to = add_model_patch_option(model)
@@ -101,8 +76,56 @@ def patch_model_function_wrapper(model, forward_patch):
     mp['step'] = 0
     mp['total_steps'] = 1
 
-    # apply patch to code
+    # 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
 
+    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
+
+
+# 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)
+
+    ### Modified part ######################################################################
+    #
+    to = add_model_patch_option(model)
+    to['model_patch']['all_sigmas'] = sigmas
+
+    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)
+
+
+# To use Controlnet with RAUNet it is much easier to modify apply_control a little
+def modified_apply_control(h, control, name):
+    '''
+    Modified by BrushNet nodes
+    '''
+    if control is not None and name in control and len(control[name]) > 0:
+        ctrl = control[name].pop()
+        if ctrl is not None:
+            if h.shape[2] != ctrl.shape[2] or h.shape[3] != ctrl.shape[3]:
+                ctrl = torch.nn.functional.interpolate(ctrl, size=(h.shape[2], h.shape[3]), mode='bicubic').to(h.dtype).to(h.device)                    
+            try:
+                h += ctrl
+            except:
+                print.warning("warning control could not be applied {} {}".format(h.shape, ctrl.shape))
+    return h    
+