diff --git a/README.md b/README.md index 358bc52..ba8b994 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 (matching legacy functionality), you must use the same prompt for `positive` and `reference_cond`. diff --git a/nodes.py b/nodes.py index 6507d6e..b788bca 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",), @@ -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}) 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)