From c2d5b29040132c171bc4d77f1f63da972306f22c Mon Sep 17 00:00:00 2001
From: Jairo Correa <jn.j41r0@gmail.com>
Date: Thu, 29 Sep 2022 01:14:54 -0300
Subject: Move silu to sd_hijack

---
 modules/sd_hijack.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

(limited to 'modules/sd_hijack.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index bfbd07f9..4bc58fa2 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -12,6 +12,7 @@ from ldm.util import default
 from einops import rearrange
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
+from torch.nn.functional import silu
 
 
 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
@@ -100,14 +101,6 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 
     return self.to_out(r2)
 
-def nonlinearity_hijack(x):
-    # swish
-    t = torch.sigmoid(x)
-    x *= t
-    del t
-
-    return x
-
 def cross_attention_attnblock_forward(self, x):
         h_ = x
         h_ = self.norm(h_)
@@ -245,11 +238,12 @@ class StableDiffusionModelHijack:
         m.cond_stage_model = FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
         self.clip = m.cond_stage_model
 
+        ldm.modules.diffusionmodules.model.nonlinearity = silu
+
         if cmd_opts.opt_split_attention_v1:
             ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1
         elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention or torch.cuda.is_available()):
             ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward
-            ldm.modules.diffusionmodules.model.nonlinearity = nonlinearity_hijack
             ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
 
         def flatten(el):
-- 
cgit v1.2.3


From f7c787eb7c295c27439f4fbdf78c26b8389560be Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 7 Oct 2022 16:39:51 +0300
Subject: make it possible to use hypernetworks without opt split attention

---
 modules/hypernetwork.py | 42 ++++++++++++++++++++++++++++++++++--------
 modules/sd_hijack.py    |  6 ++++--
 2 files changed, 38 insertions(+), 10 deletions(-)

(limited to 'modules/sd_hijack.py')

diff --git a/modules/hypernetwork.py b/modules/hypernetwork.py
index c5cf4afa..c7b86682 100644
--- a/modules/hypernetwork.py
+++ b/modules/hypernetwork.py
@@ -4,7 +4,12 @@ import sys
 import traceback
 
 import torch
-from modules import devices
+
+from ldm.util import default
+from modules import devices, shared
+import torch
+from torch import einsum
+from einops import rearrange, repeat
 
 
 class HypernetworkModule(torch.nn.Module):
@@ -48,15 +53,36 @@ def load_hypernetworks(path):
 
     return res
 
-def apply(self, x, context=None, mask=None, original=None):
 
+def attention_CrossAttention_forward(self, x, context=None, mask=None):
+    h = self.heads
+
+    q = self.to_q(x)
+    context = default(context, x)
 
-    if CrossAttention.hypernetwork is not None and context.shape[2] in CrossAttention.hypernetwork:
-        if context.shape[1] == 77 and CrossAttention.noise_cond:
-            context = context + (torch.randn_like(context) * 0.1)
-        h_k, h_v = CrossAttention.hypernetwork[context.shape[2]]
-        k = self.to_k(h_k(context))
-        v = self.to_v(h_v(context))
+    hypernetwork = shared.selected_hypernetwork()
+    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
+
+    if hypernetwork_layers is not None:
+        k = self.to_k(hypernetwork_layers[0](context))
+        v = self.to_v(hypernetwork_layers[1](context))
     else:
         k = self.to_k(context)
         v = self.to_v(context)
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+    sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+    if mask is not None:
+        mask = rearrange(mask, 'b ... -> b (...)')
+        max_neg_value = -torch.finfo(sim.dtype).max
+        mask = repeat(mask, 'b j -> (b h) () j', h=h)
+        sim.masked_fill_(~mask, max_neg_value)
+
+    # attention, what we cannot get enough of
+    attn = sim.softmax(dim=-1)
+
+    out = einsum('b i j, b j d -> b i d', attn, v)
+    out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+    return self.to_out(out)
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index a6fa890c..d68f89cc 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -8,7 +8,7 @@ from torch import einsum
 from torch.nn.functional import silu
 
 import modules.textual_inversion.textual_inversion
-from modules import prompt_parser, devices, sd_hijack_optimizations, shared
+from modules import prompt_parser, devices, sd_hijack_optimizations, shared, hypernetwork
 from modules.shared import opts, device, cmd_opts
 
 import ldm.modules.attention
@@ -20,6 +20,8 @@ diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.At
 
 
 def apply_optimizations():
+    undo_optimizations()
+
     ldm.modules.diffusionmodules.model.nonlinearity = silu
 
     if cmd_opts.opt_split_attention_v1:
@@ -30,7 +32,7 @@ def apply_optimizations():
 
 
 def undo_optimizations():
-    ldm.modules.attention.CrossAttention.forward = attention_CrossAttention_forward
+    ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
     ldm.modules.diffusionmodules.model.nonlinearity = diffusionmodules_model_nonlinearity
     ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward
 
-- 
cgit v1.2.3