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. diff --git a/nodes.py b/nodes.py index 88855d5..ea50486 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,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}) 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)