Update MagicQuill/brushnet_nodes.py

MagicQuill/brushnet_nodes.py

@@ -149,7 +149,13 @@ class PowerPaintCLIPLoader:
 
         print('PowerPaint base CLIP file: ', base_CLIP_file)
 
-        pp_tokenizer = TokenizerWrapper(pp_clip.tokenizer.clip_l.tokenizer)
+        # Reuse TokenizerWrapper if already wrapped
+        if isinstance(pp_clip.tokenizer.clip_l.tokenizer, TokenizerWrapper):
+            pp_tokenizer = pp_clip.tokenizer.clip_l.tokenizer
+        else:
+            pp_tokenizer = TokenizerWrapper(pp_clip.tokenizer.clip_l.tokenizer)
+            pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer
+
         pp_text_encoder = pp_clip.patcher.model.clip_l.transformer
 
         add_tokens(
@@ -164,7 +170,8 @@ class PowerPaintCLIPLoader:
 
         print('PowerPaint CLIP file: ', pp_CLIP_file)
 
-        pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer
+        # Already assigned above if new, or reused if existing
+        # pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer 
         pp_clip.patcher.model.clip_l.transformer = pp_text_encoder
 
         return (pp_clip,)
@@ -941,6 +948,20 @@ def brushnet_inference(x, timesteps, transformer_options, debug):
                     debug=debug,
                 )
 
+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
+
 
 # This is main patch function
 def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, 
@@ -952,7 +973,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
     is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL)
 
     if is_SDXL:
-        input_blocks = [[0, comfy.ops.disable_weight_init.Conv2d],
+        input_blocks = [[0, comfy.ops.manual_cast.Conv2d],
                         [1, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
                         [2, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
                         [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
@@ -974,7 +995,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, comfy.ops.manual_cast.Conv2d],
                         [1, comfy.ldm.modules.attention.SpatialTransformer],
                         [2, comfy.ldm.modules.attention.SpatialTransformer],
                         [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
@@ -1059,36 +1080,33 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
     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
-
+    # Modified to prevent re-patching leak and closure creation
     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)
+            
+            # Only patch if not already patched by us
+            if getattr(layer.forward, '__func__', None) != forward_patched_by_brushnet:
+                 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)
+        
+        if getattr(layer.forward, '__func__', None) != forward_patched_by_brushnet:
+             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)
+            
+            if getattr(layer.forward, '__func__', None) != forward_patched_by_brushnet:
+                 layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
+                 
             layer.add_sample_after = 0