From 820f1dc96b1979d7e92170c161db281ee8bd988b Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 2 Oct 2022 15:03:39 +0300
Subject: initial support for training textual inversion

---
 modules/sd_hijack_optimizations.py | 164 +++++++++++++++++++++++++++++++++++++
 1 file changed, 164 insertions(+)
 create mode 100644 modules/sd_hijack_optimizations.py

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
new file mode 100644
index 00000000..9c079e57
--- /dev/null
+++ b/modules/sd_hijack_optimizations.py
@@ -0,0 +1,164 @@
+import math
+import torch
+from torch import einsum
+
+from ldm.util import default
+from einops import rearrange
+
+
+# see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
+def split_cross_attention_forward_v1(self, x, context=None, mask=None):
+    h = self.heads
+
+    q = self.to_q(x)
+    context = default(context, x)
+    k = self.to_k(context)
+    v = self.to_v(context)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
+    for i in range(0, q.shape[0], 2):
+        end = i + 2
+        s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])
+        s1 *= self.scale
+
+        s2 = s1.softmax(dim=-1)
+        del s1
+
+        r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])
+        del s2
+
+    r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
+    del r1
+
+    return self.to_out(r2)
+
+
+# taken from https://github.com/Doggettx/stable-diffusion
+def split_cross_attention_forward(self, x, context=None, mask=None):
+    h = self.heads
+
+    q_in = self.to_q(x)
+    context = default(context, x)
+    k_in = self.to_k(context) * self.scale
+    v_in = self.to_v(context)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
+
+    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
+
+    stats = torch.cuda.memory_stats(q.device)
+    mem_active = stats['active_bytes.all.current']
+    mem_reserved = stats['reserved_bytes.all.current']
+    mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
+    mem_free_torch = mem_reserved - mem_active
+    mem_free_total = mem_free_cuda + mem_free_torch
+
+    gb = 1024 ** 3
+    tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
+    modifier = 3 if q.element_size() == 2 else 2.5
+    mem_required = tensor_size * modifier
+    steps = 1
+
+    if mem_required > mem_free_total:
+        steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))
+        # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
+        #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
+
+    if steps > 64:
+        max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
+        raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
+                           f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
+
+    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)
+
+        s2 = s1.softmax(dim=-1, dtype=q.dtype)
+        del s1
+
+        r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
+        del s2
+
+    del q, k, v
+
+    r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
+    del r1
+
+    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_)
+        q1 = self.q(h_)
+        k1 = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q1.shape
+
+        q2 = q1.reshape(b, c, h*w)
+        del q1
+
+        q = q2.permute(0, 2, 1)   # b,hw,c
+        del q2
+
+        k = k1.reshape(b, c, h*w) # b,c,hw
+        del k1
+
+        h_ = torch.zeros_like(k, device=q.device)
+
+        stats = torch.cuda.memory_stats(q.device)
+        mem_active = stats['active_bytes.all.current']
+        mem_reserved = stats['reserved_bytes.all.current']
+        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
+        mem_free_torch = mem_reserved - mem_active
+        mem_free_total = mem_free_cuda + mem_free_torch
+
+        tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
+        mem_required = tensor_size * 2.5
+        steps = 1
+
+        if mem_required > mem_free_total:
+            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
+
+        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
+
+            w1 = torch.bmm(q[:, i:end], k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+            w2 = w1 * (int(c)**(-0.5))
+            del w1
+            w3 = torch.nn.functional.softmax(w2, dim=2, dtype=q.dtype)
+            del w2
+
+            # attend to values
+            v1 = v.reshape(b, c, h*w)
+            w4 = w3.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+            del w3
+
+            h_[:, :, i:end] = torch.bmm(v1, w4)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+            del v1, w4
+
+        h2 = h_.reshape(b, c, h, w)
+        del h_
+
+        h3 = self.proj_out(h2)
+        del h2
+
+        h3 += x
+
+        return h3
-- 
cgit v1.2.3


From f174fb29228a04955fb951b32b0bab79e33ec2b8 Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Fri, 7 Oct 2022 05:21:49 +0300
Subject: add xformers attention

---
 modules/sd_hijack_optimizations.py | 39 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index ea4cfdfc..da1b76e1 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,7 +1,9 @@
 import math
 import torch
 from torch import einsum
-
+import xformers.ops
+import functorch
+xformers._is_functorch_available=True
 from ldm.util import default
 from einops import rearrange
 
@@ -92,6 +94,41 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 
     return self.to_out(r2)
 
+def _maybe_init(self, x):
+    """
+    Initialize the attention operator, if required We expect the head dimension to be exposed here, meaning that x
+    : B, Head, Length
+    """
+    if self.attention_op is not None:
+        return
+    _, M, K = x.shape
+    try:
+        self.attention_op = xformers.ops.AttentionOpDispatch(
+            dtype=x.dtype,
+            device=x.device,
+            k=K,
+            attn_bias_type=type(None),
+            has_dropout=False,
+            kv_len=M,
+            q_len=M,
+        ).op
+    except NotImplementedError as err:
+        raise NotImplementedError(f"Please install xformers with the flash attention / cutlass components.\n{err}")
+
+def xformers_attention_forward(self, x, context=None, mask=None):
+    h = self.heads
+    q_in = self.to_q(x)
+    context = default(context, x)
+    k_in = self.to_k(context)
+    v_in = self.to_v(context)
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
+    self._maybe_init(q)
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+
+    out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+    return self.to_out(out)
+
 def cross_attention_attnblock_forward(self, x):
         h_ = x
         h_ = self.norm(h_)
-- 
cgit v1.2.3


From bad7cb29cecac51c5c0f39afec332b007ed73133 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 7 Oct 2022 10:17:52 +0300
Subject: added support for hypernetworks (???)

---
 modules/sd_hijack_optimizations.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index ea4cfdfc..d9cca485 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -5,6 +5,8 @@ from torch import einsum
 from ldm.util import default
 from einops import rearrange
 
+from modules import shared
+
 
 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
 def split_cross_attention_forward_v1(self, x, context=None, mask=None):
@@ -42,8 +44,19 @@ 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) * self.scale
-    v_in = self.to_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_in = self.to_k(hypernetwork_layers[0](context))
+        v_in = self.to_v(hypernetwork_layers[1](context))
+    else:
+        k_in = self.to_k(context)
+        v_in = self.to_v(context)
+
+    k_in *= self.scale
+
     del context, x
 
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
-- 
cgit v1.2.3


From c9cc65b201679ea43c763b0d85e749d40bbc5433 Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Sat, 8 Oct 2022 04:09:18 +0300
Subject: switch to the proper way of calling xformers

---
 modules/sd_hijack_optimizations.py | 28 +++-------------------------
 1 file changed, 3 insertions(+), 25 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index da1b76e1..7fb4a45e 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -94,39 +94,17 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 
     return self.to_out(r2)
 
-def _maybe_init(self, x):
-    """
-    Initialize the attention operator, if required We expect the head dimension to be exposed here, meaning that x
-    : B, Head, Length
-    """
-    if self.attention_op is not None:
-        return
-    _, M, K = x.shape
-    try:
-        self.attention_op = xformers.ops.AttentionOpDispatch(
-            dtype=x.dtype,
-            device=x.device,
-            k=K,
-            attn_bias_type=type(None),
-            has_dropout=False,
-            kv_len=M,
-            q_len=M,
-        ).op
-    except NotImplementedError as err:
-        raise NotImplementedError(f"Please install xformers with the flash attention / cutlass components.\n{err}")
-
 def xformers_attention_forward(self, x, context=None, mask=None):
     h = self.heads
     q_in = self.to_q(x)
     context = default(context, x)
     k_in = self.to_k(context)
     v_in = self.to_v(context)
-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
     del q_in, k_in, v_in
-    self._maybe_init(q)
-    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
 
-    out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+    out = rearrange(out, 'b n h d -> b n (h d)', h=h)
     return self.to_out(out)
 
 def cross_attention_attnblock_forward(self, x):
-- 
cgit v1.2.3


From f2055cb1d4ce45d7aaacc49d8ab5bec7791a8f47 Mon Sep 17 00:00:00 2001
From: brkirch <brkirch@users.noreply.github.com>
Date: Sat, 8 Oct 2022 01:47:02 -0400
Subject: Add hypernetwork support to split cross attention v1

* Add hypernetwork support to split_cross_attention_forward_v1
* Fix device check in esrgan_model.py to use devices.device_esrgan instead of shared.device
---
 modules/sd_hijack_optimizations.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index d9cca485..3351c740 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -12,13 +12,22 @@ from modules import shared
 def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     h = self.heads
 
-    q = self.to_q(x)
+    q_in = self.to_q(x)
     context = default(context, x)
-    k = self.to_k(context)
-    v = self.to_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_in = self.to_k(hypernetwork_layers[0](context))
+        v_in = self.to_v(hypernetwork_layers[1](context))
+    else:
+        k_in = self.to_k(context)
+        v_in = self.to_v(context)
     del context, x
 
-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
 
     r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
     for i in range(0, q.shape[0], 2):
@@ -31,6 +40,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
 
         r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])
         del s2
+    del q, k, v
 
     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
     del r1
-- 
cgit v1.2.3


From 5d54f35c583bd5a3b0ee271a862827f1ca81ef09 Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Sat, 8 Oct 2022 11:55:02 +0300
Subject: add xformers attnblock and hypernetwork support

---
 modules/sd_hijack_optimizations.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 7fb4a45e..c78d5838 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -98,8 +98,14 @@ def xformers_attention_forward(self, x, context=None, mask=None):
     h = self.heads
     q_in = self.to_q(x)
     context = default(context, x)
-    k_in = self.to_k(context)
-    v_in = self.to_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_in = self.to_k(hypernetwork_layers[0](context))
+        v_in = self.to_v(hypernetwork_layers[1](context))
+    else:
+        k_in = self.to_k(context)
+        v_in = self.to_v(context)
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
     del q_in, k_in, v_in
     out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
@@ -169,3 +175,13 @@ def cross_attention_attnblock_forward(self, x):
         h3 += x
 
         return h3
+    
+    def xformers_attnblock_forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q1 = self.q(h_).contiguous()
+        k1 = self.k(h_).contiguous()
+        v = self.v(h_).contiguous()
+        out = xformers.ops.memory_efficient_attention(q1, k1, v)
+        out = self.proj_out(out)
+        return x+out
-- 
cgit v1.2.3


From 76a616fa6b814c681eaf6edc87eb3001b8c2b6be Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Sat, 8 Oct 2022 11:55:38 +0300
Subject: Update sd_hijack_optimizations.py

---
 modules/sd_hijack_optimizations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index c78d5838..ee58c7e4 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -176,7 +176,7 @@ def cross_attention_attnblock_forward(self, x):
 
         return h3
     
-    def xformers_attnblock_forward(self, x):
+def xformers_attnblock_forward(self, x):
         h_ = x
         h_ = self.norm(h_)
         q1 = self.q(h_).contiguous()
-- 
cgit v1.2.3


From 69d0053583757ce2942d62de81e8b89e6be07840 Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Sat, 8 Oct 2022 16:21:40 +0300
Subject: update sd_hijack_opt to respect new env variables

---
 modules/sd_hijack_optimizations.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index ee58c7e4..be09ec8f 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,9 +1,14 @@
 import math
 import torch
 from torch import einsum
-import xformers.ops
-import functorch
-xformers._is_functorch_available=True
+try:
+    import xformers.ops
+    import functorch
+    xformers._is_functorch_available = True
+    shared.xformers_available = True
+except:
+    print('Cannot find xformers, defaulting to split attention. Try setting --xformers in your webui-user file if you wish to install it.')
+    continue
 from ldm.util import default
 from einops import rearrange
 
-- 
cgit v1.2.3


From 7ff1170a2e11b6f00f587407326db0b9f8f51adf Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 8 Oct 2022 16:33:39 +0300
Subject: emergency fix for xformers (continue + shared)

---
 modules/sd_hijack_optimizations.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index e43e2c7a..05023b6f 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,19 +1,19 @@
 import math
 import torch
 from torch import einsum
-try:
-    import xformers.ops
-    import functorch
-    xformers._is_functorch_available = True
-    shared.xformers_available = True
-except:
-    print('Cannot find xformers, defaulting to split attention. Try setting --xformers in your webui-user file if you wish to install it.')
-    continue
+
 from ldm.util import default
 from einops import rearrange
 
 from modules import shared
 
+try:
+    import xformers.ops
+    import functorch
+    xformers._is_functorch_available = True
+    shared.xformers_available = True
+except Exception:
+    print('Cannot find xformers, defaulting to split attention. Try adding --xformers commandline argument to your webui-user file if you wish to install it.')
 
 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
 def split_cross_attention_forward_v1(self, x, context=None, mask=None):
-- 
cgit v1.2.3


From dc1117233ef8f9b25ff1ac40b158f20b70ba2fcb Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 8 Oct 2022 17:02:18 +0300
Subject: simplify xfrmers options: --xformers to enable and that's it

---
 modules/sd_hijack_optimizations.py | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 05023b6f..d23d733b 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,4 +1,7 @@
 import math
+import sys
+import traceback
+
 import torch
 from torch import einsum
 
@@ -7,13 +10,16 @@ from einops import rearrange
 
 from modules import shared
 
-try:
-    import xformers.ops
-    import functorch
-    xformers._is_functorch_available = True
-    shared.xformers_available = True
-except Exception:
-    print('Cannot find xformers, defaulting to split attention. Try adding --xformers commandline argument to your webui-user file if you wish to install it.')
+if shared.cmd_opts.xformers:
+    try:
+        import xformers.ops
+        import functorch
+        xformers._is_functorch_available = True
+        shared.xformers_available = True
+    except Exception:
+        print("Cannot import xformers", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+
 
 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
 def split_cross_attention_forward_v1(self, x, context=None, mask=None):
-- 
cgit v1.2.3


From f9c5da159245bb1e7603b3c8b9e0703bcb1c2ff5 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 8 Oct 2022 19:05:19 +0300
Subject: add fallback for xformers_attnblock_forward

---
 modules/sd_hijack_optimizations.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index d23d733b..dba21192 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -211,6 +211,7 @@ def cross_attention_attnblock_forward(self, x):
         return h3
     
 def xformers_attnblock_forward(self, x):
+    try:
         h_ = x
         h_ = self.norm(h_)
         q1 = self.q(h_).contiguous()
@@ -218,4 +219,6 @@ def xformers_attnblock_forward(self, x):
         v = self.v(h_).contiguous()
         out = xformers.ops.memory_efficient_attention(q1, k1, v)
         out = self.proj_out(out)
-        return x+out
+        return x + out
+    except NotImplementedError:
+        return cross_attention_attnblock_forward(self, x)
-- 
cgit v1.2.3


From e6e42f98df2c928c4f49351ad6b466387ce87d42 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 8 Oct 2022 19:25:10 +0300
Subject: make --force-enable-xformers work without needing --xformers

---
 modules/sd_hijack_optimizations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index dba21192..c4396bb9 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -10,7 +10,7 @@ from einops import rearrange
 
 from modules import shared
 
-if shared.cmd_opts.xformers:
+if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
     try:
         import xformers.ops
         import functorch
-- 
cgit v1.2.3


From 122d42687b97ec4df4c2a8c335d2de385cd1f1a1 Mon Sep 17 00:00:00 2001
From: Fampai <unknown>
Date: Sat, 8 Oct 2022 22:37:35 -0400
Subject: Fix VRAM Issue by only loading in hypernetwork when selected in
 settings

---
 modules/sd_hijack_optimizations.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index c4396bb9..634fb4b2 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -28,7 +28,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     q_in = self.to_q(x)
     context = default(context, x)
 
-    hypernetwork = shared.selected_hypernetwork()
+    hypernetwork = shared.loaded_hypernetwork
     hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
 
     if hypernetwork_layers is not None:
@@ -68,7 +68,7 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
     q_in = self.to_q(x)
     context = default(context, x)
 
-    hypernetwork = shared.selected_hypernetwork()
+    hypernetwork = shared.loaded_hypernetwork
     hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
 
     if hypernetwork_layers is not None:
@@ -132,7 +132,7 @@ def xformers_attention_forward(self, x, context=None, mask=None):
     h = self.heads
     q_in = self.to_q(x)
     context = default(context, x)
-    hypernetwork = shared.selected_hypernetwork()
+    hypernetwork = shared.loaded_hypernetwork
     hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
     if hypernetwork_layers is not None:
         k_in = self.to_k(hypernetwork_layers[0](context))
-- 
cgit v1.2.3


From 3e7a981194ed9c454e951365846e4eba66fa7095 Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Mon, 10 Oct 2022 17:51:05 +0300
Subject: remove functorch

---
 modules/sd_hijack_optimizations.py | 2 --
 1 file changed, 2 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 634fb4b2..18408e62 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -13,8 +13,6 @@ from modules import shared
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
     try:
         import xformers.ops
-        import functorch
-        xformers._is_functorch_available = True
         shared.xformers_available = True
     except Exception:
         print("Cannot import xformers", file=sys.stderr)
-- 
cgit v1.2.3


From 948533950c9db5069a874d925fadd50bac00fdb5 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 11 Oct 2022 11:09:51 +0300
Subject: replace duplicate code with a function

---
 modules/sd_hijack_optimizations.py | 44 +++++++++++++-------------------------
 1 file changed, 15 insertions(+), 29 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 18408e62..25cb67a4 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -8,7 +8,8 @@ from torch import einsum
 from ldm.util import default
 from einops import rearrange
 
-from modules import shared
+from modules import shared, hypernetwork
+
 
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
     try:
@@ -26,16 +27,10 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     q_in = self.to_q(x)
     context = default(context, x)
 
-    hypernetwork = shared.loaded_hypernetwork
-    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
-
-    if hypernetwork_layers is not None:
-        k_in = self.to_k(hypernetwork_layers[0](context))
-        v_in = self.to_v(hypernetwork_layers[1](context))
-    else:
-        k_in = self.to_k(context)
-        v_in = self.to_v(context)
-    del context, x
+    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)
+    k_in = self.to_k(context_k)
+    v_in = self.to_v(context_v)
+    del context, context_k, context_v, x
 
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
     del q_in, k_in, v_in
@@ -59,22 +54,16 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     return self.to_out(r2)
 
 
-# taken from https://github.com/Doggettx/stable-diffusion
+# taken from https://github.com/Doggettx/stable-diffusion and modified
 def split_cross_attention_forward(self, x, context=None, mask=None):
     h = self.heads
 
     q_in = self.to_q(x)
     context = default(context, x)
 
-    hypernetwork = shared.loaded_hypernetwork
-    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
-
-    if hypernetwork_layers is not None:
-        k_in = self.to_k(hypernetwork_layers[0](context))
-        v_in = self.to_v(hypernetwork_layers[1](context))
-    else:
-        k_in = self.to_k(context)
-        v_in = self.to_v(context)
+    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)
+    k_in = self.to_k(context_k)
+    v_in = self.to_v(context_v)
 
     k_in *= self.scale
 
@@ -130,14 +119,11 @@ def xformers_attention_forward(self, x, context=None, mask=None):
     h = self.heads
     q_in = self.to_q(x)
     context = default(context, x)
-    hypernetwork = shared.loaded_hypernetwork
-    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
-    if hypernetwork_layers is not None:
-        k_in = self.to_k(hypernetwork_layers[0](context))
-        v_in = self.to_v(hypernetwork_layers[1](context))
-    else:
-        k_in = self.to_k(context)
-        v_in = self.to_v(context)
+
+    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)
+    k_in = self.to_k(context_k)
+    v_in = self.to_v(context_v)
+
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
     del q_in, k_in, v_in
     out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
-- 
cgit v1.2.3


From 530103b586109c11fd068eb70ef09503ec6a4caf Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 11 Oct 2022 14:53:02 +0300
Subject: fixes related to merge

---
 modules/sd_hijack_optimizations.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 25cb67a4..27e571fc 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -8,7 +8,8 @@ from torch import einsum
 from ldm.util import default
 from einops import rearrange
 
-from modules import shared, hypernetwork
+from modules import shared
+from modules.hypernetwork import hypernetwork
 
 
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
-- 
cgit v1.2.3


From 873efeed49bb5197a42da18272115b326c5d68f3 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 11 Oct 2022 15:51:22 +0300
Subject: rename hypernetwork dir to hypernetworks to prevent clash with an old
 filename that people who use zip instead of git clone will have

---
 modules/sd_hijack_optimizations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 27e571fc..3349b9c3 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -9,7 +9,7 @@ from ldm.util import default
 from einops import rearrange
 
 from modules import shared
-from modules.hypernetwork import hypernetwork
+from modules.hypernetworks import hypernetwork
 
 
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
-- 
cgit v1.2.3


From c0484f1b986ce7acb0e3596f6089a191279f5442 Mon Sep 17 00:00:00 2001
From: brkirch <brkirch@users.noreply.github.com>
Date: Mon, 10 Oct 2022 22:48:54 -0400
Subject: Add cross-attention optimization from InvokeAI

* Add cross-attention optimization from InvokeAI (~30% speed improvement on MPS)
* Add command line option for it
* Make it default when CUDA is unavailable
---
 modules/sd_hijack_optimizations.py | 79 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 3349b9c3..870226c5 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,6 +1,7 @@
 import math
 import sys
 import traceback
+import psutil
 
 import torch
 from torch import einsum
@@ -116,6 +117,84 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 
     return self.to_out(r2)
 
+# -- From https://github.com/invoke-ai/InvokeAI/blob/main/ldm/modules/attention.py (with hypernetworks support added) --
+
+mem_total_gb = psutil.virtual_memory().total // (1 << 30)
+
+def einsum_op_compvis(q, k, v):
+    s = einsum('b i d, b j d -> b i j', q, k)
+    s = s.softmax(dim=-1, dtype=s.dtype)
+    return einsum('b i j, b j d -> b i d', s, v)
+
+def einsum_op_slice_0(q, k, v, slice_size):
+    r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
+    for i in range(0, q.shape[0], slice_size):
+        end = i + slice_size
+        r[i:end] = einsum_op_compvis(q[i:end], k[i:end], v[i:end])
+    return r
+
+def einsum_op_slice_1(q, k, v, slice_size):
+    r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
+    for i in range(0, q.shape[1], slice_size):
+        end = i + slice_size
+        r[:, i:end] = einsum_op_compvis(q[:, i:end], k, v)
+    return r
+
+def einsum_op_mps_v1(q, k, v):
+    if q.shape[1] <= 4096: # (512x512) max q.shape[1]: 4096
+        return einsum_op_compvis(q, k, v)
+    else:
+        slice_size = math.floor(2**30 / (q.shape[0] * q.shape[1]))
+        return einsum_op_slice_1(q, k, v, slice_size)
+
+def einsum_op_mps_v2(q, k, v):
+    if mem_total_gb > 8 and q.shape[1] <= 4096:
+        return einsum_op_compvis(q, k, v)
+    else:
+        return einsum_op_slice_0(q, k, v, 1)
+
+def einsum_op_tensor_mem(q, k, v, max_tensor_mb):
+    size_mb = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() // (1 << 20)
+    if size_mb <= max_tensor_mb:
+        return einsum_op_compvis(q, k, v)
+    div = 1 << int((size_mb - 1) / max_tensor_mb).bit_length()
+    if div <= q.shape[0]:
+        return einsum_op_slice_0(q, k, v, q.shape[0] // div)
+    return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1))
+
+def einsum_op(q, k, v):
+    if q.device.type == 'mps':
+        if mem_total_gb >= 32:
+            return einsum_op_mps_v1(q, k, v)
+        return einsum_op_mps_v2(q, k, v)
+
+    # Smaller slices are faster due to L2/L3/SLC caches.
+    # Tested on i7 with 8MB L3 cache.
+    return einsum_op_tensor_mem(q, k, v, 32)
+
+def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
+    h = self.heads
+
+    q = self.to_q(x)
+    context = default(context, x)
+
+    hypernetwork = shared.loaded_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)) * self.scale
+        v = self.to_v(hypernetwork_layers[1](context))
+    else:
+        k = self.to_k(context) * self.scale
+        v = self.to_v(context)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+    r = einsum_op(q, k, v)
+    return self.to_out(rearrange(r, '(b h) n d -> b n (h d)', h=h))
+
+# -- End of code from https://github.com/invoke-ai/InvokeAI/blob/main/ldm/modules/attention.py --
+
 def xformers_attention_forward(self, x, context=None, mask=None):
     h = self.heads
     q_in = self.to_q(x)
-- 
cgit v1.2.3


From 98fd5cde72d5bda1620ab78416c7828fdc3dc10b Mon Sep 17 00:00:00 2001
From: brkirch <brkirch@users.noreply.github.com>
Date: Mon, 10 Oct 2022 23:55:48 -0400
Subject: Add check for psutil

---
 modules/sd_hijack_optimizations.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 870226c5..2a4ac7e0 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,7 +1,7 @@
 import math
 import sys
 import traceback
-import psutil
+import importlib
 
 import torch
 from torch import einsum
@@ -117,9 +117,20 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 
     return self.to_out(r2)
 
-# -- From https://github.com/invoke-ai/InvokeAI/blob/main/ldm/modules/attention.py (with hypernetworks support added) --
 
-mem_total_gb = psutil.virtual_memory().total // (1 << 30)
+def check_for_psutil():
+    try:
+        spec = importlib.util.find_spec('psutil')
+        return spec is not None
+    except ModuleNotFoundError:
+        return False
+
+invokeAI_mps_available = check_for_psutil()
+
+# -- Taken from https://github.com/invoke-ai/InvokeAI --
+if invokeAI_mps_available:
+    import psutil
+    mem_total_gb = psutil.virtual_memory().total // (1 << 30)
 
 def einsum_op_compvis(q, k, v):
     s = einsum('b i d, b j d -> b i j', q, k)
@@ -193,7 +204,7 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
     r = einsum_op(q, k, v)
     return self.to_out(rearrange(r, '(b h) n d -> b n (h d)', h=h))
 
-# -- End of code from https://github.com/invoke-ai/InvokeAI/blob/main/ldm/modules/attention.py --
+# -- End of code from https://github.com/invoke-ai/InvokeAI --
 
 def xformers_attention_forward(self, x, context=None, mask=None):
     h = self.heads
-- 
cgit v1.2.3


From 574c8e554a5371eca2cbf344764cb241c6ec4efc Mon Sep 17 00:00:00 2001
From: brkirch <brkirch@users.noreply.github.com>
Date: Tue, 11 Oct 2022 03:32:11 -0400
Subject: Add InvokeAI and lstein to credits, add back CUDA support

---
 modules/sd_hijack_optimizations.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 2a4ac7e0..f006427f 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -173,7 +173,20 @@ def einsum_op_tensor_mem(q, k, v, max_tensor_mb):
         return einsum_op_slice_0(q, k, v, q.shape[0] // div)
     return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1))
 
+def einsum_op_cuda(q, k, v):
+    stats = torch.cuda.memory_stats(q.device)
+    mem_active = stats['active_bytes.all.current']
+    mem_reserved = stats['reserved_bytes.all.current']
+    mem_free_cuda, _ = torch.cuda.mem_get_info(q.device)
+    mem_free_torch = mem_reserved - mem_active
+    mem_free_total = mem_free_cuda + mem_free_torch
+    # Divide factor of safety as there's copying and fragmentation
+    return self.einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
+
 def einsum_op(q, k, v):
+    if q.device.type == 'cuda':
+        return einsum_op_cuda(q, k, v)
+
     if q.device.type == 'mps':
         if mem_total_gb >= 32:
             return einsum_op_mps_v1(q, k, v)
-- 
cgit v1.2.3


From 861db783c7acfcb93cf0b5191db3d50f9a9bc531 Mon Sep 17 00:00:00 2001
From: brkirch <brkirch@users.noreply.github.com>
Date: Tue, 11 Oct 2022 05:13:17 -0400
Subject: Use apply_hypernetwork function

---
 modules/sd_hijack_optimizations.py | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index f006427f..79405525 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -202,16 +202,10 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
     q = self.to_q(x)
     context = default(context, x)
 
-    hypernetwork = shared.loaded_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)) * self.scale
-        v = self.to_v(hypernetwork_layers[1](context))
-    else:
-        k = self.to_k(context) * self.scale
-        v = self.to_v(context)
-    del context, x
+    context_k, context_v = hypernetwork.apply_hypernetwork(shared.loaded_hypernetwork, context)
+    k = self.to_k(context_k) * self.scale
+    v = self.to_v(context_v)
+    del context, context_k, context_v, x
 
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
     r = einsum_op(q, k, v)
-- 
cgit v1.2.3


From 2043c4a231eef838bb15044f502b864b55885037 Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Mon, 17 Oct 2022 19:49:11 +0300
Subject: delete xformers attnblock

---
 modules/sd_hijack_optimizations.py | 12 ------------
 1 file changed, 12 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 79405525..60da7459 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -292,15 +292,3 @@ def cross_attention_attnblock_forward(self, x):
 
         return h3
     
-def xformers_attnblock_forward(self, x):
-    try:
-        h_ = x
-        h_ = self.norm(h_)
-        q1 = self.q(h_).contiguous()
-        k1 = self.k(h_).contiguous()
-        v = self.v(h_).contiguous()
-        out = xformers.ops.memory_efficient_attention(q1, k1, v)
-        out = self.proj_out(out)
-        return x + out
-    except NotImplementedError:
-        return cross_attention_attnblock_forward(self, x)
-- 
cgit v1.2.3


From 84823275e896bcc1f7cb4ce098ae3c5d05e17b9a Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Mon, 17 Oct 2022 22:18:59 +0300
Subject: readd xformers attnblock

---
 modules/sd_hijack_optimizations.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 60da7459..7ebef3f0 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -292,3 +292,18 @@ def cross_attention_attnblock_forward(self, x):
 
         return h3
     
+def xformers_attnblock_forward(self, x):
+    try:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+        b, c, h, w = q.shape
+        q, k, v = map(lambda t: rearrange(t, 'b c h w -> b (h w) c'), (q, k, v))
+        out = xformers.ops.memory_efficient_attention(q, k, v)
+        out = rearrange(out, 'b (h w) c -> b c h w', h=h)
+        out = self.proj_out(out)
+        return x + out
+    except NotImplementedError:
+        return cross_attention_attnblock_forward(self, x)
-- 
cgit v1.2.3


From c71008c74156635558bb2e877d1628913f6f781e Mon Sep 17 00:00:00 2001
From: C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com>
Date: Tue, 18 Oct 2022 00:02:50 +0300
Subject: Update sd_hijack_optimizations.py

---
 modules/sd_hijack_optimizations.py | 3 +++
 1 file changed, 3 insertions(+)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 7ebef3f0..a3345bb9 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -301,6 +301,9 @@ def xformers_attnblock_forward(self, x):
         v = self.v(h_)
         b, c, h, w = q.shape
         q, k, v = map(lambda t: rearrange(t, 'b c h w -> b (h w) c'), (q, k, v))
+        q = q.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
         out = xformers.ops.memory_efficient_attention(q, k, v)
         out = rearrange(out, 'b (h w) c -> b c h w', h=h)
         out = self.proj_out(out)
-- 
cgit v1.2.3


From 2fd7935ef4ed296db5dfd8c7fea99244816f8cf0 Mon Sep 17 00:00:00 2001
From: Cheka <carlos93@hotmail.cl>
Date: Tue, 18 Oct 2022 20:28:28 -0300
Subject: Remove wrong self reference in CUDA support for invokeai

---
 modules/sd_hijack_optimizations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index a3345bb9..98123fbf 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -181,7 +181,7 @@ def einsum_op_cuda(q, k, v):
     mem_free_torch = mem_reserved - mem_active
     mem_free_total = mem_free_cuda + mem_free_torch
     # Divide factor of safety as there's copying and fragmentation
-    return self.einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
+    return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
 
 def einsum_op(q, k, v):
     if q.device.type == 'cuda':
-- 
cgit v1.2.3


From adb6cb7619989cbc7a271cc6c2ae27bb936c43d9 Mon Sep 17 00:00:00 2001
From: Billy Cao <aliencaocao@gmail.com>
Date: Wed, 23 Nov 2022 18:11:24 +0800
Subject: Patch UNet Forward to support resolutions that are not multiples of
 64 Also modifed the UI to no longer step in 64

---
 modules/sd_hijack_optimizations.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 98123fbf..8cd4c954 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -5,6 +5,7 @@ import importlib
 
 import torch
 from torch import einsum
+import torch.nn.functional as F
 
 from ldm.util import default
 from einops import rearrange
@@ -12,6 +13,8 @@ from einops import rearrange
 from modules import shared
 from modules.hypernetworks import hypernetwork
 
+from ldm.modules.diffusionmodules.util import timestep_embedding
+
 
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
     try:
@@ -310,3 +313,31 @@ def xformers_attnblock_forward(self, x):
         return x + out
     except NotImplementedError:
         return cross_attention_attnblock_forward(self, x)
+
+def patched_unet_forward(self, x, timesteps=None, context=None, y=None,**kwargs):
+    assert (y is not None) == (
+        self.num_classes is not None
+    ), "must specify y if and only if the model is class-conditional"
+    hs = []
+    t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+    emb = self.time_embed(t_emb)
+
+    if self.num_classes is not None:
+        assert y.shape == (x.shape[0],)
+        emb = emb + self.label_emb(y)
+
+    h = x.type(self.dtype)
+    for module in self.input_blocks:
+        h = module(h, emb, context)
+        hs.append(h)
+    h = self.middle_block(h, emb, context)
+    for module in self.output_blocks:
+        if h.shape[-2:] != hs[-1].shape[-2:]:
+            h = F.interpolate(h, hs[-1].shape[-2:], mode="nearest")
+        h = torch.cat([h, hs.pop()], dim=1)
+        h = module(h, emb, context)
+    h = h.type(x.dtype)
+    if self.predict_codebook_ids:
+        return self.id_predictor(h)
+    else:
+        return self.out(h)
-- 
cgit v1.2.3


From 7dbfd8a7d8aefec7283b456c6f5b000ae4d3496d Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 10 Dec 2022 09:14:30 +0300
Subject: do not replace entire unet for the resolution hack

---
 modules/sd_hijack_optimizations.py | 28 ----------------------------
 1 file changed, 28 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 8cd4c954..85909eb9 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -313,31 +313,3 @@ def xformers_attnblock_forward(self, x):
         return x + out
     except NotImplementedError:
         return cross_attention_attnblock_forward(self, x)
-
-def patched_unet_forward(self, x, timesteps=None, context=None, y=None,**kwargs):
-    assert (y is not None) == (
-        self.num_classes is not None
-    ), "must specify y if and only if the model is class-conditional"
-    hs = []
-    t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
-    emb = self.time_embed(t_emb)
-
-    if self.num_classes is not None:
-        assert y.shape == (x.shape[0],)
-        emb = emb + self.label_emb(y)
-
-    h = x.type(self.dtype)
-    for module in self.input_blocks:
-        h = module(h, emb, context)
-        hs.append(h)
-    h = self.middle_block(h, emb, context)
-    for module in self.output_blocks:
-        if h.shape[-2:] != hs[-1].shape[-2:]:
-            h = F.interpolate(h, hs[-1].shape[-2:], mode="nearest")
-        h = torch.cat([h, hs.pop()], dim=1)
-        h = module(h, emb, context)
-    h = h.type(x.dtype)
-    if self.predict_codebook_ids:
-        return self.id_predictor(h)
-    else:
-        return self.out(h)
-- 
cgit v1.2.3


From 505ec7e4d960e7bea579182509050fafb10bd00c Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 10 Dec 2022 09:17:39 +0300
Subject: cleanup some unneeded imports for hijack files

---
 modules/sd_hijack_optimizations.py | 3 ---
 1 file changed, 3 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 85909eb9..98123fbf 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -5,7 +5,6 @@ import importlib
 
 import torch
 from torch import einsum
-import torch.nn.functional as F
 
 from ldm.util import default
 from einops import rearrange
@@ -13,8 +12,6 @@ from einops import rearrange
 from modules import shared
 from modules.hypernetworks import hypernetwork
 
-from ldm.modules.diffusionmodules.util import timestep_embedding
-
 
 if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
     try:
-- 
cgit v1.2.3


From 35b1775b32a07f1b7c9dccad61f7aa77027a00fa Mon Sep 17 00:00:00 2001
From: brkirch <brkirch@users.noreply.github.com>
Date: Mon, 19 Dec 2022 17:25:14 -0500
Subject: Use other MPS optimization for large q.shape[0] * q.shape[1]

Check if q.shape[0] * q.shape[1] is 2**18 or larger and use the lower memory usage MPS optimization if it is. This should prevent most crashes that were occurring at certain resolutions (e.g. 1024x1024, 2048x512, 512x2048).

Also included is a change to check slice_size and prevent it from being divisible by 4096 which also results in a crash. Otherwise a crash can occur at 1024x512 or 512x1024 resolution.
---
 modules/sd_hijack_optimizations.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 98123fbf..02c87f40 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -127,7 +127,7 @@ def check_for_psutil():
 
 invokeAI_mps_available = check_for_psutil()
 
-# -- Taken from https://github.com/invoke-ai/InvokeAI --
+# -- Taken from https://github.com/invoke-ai/InvokeAI and modified --
 if invokeAI_mps_available:
     import psutil
     mem_total_gb = psutil.virtual_memory().total // (1 << 30)
@@ -152,14 +152,16 @@ def einsum_op_slice_1(q, k, v, slice_size):
     return r
 
 def einsum_op_mps_v1(q, k, v):
-    if q.shape[1] <= 4096: # (512x512) max q.shape[1]: 4096
+    if q.shape[0] * q.shape[1] <= 2**16: # (512x512) max q.shape[1]: 4096
         return einsum_op_compvis(q, k, v)
     else:
         slice_size = math.floor(2**30 / (q.shape[0] * q.shape[1]))
+        if slice_size % 4096 == 0:
+            slice_size -= 1
         return einsum_op_slice_1(q, k, v, slice_size)
 
 def einsum_op_mps_v2(q, k, v):
-    if mem_total_gb > 8 and q.shape[1] <= 4096:
+    if mem_total_gb > 8 and q.shape[0] * q.shape[1] <= 2**16:
         return einsum_op_compvis(q, k, v)
     else:
         return einsum_op_slice_0(q, k, v, 1)
@@ -188,7 +190,7 @@ def einsum_op(q, k, v):
         return einsum_op_cuda(q, k, v)
 
     if q.device.type == 'mps':
-        if mem_total_gb >= 32:
+        if mem_total_gb >= 32 and q.shape[0] % 32 != 0 and q.shape[0] * q.shape[1] < 2**18:
             return einsum_op_mps_v1(q, k, v)
         return einsum_op_mps_v2(q, k, v)
 
-- 
cgit v1.2.3