Mac debug

__init__.py

@@ -1,5 +1,6 @@
 from .brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, PowerPaintCLIPLoader, PowerPaint, CutForInpaint
 from .raunet_nodes import RAUNet
+from subprocess import getoutput
 
 """
 @author: nullquant
@@ -8,6 +9,32 @@ from .raunet_nodes import RAUNet
 @description: These are custom nodes for ComfyUI native implementation of BrushNet, PowerPaint and RAUNet models
 """
 
+class Terminal:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return { "required": { 
+                    "text": ("STRING", {"multiline": True})
+                    }
+                }
+    
+    CATEGORY = "Tools"
+    RETURN_TYPES = ("STRING", )
+    RETURN_NAMES = ("text", )
+    OUTPUT_NODE = True
+
+    FUNCTION = "execute"
+
+    def execute(self, text):
+        if text[0] == '"' and text[-1] == '"':
+            out = getoutput(f"{text[1:-1]}")
+            print(out)
+        else:
+            exec(f"{text}")
+        return (text, )
+
+
+
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
@@ -18,6 +45,7 @@ NODE_CLASS_MAPPINGS = {
     "PowerPaint": PowerPaint,
     "CutForInpaint": CutForInpaint,
     "RAUNet": RAUNet,
+    "Terminal": Terminal,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -29,4 +57,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "PowerPaint": "PowerPaint",
     "CutForInpaint": "Cut For Inpaint",
     "RAUNet": "RAUNet",
+    "Terminal": "Terminal",
 }

brushnet/brushnet.py

@@ -691,6 +691,7 @@ class BrushNetModel(ModelMixin, ConfigMixin):
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         guess_mode: bool = False,
         return_dict: bool = True,
+        debug = False,
     ) -> Union[BrushNetOutput, Tuple[Tuple[torch.FloatTensor, ...], torch.FloatTensor]]:
         """
         The [`BrushNetModel`] forward method.

brushnet/brushnet_ca.py

@@ -699,6 +699,7 @@ class BrushNetModel(ModelMixin, ConfigMixin):
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         guess_mode: bool = False,
         return_dict: bool = True,
+        debug=False,
     ) -> Union[BrushNetOutput, Tuple[Tuple[torch.FloatTensor, ...], torch.FloatTensor]]:
         """
         The [`BrushNetModel`] forward method.
@@ -750,11 +751,15 @@ class BrushNetModel(ModelMixin, ConfigMixin):
         else:
             raise ValueError(f"unknown `brushnet_conditioning_channel_order`: {channel_order}")
 
+        if debug: print('BrushNet CA: attn mask')
+
         # prepare attention_mask
         if attention_mask is not None:
             attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
             attention_mask = attention_mask.unsqueeze(1)
 
+        if debug: print('BrushNet CA: time')
+
         # 1. time
         timesteps = timestep
         if not torch.is_tensor(timesteps):
@@ -816,10 +821,15 @@ class BrushNetModel(ModelMixin, ConfigMixin):
 
         emb = emb + aug_emb if aug_emb is not None else emb
 
+        if debug: print('BrushNet CA: pre-process')
+
+
         # 2. pre-process
         brushnet_cond = torch.concat([sample, brushnet_cond], 1)
         sample = self.conv_in_condition(brushnet_cond)
 
+        if debug: print('BrushNet CA: down')
+
         # 3. down
         down_block_res_samples = (sample,)
         for downsample_block in self.down_blocks:
@@ -836,12 +846,16 @@ class BrushNetModel(ModelMixin, ConfigMixin):
 
             down_block_res_samples += res_samples
 
+        if debug: print('BrushNet CA: PP down')
+
         # 4. PaintingNet down blocks
         brushnet_down_block_res_samples = ()
         for down_block_res_sample, brushnet_down_block in zip(down_block_res_samples, self.brushnet_down_blocks):
             down_block_res_sample = brushnet_down_block(down_block_res_sample)
             brushnet_down_block_res_samples = brushnet_down_block_res_samples + (down_block_res_sample,)
 
+        if debug: print('BrushNet CA: PP mid')
+
         # 5. mid
         if self.mid_block is not None:
             if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
@@ -855,9 +869,13 @@ class BrushNetModel(ModelMixin, ConfigMixin):
             else:
                 sample = self.mid_block(sample, emb)
 
+        if debug: print('BrushNet CA: mid')
+
         # 6. BrushNet mid blocks
         brushnet_mid_block_res_sample = self.brushnet_mid_block(sample)
 
+        if debug: print('BrushNet CA: PP up')
+
         # 7. up
         up_block_res_samples = ()
         for i, upsample_block in enumerate(self.up_blocks):
@@ -893,12 +911,16 @@ class BrushNetModel(ModelMixin, ConfigMixin):
 
             up_block_res_samples += up_res_samples
 
+        if debug: print('BrushNet CA: up')
+
         # 8. BrushNet up blocks
         brushnet_up_block_res_samples = ()
         for up_block_res_sample, brushnet_up_block in zip(up_block_res_samples, self.brushnet_up_blocks):
             up_block_res_sample = brushnet_up_block(up_block_res_sample)
             brushnet_up_block_res_samples = brushnet_up_block_res_samples + (up_block_res_sample,)
 
+        if debug: print('BrushNet CA: scaling')
+
         # 6. scaling
         if guess_mode and not self.config.global_pool_conditions:
             scales = torch.logspace(
@@ -940,6 +962,8 @@ class BrushNetModel(ModelMixin, ConfigMixin):
                 torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_up_block_res_samples
             ]
 
+        if debug: print('BrushNet CA: finish')
+
         if not return_dict:
             return (brushnet_down_block_res_samples, brushnet_mid_block_res_sample, brushnet_up_block_res_samples)
 

brushnet_nodes.py

@@ -189,6 +189,7 @@ class PowerPaint:
                         "scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}),
                         "start_at": ("INT", {"default": 0, "min": 0, "max": 10000}),
                         "end_at": ("INT", {"default": 10000, "min": 0, "max": 10000}),
+                        "debug": (['False', 'True'], ),
                      },
         }
     
@@ -198,7 +199,7 @@ class PowerPaint:
 
     FUNCTION = "model_update"
 
-    def model_update(self, model, vae, image, mask, powerpaint, clip, positive, negative, fitting, function, scale, start_at, end_at):
+    def model_update(self, model, vae, image, mask, powerpaint, clip, positive, negative, fitting, function, scale, start_at, end_at, debug):
 
         is_SDXL, is_PP = check_compatibilty(model, powerpaint)
         if not is_PP:
@@ -297,7 +298,8 @@ class PowerPaint:
                            conditioning_latents, 
                            (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), 
                            negative_prompt_embeds_pp, prompt_embeds_pp, 
-                           None, None, None)
+                           None, None, None,
+                           'True' in debug)
 
         latent = torch.zeros([batch, 4, conditioning_latents[0].shape[2], conditioning_latents[0].shape[3]], device=powerpaint['brushnet'].device)
 
@@ -419,7 +421,8 @@ class BrushNet:
                            conditioning_latents, 
                            (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), 
                            prompt_embeds, negative_prompt_embeds,
-                           pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids)
+                           pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids,
+                           False)
 
         latent = torch.zeros([batch, 4, conditioning_latents[0].shape[2], conditioning_latents[0].shape[3]], device=brushnet['brushnet'].device)
 
@@ -752,7 +755,7 @@ 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):
+def brushnet_inference(x, timesteps, transformer_options, debug):
     if 'model_patch' not in transformer_options:
         print('BrushNet inference: there is no model_patch key in transformer_options')
         return ([], 0, [])
@@ -911,6 +914,8 @@ def brushnet_inference(x, timesteps, transformer_options):
     if step == 0:
         print('BrushNet inference: sample', x.shape, ', CL', conditioning_latents.shape)
 
+    if debug: print('BrushNet: step =', step)
+
     if step < control_guidance_start or step > control_guidance_end:
         cond_scale = 0.0
     else:
@@ -924,6 +929,7 @@ def brushnet_inference(x, timesteps, transformer_options):
                     guess_mode=False,
                     added_cond_kwargs=added_cond_kwargs,
                     return_dict=False,
+                    debug=debug,
                 )
 
 
@@ -931,7 +937,8 @@ def brushnet_inference(x, timesteps, transformer_options):
 def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, 
                        controls, 
                        prompt_embeds, negative_prompt_embeds,
-                       pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids):
+                       pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids,
+                       debug):
     
     is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL)
 
@@ -1003,7 +1010,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
             output_samples = []
         else:    
             # brushnet inference
-            input_samples, mid_sample, output_samples = brushnet_inference(x, timesteps, transformer_options)
+            input_samples, mid_sample, output_samples = brushnet_inference(x, timesteps, transformer_options, debug)
 
         # give additional samples to blocks
         for i, tp in input_blocks: