From 5775fedb7b43f9d50200e6cc72a6a7e93e221ec8 Mon Sep 17 00:00:00 2001 From: PLEXATIC Date: Sat, 6 Apr 2024 17:41:32 +0200 Subject: [PATCH 1/4] Added logic to skip first n layers per block The paper states that the process should only be applied for the layers 24 and later --- nodes.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/nodes.py b/nodes.py index 88855d5..c564761 100644 --- a/nodes.py +++ b/nodes.py @@ -64,7 +64,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 +78,27 @@ class ApplyVisualStyle: block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks) + + layer_indexes = { + "input": 0, + "middle": 0, + "output": 0 + } + + n_skip_per_block = { + "input": skip_input_layers, + "middle": skip_middle_layers, + "output": skip_output_layers + } + 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: From 64ccb14722e3b0eba4d128f8014a7b5749ba521a Mon Sep 17 00:00:00 2001 From: PLEXATIC Date: Sat, 6 Apr 2024 18:00:29 +0200 Subject: [PATCH 2/4] Added UI Elements to control number of skipped layers --- nodes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/nodes.py b/nodes.py index c564761..8ab9edc 100644 --- a/nodes.py +++ b/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",), From 92fc9d8765abf8e9b1e2dfb9b729834a335aa0e8 Mon Sep 17 00:00:00 2001 From: PLEXATIC Date: Sat, 6 Apr 2024 18:12:46 +0200 Subject: [PATCH 3/4] Update README.md to reflect UI changes --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 481c8fb..5248c34 100644 --- a/README.md +++ b/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. From 1f4d1363cac4d23b61bdb3df874ed708c249f9e2 Mon Sep 17 00:00:00 2001 From: PLEXATIC Date: Sun, 7 Apr 2024 22:20:00 +0200 Subject: [PATCH 4/4] fixed algo --- nodes.py | 11 ++++------- utils/attention_functions.py | 10 ++++++---- utils/style_functions.py | 12 ------------ 3 files changed, 10 insertions(+), 23 deletions(-) diff --git a/nodes.py b/nodes.py index 8ab9edc..ea50486 100644 --- a/nodes.py +++ b/nodes.py @@ -94,8 +94,9 @@ class ApplyVisualStyle: "output": skip_output_layers } + class_names = set() for n, m in model.model.diffusion_model.named_modules(): - if m.__class__.__name__ == "CrossAttention": + if m.__class__.__name__ == "CrossAttention": is_enabled = self.activate_block_choice(n, block_choices) if is_enabled: block_name = n.split("_")[0] @@ -113,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}) diff --git a/utils/attention_functions.py b/utils/attention_functions.py index 65ab5f7..2424c53 100644 --- a/utils/attention_functions.py +++ b/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) diff --git a/utils/style_functions.py b/utils/style_functions.py index a41321e..75be7f8 100644 --- a/utils/style_functions.py +++ b/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)