Merge remote-tracking branch 'source/main'

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 (matching legacy functionality), you must use the same prompt for `positive` and `reference_cond`.

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,15 +67,37 @@ 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
+
     ):
         reference_samples = reference_latent["samples"]
         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:
@@ -84,15 +109,11 @@ class ApplyVisualStyle:
 
         latents = torch.zeros_like(reference_samples)
         latents = torch.cat([latents] * 2)
-
-        if denoise < 1.0:
-            latents[::1] = reference_samples[:1]
-        else:
-            latents[::2] = reference_samples
+        latents[0] = reference_samples
 
         denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise
 
-        denoise_mask[::2] = 0.
+        denoise_mask[0] = 0.
 
         return (model, positive_cat, negative_cat, {"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)