SDXL

README.md

@@ -6,6 +6,8 @@ My contribution is limited to the ComfyUI adaptation, and all credit goes to the
 
 ## Updates
 
+May 2, 2024. BrushNet SDXL is live. It needs positive and negative conditioning though, so workflow changes a little, see example.
+
 Apr 28, 2014. Another rework, sorry for inconvenience. But now BrushNet is native to ComfyUI. Famous cubiq's [IPAdapter Plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus) is now working with BrushNet! I hope... :) Please, report any bugs you found.
 
 Apr 18, 2024. Complete rework, no more custom `diffusers` library. It is possible to use LoRA models.
@@ -14,7 +16,7 @@ Apr 11, 2024. Initial commit.
 
 ## Plans
 
-- [ ] BrushNet SDXL
+- [x] BrushNet SDXL
 - [ ] PowerPaint v2
 - [ ] Compatibility with `ComfyUI-AnimateDiff-Evolved` and similar nodes
 
@@ -31,22 +33,32 @@ Checkpoints of BrushNet can be downloaded from [here](https://drive.google.com/d
 
 The checkpoint in `segmentation_mask_brushnet_ckpt` provides checkpoints trained on BrushData, which has segmentation prior (mask are with the same shape of objects). The `random_mask_brushnet_ckpt` provides a more general ckpt for random mask shape.
 
-Both `segmentation_mask_brushnet_ckpt` and `random_mask_brushnet_ckpt` or only one of them should be placed inside `models/inpaint/brushnet` folder.
+`segmentation_mask_brushnet_ckpt` and `random_mask_brushnet_ckpt` contains BrushNet for SD 1.5 models while 
+`segmentation_mask_brushnet_ckpt_sdxl_v0` and `random_mask_brushnet_ckpt_sdxl_v0` for SDXL.
 
-Base model should be SD1.5 type (realisticVisionV60B1_v51VAE, for example).
+You should place `diffusion_pytorch_model.safetensors` files to your `models/inpaint` folder.
 
 ## Usage
 
-Below is an example for the intended workflow. The [json file](example/BrushNet_basic.json) for the example can be found inside the 'example' directory.
+Below is an example for the intended workflow. The [workflow](example/BrushNet_basic.json) for the example can be found inside the 'example' directory.
 
 ![example workflow](example/BrushNet_basic.png?raw=true)
 
 The latent image can be from BrushNet node or not, but it should be the same size as original image (divided by 8 in latent space).
 
+<details>
+  <summary>SDXL</summary>
+  
+[workflow](example/BrushNet_SDXL_basic.json)
+
+![example workflow](example/BrushNet_SDXL_basic.png?raw=true)
+
+</details>
+
 <details>
   <summary>IPAdapter plus</summary>
   
-[json file](example/BrushNet_with_IPA.json)
+[workflow](example/BrushNet_with_IPA.json)
 
 ![example workflow](example/BrushNet_with_IPA.png?raw=true)
 
@@ -55,7 +67,7 @@ The latent image can be from BrushNet node or not, but it should be the same siz
 <details>
   <summary>LoRA</summary>
   
-[json file](example/BrushNet_with_LoRA.json)
+[workflow](example/BrushNet_with_LoRA.json)
 
 ![example workflow](example/BrushNet_with_LoRA.png?raw=true)
 
@@ -64,7 +76,7 @@ The latent image can be from BrushNet node or not, but it should be the same siz
 <details>
   <summary>Blending inpaint</summary>
 
-Sometimes inference and VAE broke image, so you need to blend inpaint image with the original: [json file](example/BrushNet_inpaint.json)
+Sometimes inference and VAE broke image, so you need to blend inpaint image with the original: [workflow](example/BrushNet_inpaint.json)
 
 ![example workflow](example/BrushNet_inpaint.png?raw=true)
 
@@ -75,7 +87,7 @@ You can see blurred and broken text after inpainting in the first image and how
 <details>
   <summary>ControlNet</summary>
 
-[json file](example/BrushNet_with_CN.json)
+[workflow](example/BrushNet_with_CN.json)
 
 ![example workflow](example/BrushNet_with_CN.png?raw=true)
 
@@ -90,7 +102,6 @@ List of known uncompartible nodes.
 
 - [WASasquatch's FreeU_Advanced](https://github.com/WASasquatch/FreeU_Advanced/tree/main)
 - [blepping's jank HiDiffusion](https://github.com/blepping/comfyui_jankhidiffusion)
-- [Kosinkadink's ComfyUI-AnimateDiff-Evolved](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved)
 
 I will think how to avoid it.
 

brushnet_nodes.py

@@ -25,6 +25,8 @@ from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 from .brushnet.brushnet import BrushNetModel
 
+import types
+
 from typing import Tuple
 
 import warnings
@@ -32,7 +34,6 @@ warnings.filterwarnings("ignore", category=UserWarning, module="torch")
 warnings.filterwarnings("ignore", category=UserWarning, module="safetensors")
 
 current_directory = os.path.dirname(os.path.abspath(__file__))
-#original_config_file = os.path.join(current_directory, 'brushnet', 'v1-inference.yaml')
 brushnet_config_file = os.path.join(current_directory, 'brushnet', 'brushnet.json')
 brushnet_xl_config_file = os.path.join(current_directory, 'brushnet', 'brushnet_xl.json')
 
@@ -64,10 +65,10 @@ class BrushNetLoader:
         brushnet_down_block, brushnet_mid_block, brushnet_up_block = brushnet_blocks(sd)
         del sd
         if brushnet_down_block == 24 and brushnet_mid_block == 2 and brushnet_up_block == 30:
-            print('Loading SD1.5 BrushNet model')
+            print('BrushNet model type: SD1.5')
             is_SDXL = False
         elif brushnet_down_block == 18 and brushnet_mid_block == 2 and brushnet_up_block == 22:
-            print('Loading SDXL BrushNet model')
+            print('BrushNet model type: Loading SDXL')
             is_SDXL = True
         else:
             raise Exception("Unknown BrushNet model")
@@ -79,7 +80,7 @@ class BrushNetLoader:
                 brushnet_config = BrushNetModel.load_config(brushnet_config_file)
             brushnet_model = BrushNetModel.from_config(brushnet_config)
 
-        print("BrushNet file", brushnet_file)
+        print("BrushNet model file:", brushnet_file)
 
         brushnet_model = load_checkpoint_and_dispatch(
             brushnet_model,
@@ -143,6 +144,8 @@ class BrushNet:
                         "image": ("IMAGE",),
                         "mask": ("MASK",),
                         "brushnet": ("BRMODEL", ),
+                        "positive": ("CONDITIONING", ),
+                        "negative": ("CONDITIONING", ),
                         "scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}),
                         "start_at": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
                         "end_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
@@ -155,7 +158,7 @@ class BrushNet:
 
     FUNCTION = "model_update"
 
-    def model_update(self, model, vae, image, mask, brushnet, scale, start_at, end_at):
+    def model_update(self, model, vae, image, mask, brushnet, positive, negative, scale, start_at, end_at):
 
         is_SDXL = False
         if isinstance(model.model.model_config, comfy.supported_models.SD15):
@@ -168,7 +171,7 @@ class BrushNet:
             is_SDXL = True
             if not brushnet["SDXL"]:
                 raise Exception("Base model is SDXL, but BrushNet is SD15 type")    
-            raise Exception("SDXL is not implemented yet")
+            #raise Exception("SDXL is not implemented yet")
         else:
             print('Base model type: ', type(model.model.model_config))
             raise Exception("Unsupported model type: " + str(type(model.model.model_config)))
@@ -198,66 +201,63 @@ class BrushNet:
         else:
             scaling_factor = sd15_scaling_factor
 
-        processed_image = torch.cat([masked_image] * 2).to(vae.device)
-        image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor
+        # prepare conditioning latents
 
-        processed_mask = torch.cat([1. - mask[None,:,:,:]] * 2)
-        interpolated_mask = torch.nn.functional.interpolate(
-                    processed_mask, 
-                    size=(
-                        image_latents.shape[-2], 
-                        image_latents.shape[-1]
+        with torch.no_grad():
+            processed_image = torch.cat([masked_image] * 2).to(vae.device)
+            image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor
+
+            processed_mask = torch.cat([1. - mask[None,:,:,:]] * 2)
+            interpolated_mask = torch.nn.functional.interpolate(
+                        processed_mask, 
+                        size=(
+                            image_latents.shape[-2], 
+                            image_latents.shape[-1]
+                        )
                     )
-                )
-        interpolated_mask = interpolated_mask.to(image_latents.device)
+            interpolated_mask = interpolated_mask.to(image_latents.device)
+            conditioning_latents = torch.concat([image_latents, interpolated_mask], 1).to(dtype=torch_dtype).to(brushnet['brushnet'].device)
 
-        print('BrushNet: image latents shape =', image_latents.shape, 'interpolated mask shape =', interpolated_mask.shape)
+        print('BrushNet CL: image_latents shape =', image_latents.shape, 'interpolated_mask shape =', interpolated_mask.shape)
 
-        conditioning_latents = torch.concat([image_latents, interpolated_mask], 1).to(dtype=torch_dtype).to(brushnet['brushnet'].device)
+        # prepare embeddings
 
-        # apply patches to code
+        prompt_embeds = positive[0][0][0].to(brushnet['brushnet'].device)
+        negative_prompt_embeds = negative[0][0][0].to(brushnet['brushnet'].device)
+
+        add_text_embeds = positive[0][1]['pooled_output'].to(brushnet['brushnet'].device)
+        negative_pooled_prompt_embeds = negative[0][1]['pooled_output'].to(brushnet['brushnet'].device)
+
+        add_time_ids = torch.FloatTensor([[height, width, 0., 0., height, width]]).to(brushnet['brushnet'].device)
+        negative_add_time_ids = add_time_ids
+
+        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+        add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+        add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        # apply patch to code
 
         if nodes.common_ksampler.__doc__ is None or 'BrushNet' not in nodes.common_ksampler.__doc__:
             nodes.original_common_ksampler = nodes.common_ksampler
             nodes.common_ksampler = modified_common_ksampler
 
-        if comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed.__doc__ is None or \
-            'BrushNet' not in comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed.__doc__:
-            comfy.ldm.modules.diffusionmodules.openaimodel.original_forward_timestep_embed =  \
-            comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed
-            comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed = modified_forward_timestep_embed
-
         # apply patch to model
 
         brushnet_conditioning_scale = scale
         control_guidance_start = start_at
         control_guidance_end = end_at
 
-        add_brushnet_patch(model, brushnet["brushnet"], conditioning_latents, 
-                           [brushnet_conditioning_scale, control_guidance_start, control_guidance_end])
-        
-        if is_SDXL:
-            add_model_patch(model, brushnet_inference, ('input', 0), (0, 'before'))
-            input_blocks = [[0,0],[1,0],[2,0],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1]]
-            for i, j in input_blocks:
-                add_model_patch(model, apply_brushnet, ('input', i), (j, 'after'))
-            add_model_patch(model, apply_brushnet, ('middle', 0), (2, 'after'))
-            output_blocks = [[0,1],[1,1],[2,1],[2,2],[3,1],[4,1],[5,1],[5,2],[6,0],[7,0],[8,0]]
-            for i, j in output_blocks:
-                add_model_patch(model, apply_brushnet, ('output', i), (j, 'after'))
-        else:
-            add_model_patch(model, brushnet_inference, ('input', 0), (0, 'before'))
-            input_blocks = [[0,0],[1,1],[2,1],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1],[9,0],[10,0],[11,0]]
-            for i, j in input_blocks:
-                add_model_patch(model, apply_brushnet, ('input', i), (j, 'after'))
-            add_model_patch(model, apply_brushnet, ('middle', 0), (2, 'after'))
-            output_blocks = [[0,0],[1,0],[2,0],[2,1],[3,1],[4,1],[5,1],[5,2],[6,1],[7,1],[8,1],[8,2],[9,1],[10,1],[11,1]]
-            for i, j in output_blocks:
-                add_model_patch(model, apply_brushnet, ('output', i), (j, 'after'))
+        cloned = model.clone()
+        add_brushnet_patch(cloned, 
+                           brushnet['brushnet'],
+                           conditioning_latents, 
+                           (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), 
+                           prompt_embeds, add_text_embeds, add_time_ids)
 
         latent = torch.zeros([1, 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)
 
-        return (model, {"samples":latent},)
+        return (cloned, {"samples":latent},)
 
 
 class BlendInpaint:
@@ -299,76 +299,8 @@ class BlendInpaint:
 
         return (torch.stack(ret), blurred_mask[0],)
 
-'''
-import pip
 
-class TestNode:
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {    
-                        "image": ("IMAGE",),
-                     },
-                }
-
-    CATEGORY = "inpaint"
-    RETURN_TYPES = ()
-    RETURN_NAMES = ()
-
-    FUNCTION = "test_node"
-
-    def test_node(self):
-        print('Installed packages')
-        print(pip.utils.get_installed_distributions())
-'''
-
-
-
-# Unfortunately, ModelPatcher does not have necessary hooks to patch, so we have to patch the code instead
-def modified_forward_timestep_embed(block, x, emb, context=None, transformer_options={}, 
-                           output_shape=None, time_context=None, 
-                           num_video_frames=None, image_only_indicator=None):
-    '''
-    Modified by BrushNet nodes
-    '''
-
-    if 'model_patch' not in transformer_options or transformer_options['block'] not in transformer_options['model_patch']:
-        return comfy.ldm.modules.diffusionmodules.openaimodel.original_forward_timestep_embed(block, x, emb, context, 
-                                                                                              transformer_options,
-                                                                                              output_shape, time_context, 
-                                                                                              num_video_frames, image_only_indicator)
-        
-    block_patch = transformer_options['model_patch'][transformer_options['block']]
-
-    for i, layer in enumerate(block):
-
-        if (i, 'before') in block_patch:
-            x = block_patch[(i, 'before')](x, emb, context, (transformer_options['block'], i, 'before'), transformer_options)
-
-        if isinstance(layer, comfy.ldm.modules.diffusionmodules.openaimodel.VideoResBlock):
-            x = layer(x, emb, num_video_frames, image_only_indicator)
-        elif isinstance(layer, comfy.ldm.modules.diffusionmodules.openaimodel.TimestepBlock):
-            x = layer(x, emb)
-        elif isinstance(layer, comfy.ldm.modules.attention.SpatialVideoTransformer):
-            x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options)
-            if "transformer_index" in transformer_options:
-                transformer_options["transformer_index"] += 1
-        elif isinstance(layer, comfy.ldm.modules.attention.SpatialTransformer):
-            x = layer(x, context, transformer_options)
-            if "transformer_index" in transformer_options:
-                transformer_options["transformer_index"] += 1
-        elif isinstance(layer, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample):
-            x = layer(x, output_shape=output_shape)
-        else:
-            x = layer(x)
-
-        if (i, 'after') in block_patch:
-            x = block_patch[(i, 'after')](x, emb, context, (transformer_options['block'], i, 'after'), transformer_options)
-        
-    return x
-
-# Model needs current step number at inference stepю. It is possible to write a custom KSampler but we'd like to use ComfyUI's one.
+# Model needs current step number at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one.
 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):
     '''
@@ -408,6 +340,43 @@ def modified_common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, p
     out["samples"] = samples
     return (out, )
 
+def brushnet_inference(x, timesteps, transformer_options):
+    if 'model_patch' not in transformer_options:
+        print('BrushNet inference: there is no model_patch in transformer_options')
+        return ([], [], [])
+    mp = transformer_options['model_patch']
+    brushnet = mp['brushnet_model']
+    if isinstance(brushnet, BrushNetModel):
+        conditioning_latents = mp['brushnet_latents']
+        brushnet_conditioning_scale, control_guidance_start, control_guidance_end = mp['brushnet_controls']
+        prompt_embeds = mp['brushnet_prompt']
+        added_cond_kwargs = mp['brushnet_add_embeds']
+        step = mp['step']
+        total_steps = mp['total_steps']
+        
+        brushnet_keep = []
+        for i in range(total_steps):
+            keeps = [
+                1.0 - float(i / total_steps < s or (i + 1) / total_steps > e)
+                for s, e in zip([control_guidance_start], [control_guidance_end])
+            ]
+            brushnet_keep.append(keeps[0])
+
+        cond_scale = brushnet_conditioning_scale * brushnet_keep[step]
+
+        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 ([], [], [])
+
 def add_model_patch_option(model):
     if 'transformer_options' not in model.model_options:
         model.model_options['transformer_options'] = {}
@@ -416,90 +385,66 @@ def add_model_patch_option(model):
         to["model_patch"] = {}
     return to
 
-def add_model_patch(model, patch, block, key, replace=True):
-    to = add_model_patch_option(model)
-    if block not in to["model_patch"]:
-        to["model_patch"][block] = {}
-    if key not in to["model_patch"][block]:
-        to["model_patch"][block][key] = []
-
-    if replace:
-        to["model_patch"][block][key] = patch
-    else:
-        to["model_patch"][block][key].append(patch)
-
-def add_brushnet_patch(model, brushnet, conditioning_latents, controls):
+def add_brushnet_patch(model, brushnet, conditioning_latents, controls, prompt_embeds, add_text_embeds, add_time_ids):
     to = add_model_patch_option(model)
     to['model_patch']['brushnet_model'] = brushnet
     to['model_patch']['brushnet_latents'] = conditioning_latents
     to['model_patch']['brushnet_controls'] = controls
+    to['model_patch']['brushnet_prompt'] = prompt_embeds
+    to['model_patch']['brushnet_add_embeds'] = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
     to['model_patch']['input_samples'] = []
     to['model_patch']['middle_sample'] = 0
     to['model_patch']['output_samples'] = []
 
-def brushnet_inference(x, emb, context, loc, transformer_options):
-    # x : sample
-    # emb : time embedding
-    # context[0] == neg_cond[0][0][0]
-    # context[1] == pos_cond[0][0][0]
-    # loc = (('input'|'middle'|'output', i), j, 'before'|'after')
+    is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL)
 
-    if 'model_patch' not in transformer_options:
-        raise Exception('BrushNet inference: there is no model_patch in transformer_options')
+    if is_SDXL:
+        input_blocks = [[0,0],[1,0],[2,0],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1]]
+        output_blocks = [[0,1],[1,1],[2,1],[2,2],[3,1],[4,1],[5,1],[5,2],[6,0],[7,0],[8,0]]
+    else:
+        input_blocks = [[0,0],[1,1],[2,1],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1],[9,0],[10,0],[11,0]]
+        output_blocks = [[0,0],[1,0],[2,0],[2,1],[3,1],[4,1],[5,1],[5,2],[6,1],[7,1],[8,1],[8,2],[9,1],[10,1],[11,1]]
     
-    mp = transformer_options['model_patch']
-    brushnet = mp['brushnet_model']
-    if isinstance(brushnet, BrushNetModel):
-        conditioning_latents = mp['brushnet_latents']
-        step = mp['step']
-        total_steps = mp['total_steps']
-        brushnet_conditioning_scale, control_guidance_start, control_guidance_end = mp['brushnet_controls']
-        brushnet_keep = []
-        for i in range(total_steps):
-            keeps = [
-                1.0 - float(i / total_steps < s or (i + 1) / total_steps > e)
-                for s, e in zip([control_guidance_start], [control_guidance_end])
-            ]
-            brushnet_keep.append(keeps[0])
-        cond_scale = brushnet_conditioning_scale * brushnet_keep[step]
+    # patch model `forward` so we can call BrushNet inference and setup additional samples to layers
+    if not hasattr(model.model.diffusion_model, 'original_forward'):
+        model.model.diffusion_model.original_forward = model.model.diffusion_model.forward
+    def forward_patched_by_brushnet(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+        # brushnet inference
+        input_samples, mid_sample, output_samples = brushnet_inference(x, timesteps, transformer_options)
+        # give blocks additional samples
+        # no control of samples length
+        for i, block in enumerate(self.input_blocks):
+            for j, layer in enumerate(block):
+                if [i,j] in input_blocks:
+                    layer.brushnet_sample = input_samples.pop(0)
+        self.middle_block[-1].brushnet_sample = mid_sample
+        for i, block in enumerate(self.output_blocks):
+            for j, layer in enumerate(block):
+                if [i,j] in output_blocks:
+                    layer.brushnet_sample = output_samples.pop(0)
+        return self.original_forward(x, timesteps, context, y, control, transformer_options, **kwargs)
+    model.model.diffusion_model.forward = types.MethodType(forward_patched_by_brushnet, model.model.diffusion_model)
 
-        down_samples, mid_sample, up_samples = brushnet(x,
-                                                        encoder_hidden_states=context[1],
-                                                        brushnet_cond=conditioning_latents,
-                                                        time_emb=emb,
-                                                        conditioning_scale=cond_scale,
-                                                        guess_mode=False,
-                                                        return_dict=False,)
-        transformer_options['model_patch']['input_samples'] = down_samples
-        transformer_options['model_patch']['middle_sample'] = mid_sample
-        transformer_options['model_patch']['output_samples'] = up_samples
-    else:
-        raise Exception('BrushNet model is not a BrushNetModel class')
-        
-    return x
+    # 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'):
+            h += self.brushnet_sample
+        return h
 
-def apply_brushnet(x, emb, context, loc, transformer_options):
-    if loc[0][0] == 'input':
-        if len(transformer_options['model_patch']['input_samples']) > 0:
-            return x + transformer_options['model_patch']['input_samples'].pop(0)
-        else:
-            raise Exception('BrushNet: something is not right, input samples are empty, ' + str(loc))
-        
-    elif loc[0][0] == 'middle':
-        return x + transformer_options['model_patch']['middle_sample']
-    else:
-        if len(transformer_options['model_patch']['output_samples']) > 0:
-            return x + transformer_options['model_patch']['output_samples'].pop(0)
-        else:
-            raise Exception('BrushNet: something is not right, output samples are empty, ' + str(loc))
+    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)
 
-#def image_from_tensor(t: torch.Tensor) -> np.ndarray:
-#    image_np = t.numpy() 
-#    # Convert the numpy array back to the original range (0-255) and data type (uint8)
-#    image_np = (image_np * 255).astype(np.uint8)
-#    return image_np
-#
-#def numpy_to_tensor(array: np.ndarray) -> torch.Tensor:
-#    """Convert a numpy array to a tensor and scale its values from 0-255 to 0-1."""
-#    array = array.astype(np.float32) / 255.0
-#    return torch.from_numpy(array)[None,]
+    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)
+
+    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)

requirements.txt

@@ -1,4 +1,3 @@
 diffusers>=0.27.0
 accelerate>=0.29.0
-numpy
-opencv-python
+peft>=0.7.0