Merge pull request #635 from C43H66N12O12S2/attention

Move scale multiplication to the front

1 files changed, 2 insertions, 2 deletions

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 65414518..c4450ce4 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -50,7 +50,7 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 

     q_in = self.to_q(x)

     context = default(context, x)

-    k_in = self.to_k(context)

+    k_in = self.to_k(context) * self.scale

     v_in = self.to_v(context)

     del context, x

 

@@ -85,7 +85,7 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
     slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]

     for i in range(0, q.shape[1], slice_size):

         end = i + slice_size

-        s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale

+        s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)

 

         s2 = s1.softmax(dim=-1, dtype=q.dtype)

         del s1