From af6fe36f1742e73234463951e30baa2acccd520c Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Thu, 21 Mar 2024 15:14:52 -0700 Subject: [PATCH] Add lazy block selection to style prompting Adds the feature recommended from this issue: https://github.com/ExponentialML/ComfyUI_VisualStylePrompting/issues/9#issue-2199077769. Also fixes bugs that wasn't allowing style prompting to work correctly. --- nodes.py | 55 ++++++++++++++++++++++------- utils/attention_functions.py | 67 ++++++++++++++++++++++++++++++++++-- 2 files changed, 107 insertions(+), 15 deletions(-) diff --git a/nodes.py b/nodes.py index a5ae07e..b80b9e6 100644 --- a/nodes.py +++ b/nodes.py @@ -2,7 +2,7 @@ import comfy import torch import nodes -from .utils.attention_functions import VisualStyleProcessor +from .utils.attention_functions import VisualStyleProcessor, ATTENTION_LAYER_CANDIDATES class ApplyVisualStyle: @classmethod @@ -16,14 +16,39 @@ class ApplyVisualStyle: "reference_image_prompt": ("CONDITIONING",), "negative_prompt": ("CONDITIONING", ), "enabled": ("BOOLEAN", {"default": True}), - "denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}) + "denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}), + "input_blocks": ("BOOLEAN", {"default": False}), + "middle_block": ("BOOLEAN", {"default": False}), + "output_blocks": ("BOOLEAN", {"default": True}), + }, + "optional": { + "init_image": ("IMAGE",), } } + CATEGORY = "VisualStylePrompting/apply" RETURN_TYPES = ("MODEL", "CONDITIONING","CONDITIONING", "LATENT") RETURN_NAMES = ("model", "positive", "negative", "latents") FUNCTION = "apply_visual_style_prompt" + + def get_block_choices(self, input_blocks, middle_block, output_blocks): + block_choices_map = ( + [input_blocks, "input"], + [middle_block, "middle"], + [output_blocks, "output"] + ) + + block_choices = [] + + for block_choice in block_choices_map: + if block_choice[0] == True: + block_choices.append(block_choice[1]) + + return block_choices + + def activate_block_choice(self, key, block_choices): + return any([block in key for block in block_choices]) def apply_visual_style_prompt( self, @@ -34,22 +59,28 @@ class ApplyVisualStyle: reference_image_prompt, negative_prompt, enabled, - denoise + denoise, + input_blocks, + middle_block, + output_blocks, + init_image = None ): self.model = model reference_latent = vae.encode(reference_image[:,:,:,:3]) - for n, m in model.model.diffusion_model.named_modules(): - if m.__class__.__name__ == "CrossAttention": - processor = VisualStyleProcessor(m, enabled=enabled) - setattr(m, 'forward', processor.visual_style_forward) + block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks) - conditioning_prompt = reference_image_prompt + conditioning_prompt - negative_prompt = negative_prompt * 2 + for n, m in model.model.diffusion_model.named_modules(): + if m.__class__.__name__ == "CrossAttention": + processor = VisualStyleProcessor(m, enabled=self.activate_block_choice(n, block_choices)) + setattr(m, 'forward', processor.visual_style_forward) + + conditioning_prompt[0][0] = torch.cat([reference_image_prompt[0][0], conditioning_prompt[0][0]]) + negative_prompt[0][0] = torch.cat([negative_prompt[0][0]] * 2) latents = torch.zeros_like(reference_latent) latents = torch.cat([latents] * 2) - + if denoise < 1.0: latents[::1] = reference_latent[:1] else: @@ -57,10 +88,10 @@ class ApplyVisualStyle: denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise - denoise_mask[0] = 0. + denoise_mask[::2] = 0. return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask}) NODE_CLASS_MAPPINGS = {"ApplyVisualStyle": ApplyVisualStyle} -NODE_DISPLAY_NAME_MAPPINGS = {"ApplyVisualStyle": "Apply Visual Style Prompting"} +NODE_DISPLAY_NAME_MAPPINGS = {"ApplyVisualStyle": "Apply Visual Style Prompting"} \ No newline at end of file diff --git a/utils/attention_functions.py b/utils/attention_functions.py index 29be745..f1fad7f 100644 --- a/utils/attention_functions.py +++ b/utils/attention_functions.py @@ -1,11 +1,47 @@ -from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked -from .style_functions import adain, concat_first +from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked, attention_basic +from .style_functions import adain, concat_first, swapping_attention + +ATTENTION_LAYER_CANDIDATES = [ +"input_blocks.1.1.transformer_blocks.0.attn1", +"input_blocks.1.1.transformer_blocks.0.attn2", +"input_blocks.2.1.transformer_blocks.0.attn1", +"input_blocks.2.1.transformer_blocks.0.attn2", +"input_blocks.4.1.transformer_blocks.0.attn1", +"input_blocks.4.1.transformer_blocks.0.attn2", +"input_blocks.5.1.transformer_blocks.0.attn1", +"input_blocks.5.1.transformer_blocks.0.attn2", +"input_blocks.7.1.transformer_blocks.0.attn1", +"input_blocks.7.1.transformer_blocks.0.attn2", +"input_blocks.8.1.transformer_blocks.0.attn1", +"input_blocks.8.1.transformer_blocks.0.attn2", +"middle_block.1.transformer_blocks.0.attn1", +"middle_block.1.transformer_blocks.0.attn2", +"output_blocks.3.1.transformer_blocks.0.attn1", +"output_blocks.3.1.transformer_blocks.0.attn2", +"output_blocks.4.1.transformer_blocks.0.attn1", +"output_blocks.4.1.transformer_blocks.0.attn2", +"output_blocks.5.1.transformer_blocks.0.attn1", +"output_blocks.5.1.transformer_blocks.0.attn2", +"output_blocks.6.1.transformer_blocks.0.attn1", +"output_blocks.6.1.transformer_blocks.0.attn2", +"output_blocks.7.1.transformer_blocks.0.attn1", +"output_blocks.7.1.transformer_blocks.0.attn2", +"output_blocks.8.1.transformer_blocks.0.attn1", +"output_blocks.8.1.transformer_blocks.0.attn2", +"output_blocks.9.1.transformer_blocks.0.attn1", +"output_blocks.9.1.transformer_blocks.0.attn2", +"output_blocks.10.1.transformer_blocks.0.attn1", +"output_blocks.10.1.transformer_blocks.0.attn2", +"output_blocks.11.1.transformer_blocks.0.attn1", +"output_blocks.11.1.transformer_blocks.0.attn2" +] class VisualStyleProcessor(object): def __init__(self, module_self, keys_scale: float = 1.0, enabled: bool = True, + enabled_animatediff: bool = True, adain_queries: bool = True, adain_keys: bool = True, adain_values: bool = False @@ -13,6 +49,7 @@ class VisualStyleProcessor(object): self.module_self = module_self self.keys_scale = keys_scale self.enabled = enabled + self.enabled_animatediff = enabled self.adain_queries = adain_queries self.adain_keys = adain_keys self.adain_values = adain_values @@ -42,4 +79,28 @@ class VisualStyleProcessor(object): out = optimized_attention(q, k, v, self.module_self.heads) else: out = optimized_attention_masked(q, k, v, self.module_self.heads, mask) - return self.module_self.to_out(out) \ No newline at end of file + return self.module_self.to_out(out) + + def visual_style_forward_ad(self, x, context=None, value=None, mask=None): + q = module_self.to_q(x) + context = default(context, x) + k = module_self.to_k(context) + if value is not None: + v = module_self.to_v(value) + del value + else: + v = module_self.to_v(context) + + # apply custom scale by multiplying k by scale factor + if module_self.scale is not None: + k *= module_self.scale + + # apply scale mask, if present + if scale_mask is not None: + k *= scale_mask + + if self.enable_animatediff: + k, v = swapping_attention(key, value) + + out = attention_basic(q, k, v, module_self.heads, mask) + return self.to_out(out) \ No newline at end of file