Merge pull request #10 from ExponentialML/update/style_prompting_blocks

Add lazy block selection to style prompting

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,15 +16,40 @@ 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, 
         model: comfy.model_patcher.ModelPatcher, 
@@ -34,18 +59,24 @@ 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])
         
+        block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
+
         for n, m in model.model.diffusion_model.named_modules():
             if m.__class__.__name__  == "CrossAttention": 
-                processor = VisualStyleProcessor(m, enabled=enabled)
+                processor = VisualStyleProcessor(m, enabled=self.activate_block_choice(n, block_choices))
                 setattr(m, 'forward', processor.visual_style_forward)
         
-        conditioning_prompt = reference_image_prompt + conditioning_prompt
+        conditioning_prompt[0][0] = torch.cat([reference_image_prompt[0][0], conditioning_prompt[0][0]])
-        negative_prompt = negative_prompt * 2 
+        negative_prompt[0][0] = torch.cat([negative_prompt[0][0]] * 2)
 
         latents = torch.zeros_like(reference_latent) 
         latents = torch.cat([latents] * 2)
@@ -57,7 +88,7 @@ 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})    
 

utils/attention_functions.py

@@ -1,11 +1,47 @@
-from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked
+from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked, attention_basic
-from .style_functions import adain, concat_first
+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 = False,
         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
@@ -43,3 +80,28 @@ class VisualStyleProcessor(object):
         else:
             out = optimized_attention_masked(q, k, v, self.module_self.heads, mask)
         return self.module_self.to_out(out)
+
+    # TODO
+    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)