diff --git a/__init__.py b/__init__.py index 1aa35c4..59dfe92 100644 --- a/__init__.py +++ b/__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", } diff --git a/brushnet/brushnet.py b/brushnet/brushnet.py index 8ce9910..aed1cfd 100644 --- a/brushnet/brushnet.py +++ b/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. diff --git a/brushnet/brushnet_ca.py b/brushnet/brushnet_ca.py index 10366de..1e6f13a 100644 --- a/brushnet/brushnet_ca.py +++ b/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) diff --git a/brushnet_nodes.py b/brushnet_nodes.py index 2c6afe9..c6769b4 100644 --- a/brushnet_nodes.py +++ b/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: