Merge pull request #27 from PLEXATIC/layer_level_control

Layer level control

README.md

@@ -22,8 +22,11 @@ This has currently only been tested with 1.5 based models.
 - `enabled`: Enables or disables the effect. Note that this node will still be hooked even after disabling unless you remove it.
 - `denoise`: Works the same way Img2Img works, but utilized with reference and / or init images (this is experimental).
 - `input_blocks`: Focuses attention on the encoder layers.
+- `skip_input_layers`: Number of layers in the input block that will **not** have swapping self-attention applied to them.
 - `middle_block`: Focuses attention on the middle layers.
+- `skip_middle_layers`: Number of layers in the middle block that will **not** have swapping self-attention applied to them.
 - `output_blocks`: Focuses attention on the decoder layers.
+- `skip_output_layers`: Number of layers in the output block that will **not** have swapping self-attention applied to them.
 
 > [!TIP]  
 > In order to get the best results, you must engineer both the positive and reference image prompts correctly. Focus on the details you want to derive from the image reference, and the details you wish to see in the output.

nodes.py

@@ -19,8 +19,11 @@ class ApplyVisualStyle:
                 "enabled": ("BOOLEAN", {"default": True}),
                 "denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}),
                 "input_blocks": ("BOOLEAN", {"default": False}),
+                "skip_input_layers": ("INT", {"default": 24, "min": 0, "max": 48, "step": 1}),
                 "middle_block": ("BOOLEAN", {"default": False}),
+                "skip_middle_layers": ("INT", {"default": 1, "min": 0, "max": 2, "step": 1}),
                 "output_blocks": ("BOOLEAN", {"default": True}),
+                "skip_output_layers": ("INT", {"default": 24, "min": 0, "max": 72, "step": 1}),
             },
             "optional": {
                 "init_image": ("IMAGE",),
@@ -64,7 +67,11 @@ class ApplyVisualStyle:
         input_blocks,
         middle_block,
         output_blocks,
-        init_image = None
+        init_image = None,
+        skip_input_layers=0,
+        skip_middle_layers=0,
+        skip_output_layers=0
+
     ):
         positive = copy_cond(positive)
         reference_cond = copy_cond(reference_cond)
@@ -74,10 +81,28 @@ class ApplyVisualStyle:
 
         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":
-                is_enabled = self.activate_block_choice(n, block_choices)
 
+        layer_indexes = {
+            "input": 0,
+            "middle": 0,
+            "output": 0
+        }
+
+        n_skip_per_block = {
+            "input": skip_input_layers,
+            "middle": skip_middle_layers,
+            "output": skip_output_layers
+        }
+
+        class_names = set()
+        for n, m in model.model.diffusion_model.named_modules():
+            if m.__class__.__name__ == "CrossAttention":
+                is_enabled = self.activate_block_choice(n, block_choices)
+                if is_enabled:
+                    block_name = n.split("_")[0]
+                    if layer_indexes[block_name] < n_skip_per_block[block_name]:
+                        is_enabled = False
+                    layer_indexes[block_name] += 1
                 if hasattr(m.forward, 'module_self'):
                     m.forward.enabled = is_enabled and enabled
                 else:
@@ -89,15 +114,11 @@ class ApplyVisualStyle:
 
         latents = torch.zeros_like(reference_latent)
         latents = torch.cat([latents] * 2)
-
-        if denoise < 1.0:
-            latents[::1] = reference_latent[:1]
-        else:
-            latents[::2] = reference_latent
+        latents[0] = reference_latent
 
         denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise
 
-        denoise_mask[::2] = 0.
+        denoise_mask[0] = 0.
 
         return (model, positive, negative, {"samples": latents, "noise_mask": denoise_mask})
 

utils/attention_functions.py

@@ -42,8 +42,8 @@ class VisualStyleProcessor(object):
         keys_scale: float = 1.0,
         enabled: bool = True, 
         enabled_animatediff: bool = False,
-        adain_queries: bool = True,
-        adain_keys: bool = True,
+        adain_queries: bool = False,
+        adain_keys: bool = False,
         adain_values: bool = False 
     ):
         self.module_self = module_self
@@ -75,8 +75,10 @@ class VisualStyleProcessor(object):
             if self.adain_values:
                 v = adain(v)
             
-            k = concat_first(k, -2, self.keys_scale)
-            v = concat_first(v, -2)
+            k, v = swapping_attention(k, v)
+
+            #k = concat_first(k, -2, self.keys_scale)
+            #v = concat_first(v, -2)
 
         if mask is None:
             out = optimized_attention(q, k, v, self.module_self.heads)

utils/style_functions.py

@@ -5,18 +5,6 @@ from dataclasses import dataclass
 
 T = torch.Tensor
 
-@dataclass(frozen=True)
-class StyleAlignedArgs:
-    share_group_norm: bool = True
-    share_layer_norm: bool = True,
-    share_attention: bool = True
-    adain_queries: bool = True
-    adain_keys: bool = True
-    adain_values: bool = False
-    full_attention_share: bool = False
-    keys_scale: float = 1.
-    only_self_level: float = 0.
-
 def expand_first(feat: T, scale=1., ) -> T:
     b = feat.shape[0]
     feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1)