From 1f4d1363cac4d23b61bdb3df874ed708c249f9e2 Mon Sep 17 00:00:00 2001
From: PLEXATIC <plexatic@gmail.com>
Date: Sun, 7 Apr 2024 22:20:00 +0200
Subject: [PATCH] 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)