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/hypernetwork.py            | 55 ++++++++++++++++++++++++++++++++++++++
 modules/sd_hijack_optimizations.py | 17 ++++++++++--
 modules/shared.py                  |  9 ++++++-
 scripts/xy_grid.py                 | 10 +++++++
 4 files changed, 88 insertions(+), 3 deletions(-)
 create mode 100644 modules/hypernetwork.py

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/hypernetwork.py b/modules/hypernetwork.py
new file mode 100644
index 00000000..9ed1eed9
--- /dev/null
+++ b/modules/hypernetwork.py
@@ -0,0 +1,55 @@
+import glob
+import os
+import torch
+from modules import devices
+
+
+class HypernetworkModule(torch.nn.Module):
+    def __init__(self, dim, state_dict):
+        super().__init__()
+
+        self.linear1 = torch.nn.Linear(dim, dim * 2)
+        self.linear2 = torch.nn.Linear(dim * 2, dim)
+
+        self.load_state_dict(state_dict, strict=True)
+        self.to(devices.device)
+
+    def forward(self, x):
+        return x + (self.linear2(self.linear1(x)))
+
+
+class Hypernetwork:
+    filename = None
+    name = None
+
+    def __init__(self, filename):
+        self.filename = filename
+        self.name = os.path.splitext(os.path.basename(filename))[0]
+        self.layers = {}
+
+        state_dict = torch.load(filename, map_location='cpu')
+        for size, sd in state_dict.items():
+            self.layers[size] = (HypernetworkModule(size, sd[0]), HypernetworkModule(size, sd[1]))
+
+
+def load_hypernetworks(path):
+    res = {}
+
+    for filename in glob.iglob(path + '**/*.pt', recursive=True):
+        hn = Hypernetwork(filename)
+        res[hn.name] = hn
+
+    return res
+
+def apply(self, x, context=None, mask=None, original=None):
+
+
+    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))
+    else:
+        k = self.to_k(context)
+        v = self.to_v(context)
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))
diff --git a/modules/shared.py b/modules/shared.py
index 25bb6e6c..879d8424 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -13,7 +13,7 @@ import modules.memmon
 import modules.sd_models
 import modules.styles
 import modules.devices as devices
-from modules import sd_samplers
+from modules import sd_samplers, hypernetwork
 from modules.paths import models_path, script_path, sd_path
 
 sd_model_file = os.path.join(script_path, 'model.ckpt')
@@ -76,6 +76,12 @@ parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 
 config_filename = cmd_opts.ui_settings_file
 
+hypernetworks = hypernetwork.load_hypernetworks(os.path.join(models_path, 'hypernetworks'))
+
+
+def selected_hypernetwork():
+    return hypernetworks.get(opts.sd_hypernetwork, None)
+
 
 class State:
     interrupted = False
@@ -206,6 +212,7 @@ options_templates.update(options_section(('system', "System"), {
 
 options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": modules.sd_models.checkpoint_tiles()}),
+    "sd_hypernetwork": OptionInfo("None", "Stable Diffusion finetune hypernetwork", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}),
     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),
     "save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"),
     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies (normally you'd do less with less denoising)."),
diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index 6344e612..c0c364df 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -77,6 +77,11 @@ def apply_checkpoint(p, x, xs):
     modules.sd_models.reload_model_weights(shared.sd_model, info)
 
 
+def apply_hypernetwork(p, x, xs):
+    hn = shared.hypernetworks.get(x, None)
+    opts.data["sd_hypernetwork"] = hn.name if hn is not None else 'None'
+
+
 def format_value_add_label(p, opt, x):
     if type(x) == float:
         x = round(x, 8)
@@ -122,6 +127,7 @@ axis_options = [
     AxisOption("Prompt order", str_permutations, apply_order, format_value_join_list),
     AxisOption("Sampler", str, apply_sampler, format_value),
     AxisOption("Checkpoint name", str, apply_checkpoint, format_value),
+    AxisOption("Hypernetwork", str, apply_hypernetwork, format_value),
     AxisOption("Sigma Churn", float, apply_field("s_churn"), format_value_add_label),
     AxisOption("Sigma min", float, apply_field("s_tmin"), format_value_add_label),
     AxisOption("Sigma max", float, apply_field("s_tmax"), format_value_add_label),
@@ -193,6 +199,8 @@ class Script(scripts.Script):
         modules.processing.fix_seed(p)
         p.batch_size = 1
 
+        initial_hn = opts.sd_hypernetwork
+
         def process_axis(opt, vals):
             if opt.label == 'Nothing':
                 return [0]
@@ -300,4 +308,6 @@ class Script(scripts.Script):
         # restore checkpoint in case it was changed by axes
         modules.sd_models.reload_model_weights(shared.sd_model)
 
+        opts.data["sd_hypernetwork"] = initial_hn
+
         return processed
-- 
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/esrgan_model.py            |  2 +-
 modules/sd_hijack_optimizations.py | 18 ++++++++++++++----
 2 files changed, 15 insertions(+), 5 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index d17e730f..28548124 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -111,7 +111,7 @@ class UpscalerESRGAN(Upscaler):
             print("Unable to load %s from %s" % (self.model_path, filename))
             return None
 
-        pretrained_net = torch.load(filename, map_location='cpu' if shared.device.type == 'mps' else None)
+        pretrained_net = torch.load(filename, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)
         crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)
 
         pretrained_net = fix_model_layers(crt_model, pretrained_net)
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

---
 launch.py                          |  2 +-
 modules/sd_hijack.py               |  2 +-
 modules/sd_hijack_optimizations.py | 20 +++++++++++++-------
 modules/shared.py                  |  2 +-
 4 files changed, 16 insertions(+), 10 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/launch.py b/launch.py
index a592e1ba..61f62096 100644
--- a/launch.py
+++ b/launch.py
@@ -32,7 +32,7 @@ def extract_arg(args, name):
 
 
 args, skip_torch_cuda_test = extract_arg(args, '--skip-torch-cuda-test')
-args, xformers = extract_arg(args, '--xformers')
+xformers = '--xformers' in args
 
 
 def repo_dir(name):
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 5d93f7f6..91e98c16 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -22,7 +22,7 @@ def apply_optimizations():
     undo_optimizations()
 
     ldm.modules.diffusionmodules.model.nonlinearity = silu
-    if not cmd_opts.disable_opt_xformers_attention and not (cmd_opts.opt_split_attention or torch.version.hip) and shared.xformers_available:
+    if cmd_opts.xformers and shared.xformers_available and not torch.version.hip:
         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward
     elif cmd_opts.opt_split_attention_v1:
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):
diff --git a/modules/shared.py b/modules/shared.py
index d68df751..02cb2722 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -43,7 +43,7 @@ parser.add_argument("--realesrgan-models-path", type=str, help="Path to director
 parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(models_path, 'ScuNET'))
 parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(models_path, 'SwinIR'))
 parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(models_path, 'LDSR'))
-parser.add_argument("--disable-opt-xformers-attention", action='store_true', help="force-disables xformers attention optimization")
+parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")
 parser.add_argument("--opt-split-attention", action='store_true', help="force-enables cross-attention layer optimization. By default, it's on for torch.cuda and off for other torch devices.")
 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
-- 
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/hypernetwork.py            | 23 +++++++++++++++--------
 modules/sd_hijack_optimizations.py |  6 +++---
 modules/shared.py                  |  7 ++-----
 webui.py                           |  3 +++
 4 files changed, 23 insertions(+), 16 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/hypernetwork.py b/modules/hypernetwork.py
index 7f062242..19f1c227 100644
--- a/modules/hypernetwork.py
+++ b/modules/hypernetwork.py
@@ -40,18 +40,25 @@ class Hypernetwork:
             self.layers[size] = (HypernetworkModule(size, sd[0]), HypernetworkModule(size, sd[1]))
 
 
-def load_hypernetworks(path):
+def list_hypernetworks(path):
     res = {}
-
     for filename in glob.iglob(os.path.join(path, '**/*.pt'), recursive=True):
+        name = os.path.splitext(os.path.basename(filename))[0]
+        res[name] = filename
+    return res
+
+
+def load_hypernetwork(filename):
+    print(f"Loading hypernetwork {filename}")
+    path = shared.hypernetworks.get(filename, None)
+    if (path is not None):
         try:
-            hn = Hypernetwork(filename)
-            res[hn.name] = hn
+            shared.loaded_hypernetwork = Hypernetwork(path)
         except Exception:
-            print(f"Error loading hypernetwork {filename}", file=sys.stderr)
+            print(f"Error loading hypernetwork {path}", file=sys.stderr)
             print(traceback.format_exc(), file=sys.stderr)
-
-    return res
+    else:
+        shared.loaded_hypernetwork = None
 
 
 def attention_CrossAttention_forward(self, x, context=None, mask=None):
@@ -60,7 +67,7 @@ def attention_CrossAttention_forward(self, x, context=None, mask=None):
     q = 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:
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))
diff --git a/modules/shared.py b/modules/shared.py
index b2c76a32..9dce6cb7 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -79,11 +79,8 @@ parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 xformers_available = False
 config_filename = cmd_opts.ui_settings_file
 
-hypernetworks = hypernetwork.load_hypernetworks(os.path.join(models_path, 'hypernetworks'))
-
-
-def selected_hypernetwork():
-    return hypernetworks.get(opts.sd_hypernetwork, None)
+hypernetworks = hypernetwork.list_hypernetworks(os.path.join(models_path, 'hypernetworks'))
+loaded_hypernetwork = None
 
 
 class State:
diff --git a/webui.py b/webui.py
index 18de8e16..270584f7 100644
--- a/webui.py
+++ b/webui.py
@@ -82,6 +82,9 @@ modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
 shared.sd_model = modules.sd_models.load_model()
 shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))
 
+loaded_hypernetwork = modules.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)
+shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)))
+
 
 def webui():
     # make the program just exit at ctrl+c without waiting for anything
-- 
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/hypernetwork.py            | 23 ++++++++++++--------
 modules/sd_hijack_optimizations.py | 44 +++++++++++++-------------------------
 2 files changed, 29 insertions(+), 38 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/hypernetwork.py b/modules/hypernetwork.py
index 498bc9d8..7bbc443e 100644
--- a/modules/hypernetwork.py
+++ b/modules/hypernetwork.py
@@ -64,21 +64,26 @@ def load_hypernetwork(filename):
         shared.loaded_hypernetwork = None
 
 
+def apply_hypernetwork(hypernetwork, context):
+    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
+
+    if hypernetwork_layers is None:
+        return context, context
+
+    context_k = hypernetwork_layers[0](context)
+    context_v = hypernetwork_layers[1](context)
+    return context_k, context_v
+
+
 def attention_CrossAttention_forward(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))
-        v = self.to_v(hypernetwork_layers[1](context))
-    else:
-        k = self.to_k(context)
-        v = self.to_v(context)
+    context_k, context_v = apply_hypernetwork(shared.loaded_hypernetwork, context)
+    k = self.to_k(context_k)
+    v = self.to_v(context_v)
 
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
 
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/hypernetwork.py                        | 103 -------------------------
 modules/hypernetwork/hypernetwork.py           |  74 +++++++++++-------
 modules/hypernetwork/ui.py                     |  10 +--
 modules/sd_hijack_optimizations.py             |   3 +-
 modules/shared.py                              |  13 +++-
 modules/textual_inversion/textual_inversion.py |  12 +--
 modules/ui.py                                  |   5 +-
 scripts/xy_grid.py                             |   3 +-
 webui.py                                       |  15 +---
 9 files changed, 78 insertions(+), 160 deletions(-)
 delete mode 100644 modules/hypernetwork.py

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/hypernetwork.py b/modules/hypernetwork.py
deleted file mode 100644
index 7bbc443e..00000000
--- a/modules/hypernetwork.py
+++ /dev/null
@@ -1,103 +0,0 @@
-import glob
-import os
-import sys
-import traceback
-
-import torch
-
-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):
-    def __init__(self, dim, state_dict):
-        super().__init__()
-
-        self.linear1 = torch.nn.Linear(dim, dim * 2)
-        self.linear2 = torch.nn.Linear(dim * 2, dim)
-
-        self.load_state_dict(state_dict, strict=True)
-        self.to(devices.device)
-
-    def forward(self, x):
-        return x + (self.linear2(self.linear1(x)))
-
-
-class Hypernetwork:
-    filename = None
-    name = None
-
-    def __init__(self, filename):
-        self.filename = filename
-        self.name = os.path.splitext(os.path.basename(filename))[0]
-        self.layers = {}
-
-        state_dict = torch.load(filename, map_location='cpu')
-        for size, sd in state_dict.items():
-            self.layers[size] = (HypernetworkModule(size, sd[0]), HypernetworkModule(size, sd[1]))
-
-
-def list_hypernetworks(path):
-    res = {}
-    for filename in glob.iglob(os.path.join(path, '**/*.pt'), recursive=True):
-        name = os.path.splitext(os.path.basename(filename))[0]
-        res[name] = filename
-    return res
-
-
-def load_hypernetwork(filename):
-    path = shared.hypernetworks.get(filename, None)
-    if path is not None:
-        print(f"Loading hypernetwork {filename}")
-        try:
-            shared.loaded_hypernetwork = Hypernetwork(path)
-        except Exception:
-            print(f"Error loading hypernetwork {path}", file=sys.stderr)
-            print(traceback.format_exc(), file=sys.stderr)
-    else:
-        if shared.loaded_hypernetwork is not None:
-            print(f"Unloading hypernetwork")
-
-        shared.loaded_hypernetwork = None
-
-
-def apply_hypernetwork(hypernetwork, context):
-    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
-
-    if hypernetwork_layers is None:
-        return context, context
-
-    context_k = hypernetwork_layers[0](context)
-    context_v = hypernetwork_layers[1](context)
-    return context_k, context_v
-
-
-def attention_CrossAttention_forward(self, x, context=None, mask=None):
-    h = self.heads
-
-    q = self.to_q(x)
-    context = default(context, x)
-
-    context_k, context_v = apply_hypernetwork(shared.loaded_hypernetwork, context)
-    k = self.to_k(context_k)
-    v = self.to_v(context_v)
-
-    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/hypernetwork/hypernetwork.py b/modules/hypernetwork/hypernetwork.py
index a3d6a47e..aa701bda 100644
--- a/modules/hypernetwork/hypernetwork.py
+++ b/modules/hypernetwork/hypernetwork.py
@@ -26,10 +26,11 @@ class HypernetworkModule(torch.nn.Module):
         if state_dict is not None:
             self.load_state_dict(state_dict, strict=True)
         else:
-            self.linear1.weight.data.fill_(0.0001)
-            self.linear1.bias.data.fill_(0.0001)
-            self.linear2.weight.data.fill_(0.0001)
-            self.linear2.bias.data.fill_(0.0001)
+
+            self.linear1.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear1.bias.data.zero_()
+            self.linear2.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear2.bias.data.zero_()
 
         self.to(devices.device)
 
@@ -92,41 +93,54 @@ class Hypernetwork:
         self.sd_checkpoint_name = state_dict.get('sd_checkpoint_name', None)
 
 
-def load_hypernetworks(path):
+def list_hypernetworks(path):
     res = {}
+    for filename in glob.iglob(os.path.join(path, '**/*.pt'), recursive=True):
+        name = os.path.splitext(os.path.basename(filename))[0]
+        res[name] = filename
+    return res
 
-    for filename in glob.iglob(path + '**/*.pt', recursive=True):
+
+def load_hypernetwork(filename):
+    path = shared.hypernetworks.get(filename, None)
+    if path is not None:
+        print(f"Loading hypernetwork {filename}")
         try:
-            hn = Hypernetwork()
-            hn.load(filename)
-            res[hn.name] = hn
+            shared.loaded_hypernetwork = Hypernetwork()
+            shared.loaded_hypernetwork.load(path)
+
         except Exception:
-            print(f"Error loading hypernetwork {filename}", file=sys.stderr)
+            print(f"Error loading hypernetwork {path}", file=sys.stderr)
             print(traceback.format_exc(), file=sys.stderr)
+    else:
+        if shared.loaded_hypernetwork is not None:
+            print(f"Unloading hypernetwork")
 
-    return res
+        shared.loaded_hypernetwork = None
 
 
-def attention_CrossAttention_forward(self, x, context=None, mask=None):
-    h = self.heads
+def apply_hypernetwork(hypernetwork, context, layer=None):
+    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
 
-    q = self.to_q(x)
-    context = default(context, x)
+    if hypernetwork_layers is None:
+        return context, context
 
-    hypernetwork_layers = (shared.hypernetwork.layers if shared.hypernetwork is not None else {}).get(context.shape[2], None)
+    if layer is not None:
+        layer.hyper_k = hypernetwork_layers[0]
+        layer.hyper_v = hypernetwork_layers[1]
 
-    if hypernetwork_layers is not None:
-        hypernetwork_k, hypernetwork_v = hypernetwork_layers
+    context_k = hypernetwork_layers[0](context)
+    context_v = hypernetwork_layers[1](context)
+    return context_k, context_v
 
-        self.hypernetwork_k = hypernetwork_k
-        self.hypernetwork_v = hypernetwork_v
 
-        context_k = hypernetwork_k(context)
-        context_v = hypernetwork_v(context)
-    else:
-        context_k = context
-        context_v = context
+def attention_CrossAttention_forward(self, x, context=None, mask=None):
+    h = self.heads
+
+    q = self.to_q(x)
+    context = default(context, x)
 
+    context_k, context_v = apply_hypernetwork(shared.loaded_hypernetwork, context, self)
     k = self.to_k(context_k)
     v = self.to_v(context_v)
 
@@ -151,7 +165,9 @@ def attention_CrossAttention_forward(self, x, context=None, mask=None):
 def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory, steps, create_image_every, save_hypernetwork_every, template_file, preview_image_prompt):
     assert hypernetwork_name, 'embedding not selected'
 
-    shared.hypernetwork = shared.hypernetworks[hypernetwork_name]
+    path = shared.hypernetworks.get(hypernetwork_name, None)
+    shared.loaded_hypernetwork = Hypernetwork()
+    shared.loaded_hypernetwork.load(path)
 
     shared.state.textinfo = "Initializing hypernetwork training..."
     shared.state.job_count = steps
@@ -176,9 +192,9 @@ def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory,
 
     shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
     with torch.autocast("cuda"):
-        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, size=512, placeholder_token=hypernetwork_name, model=shared.sd_model, device=devices.device, template_file=template_file)
+        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=512, height=512, repeats=1, placeholder_token=hypernetwork_name, model=shared.sd_model, device=devices.device, template_file=template_file)
 
-    hypernetwork = shared.hypernetworks[hypernetwork_name]
+    hypernetwork = shared.loaded_hypernetwork
     weights = hypernetwork.weights()
     for weight in weights:
         weight.requires_grad = True
@@ -194,7 +210,7 @@ def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory,
     if ititial_step > steps:
         return hypernetwork, filename
 
-    pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step)
+    pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)
     for i, (x, text) in pbar:
         hypernetwork.step = i + ititial_step
 
diff --git a/modules/hypernetwork/ui.py b/modules/hypernetwork/ui.py
index 525f978c..f6d1d0a3 100644
--- a/modules/hypernetwork/ui.py
+++ b/modules/hypernetwork/ui.py
@@ -6,24 +6,24 @@ import gradio as gr
 import modules.textual_inversion.textual_inversion
 import modules.textual_inversion.preprocess
 from modules import sd_hijack, shared
+from modules.hypernetwork import hypernetwork
 
 
 def create_hypernetwork(name):
     fn = os.path.join(shared.cmd_opts.hypernetwork_dir, f"{name}.pt")
     assert not os.path.exists(fn), f"file {fn} already exists"
 
-    hypernetwork = modules.hypernetwork.hypernetwork.Hypernetwork(name=name)
-    hypernetwork.save(fn)
+    hypernet = modules.hypernetwork.hypernetwork.Hypernetwork(name=name)
+    hypernet.save(fn)
 
     shared.reload_hypernetworks()
-    shared.hypernetwork = shared.hypernetworks.get(shared.opts.sd_hypernetwork, None)
 
     return gr.Dropdown.update(choices=sorted([x for x in shared.hypernetworks.keys()])), f"Created: {fn}", ""
 
 
 def train_hypernetwork(*args):
 
-    initial_hypernetwork = shared.hypernetwork
+    initial_hypernetwork = shared.loaded_hypernetwork
 
     try:
         sd_hijack.undo_optimizations()
@@ -38,6 +38,6 @@ Hypernetwork saved to {html.escape(filename)}
     except Exception:
         raise
     finally:
-        shared.hypernetwork = initial_hypernetwork
+        shared.loaded_hypernetwork = initial_hypernetwork
         sd_hijack.apply_optimizations()
 
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:
diff --git a/modules/shared.py b/modules/shared.py
index 14b40d70..8753015e 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -13,7 +13,8 @@ import modules.memmon
 import modules.sd_models
 import modules.styles
 import modules.devices as devices
-from modules import sd_samplers, hypernetwork
+from modules import sd_samplers
+from modules.hypernetwork import hypernetwork
 from modules.paths import models_path, script_path, sd_path
 
 sd_model_file = os.path.join(script_path, 'model.ckpt')
@@ -29,6 +30,7 @@ parser.add_argument("--no-half-vae", action='store_true', help="do not switch th
 parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
 parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
 parser.add_argument("--embeddings-dir", type=str, default=os.path.join(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
+parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")
 parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
 parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage")
 parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage")
@@ -82,10 +84,17 @@ parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 xformers_available = False
 config_filename = cmd_opts.ui_settings_file
 
-hypernetworks = hypernetwork.list_hypernetworks(os.path.join(models_path, 'hypernetworks'))
+hypernetworks = hypernetwork.list_hypernetworks(cmd_opts.hypernetwork_dir)
 loaded_hypernetwork = None
 
 
+def reload_hypernetworks():
+    global hypernetworks
+
+    hypernetworks = hypernetwork.list_hypernetworks(cmd_opts.hypernetwork_dir)
+    hypernetwork.load_hypernetwork(opts.sd_hypernetwork)
+
+
 class State:
     skipped = False
     interrupted = False
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 5965c5a0..d6977950 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -156,7 +156,7 @@ def create_embedding(name, num_vectors_per_token, init_text='*'):
     return fn
 
 
-def train_embedding(embedding_name, learn_rate, data_root, log_directory, training_width, training_height, steps, num_repeats, create_image_every, save_embedding_every, template_file):
+def train_embedding(embedding_name, learn_rate, data_root, log_directory, training_width, training_height, steps, num_repeats, create_image_every, save_embedding_every, template_file, preview_image_prompt):
     assert embedding_name, 'embedding not selected'
 
     shared.state.textinfo = "Initializing textual inversion training..."
@@ -238,12 +238,14 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
         if embedding.step > 0 and images_dir is not None and embedding.step % create_image_every == 0:
             last_saved_image = os.path.join(images_dir, f'{embedding_name}-{embedding.step}.png')
 
+            preview_text = text if preview_image_prompt == "" else preview_image_prompt
+
             p = processing.StableDiffusionProcessingTxt2Img(
                 sd_model=shared.sd_model,
-                prompt=text,
+                prompt=preview_text,
                 steps=20,
-				height=training_height,
-				width=training_width,
+                height=training_height,
+                width=training_width,
                 do_not_save_grid=True,
                 do_not_save_samples=True,
             )
@@ -254,7 +256,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
             shared.state.current_image = image
             image.save(last_saved_image)
 
-            last_saved_image += f", prompt: {text}"
+            last_saved_image += f", prompt: {preview_text}"
 
         shared.state.job_no = embedding.step
 
diff --git a/modules/ui.py b/modules/ui.py
index 10b1ee3a..df653059 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1023,7 +1023,7 @@ def create_ui(wrap_gradio_gpu_call):
                             gr.HTML(value="")
 
                         with gr.Column():
-                            create_embedding = gr.Button(value="Create", variant='primary')
+                            create_embedding = gr.Button(value="Create embedding", variant='primary')
 
                 with gr.Group():
                     gr.HTML(value="<p style='margin-bottom: 0.7em'>Create a new hypernetwork</p>")
@@ -1035,7 +1035,7 @@ def create_ui(wrap_gradio_gpu_call):
                             gr.HTML(value="")
 
                         with gr.Column():
-                            create_hypernetwork = gr.Button(value="Create", variant='primary')
+                            create_hypernetwork = gr.Button(value="Create hypernetwork", variant='primary')
 
                 with gr.Group():
                     gr.HTML(value="<p style='margin-bottom: 0.7em'>Preprocess images</p>")
@@ -1147,6 +1147,7 @@ def create_ui(wrap_gradio_gpu_call):
                 create_image_every,
                 save_embedding_every,
                 template_file,
+                preview_image_prompt,
             ],
             outputs=[
                 ti_output,
diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index 42e1489c..0af5993c 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -10,7 +10,8 @@ import numpy as np
 import modules.scripts as scripts
 import gradio as gr
 
-from modules import images, hypernetwork
+from modules import images
+from modules.hypernetwork import hypernetwork
 from modules.processing import process_images, Processed, get_correct_sampler
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
diff --git a/webui.py b/webui.py
index 7c200551..ba2156c8 100644
--- a/webui.py
+++ b/webui.py
@@ -29,6 +29,7 @@ from modules import devices
 from modules import modelloader
 from modules.paths import script_path
 from modules.shared import cmd_opts
+import modules.hypernetwork.hypernetwork
 
 modelloader.cleanup_models()
 modules.sd_models.setup_model()
@@ -77,22 +78,12 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
     return modules.ui.wrap_gradio_call(f, extra_outputs=extra_outputs)
 
 
-def set_hypernetwork():
-    shared.hypernetwork = shared.hypernetworks.get(shared.opts.sd_hypernetwork, None)
-
-
-shared.reload_hypernetworks()
-shared.opts.onchange("sd_hypernetwork", set_hypernetwork)
-set_hypernetwork()
-
-
 modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
 
 shared.sd_model = modules.sd_models.load_model()
 shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))
 
-loaded_hypernetwork = modules.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)
-shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)))
+shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetwork.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)))
 
 
 def webui():
@@ -117,7 +108,7 @@ def webui():
             prevent_thread_lock=True
         )
         
-        app.add_middleware(GZipMiddleware,minimum_size=1000)
+        app.add_middleware(GZipMiddleware, minimum_size=1000)
 
         while 1:
             time.sleep(0.5)
-- 
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/hypernetwork/hypernetwork.py  | 283 ----------------------------------
 modules/hypernetwork/ui.py            |  43 ------
 modules/hypernetworks/hypernetwork.py | 283 ++++++++++++++++++++++++++++++++++
 modules/hypernetworks/ui.py           |  43 ++++++
 modules/sd_hijack.py                  |   2 +-
 modules/sd_hijack_optimizations.py    |   2 +-
 modules/shared.py                     |   2 +-
 modules/ui.py                         |   2 +-
 scripts/xy_grid.py                    |   2 +-
 webui.py                              |   2 +-
 10 files changed, 332 insertions(+), 332 deletions(-)
 delete mode 100644 modules/hypernetwork/hypernetwork.py
 delete mode 100644 modules/hypernetwork/ui.py
 create mode 100644 modules/hypernetworks/hypernetwork.py
 create mode 100644 modules/hypernetworks/ui.py

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/hypernetwork/hypernetwork.py b/modules/hypernetwork/hypernetwork.py
deleted file mode 100644
index aa701bda..00000000
--- a/modules/hypernetwork/hypernetwork.py
+++ /dev/null
@@ -1,283 +0,0 @@
-import datetime
-import glob
-import html
-import os
-import sys
-import traceback
-import tqdm
-
-import torch
-
-from ldm.util import default
-from modules import devices, shared, processing, sd_models
-import torch
-from torch import einsum
-from einops import rearrange, repeat
-import modules.textual_inversion.dataset
-
-
-class HypernetworkModule(torch.nn.Module):
-    def __init__(self, dim, state_dict=None):
-        super().__init__()
-
-        self.linear1 = torch.nn.Linear(dim, dim * 2)
-        self.linear2 = torch.nn.Linear(dim * 2, dim)
-
-        if state_dict is not None:
-            self.load_state_dict(state_dict, strict=True)
-        else:
-
-            self.linear1.weight.data.normal_(mean=0.0, std=0.01)
-            self.linear1.bias.data.zero_()
-            self.linear2.weight.data.normal_(mean=0.0, std=0.01)
-            self.linear2.bias.data.zero_()
-
-        self.to(devices.device)
-
-    def forward(self, x):
-        return x + (self.linear2(self.linear1(x)))
-
-
-class Hypernetwork:
-    filename = None
-    name = None
-
-    def __init__(self, name=None):
-        self.filename = None
-        self.name = name
-        self.layers = {}
-        self.step = 0
-        self.sd_checkpoint = None
-        self.sd_checkpoint_name = None
-
-        for size in [320, 640, 768, 1280]:
-            self.layers[size] = (HypernetworkModule(size), HypernetworkModule(size))
-
-    def weights(self):
-        res = []
-
-        for k, layers in self.layers.items():
-            for layer in layers:
-                layer.train()
-                res += [layer.linear1.weight, layer.linear1.bias, layer.linear2.weight, layer.linear2.bias]
-
-        return res
-
-    def save(self, filename):
-        state_dict = {}
-
-        for k, v in self.layers.items():
-            state_dict[k] = (v[0].state_dict(), v[1].state_dict())
-
-        state_dict['step'] = self.step
-        state_dict['name'] = self.name
-        state_dict['sd_checkpoint'] = self.sd_checkpoint
-        state_dict['sd_checkpoint_name'] = self.sd_checkpoint_name
-
-        torch.save(state_dict, filename)
-
-    def load(self, filename):
-        self.filename = filename
-        if self.name is None:
-            self.name = os.path.splitext(os.path.basename(filename))[0]
-
-        state_dict = torch.load(filename, map_location='cpu')
-
-        for size, sd in state_dict.items():
-            if type(size) == int:
-                self.layers[size] = (HypernetworkModule(size, sd[0]), HypernetworkModule(size, sd[1]))
-
-        self.name = state_dict.get('name', self.name)
-        self.step = state_dict.get('step', 0)
-        self.sd_checkpoint = state_dict.get('sd_checkpoint', None)
-        self.sd_checkpoint_name = state_dict.get('sd_checkpoint_name', None)
-
-
-def list_hypernetworks(path):
-    res = {}
-    for filename in glob.iglob(os.path.join(path, '**/*.pt'), recursive=True):
-        name = os.path.splitext(os.path.basename(filename))[0]
-        res[name] = filename
-    return res
-
-
-def load_hypernetwork(filename):
-    path = shared.hypernetworks.get(filename, None)
-    if path is not None:
-        print(f"Loading hypernetwork {filename}")
-        try:
-            shared.loaded_hypernetwork = Hypernetwork()
-            shared.loaded_hypernetwork.load(path)
-
-        except Exception:
-            print(f"Error loading hypernetwork {path}", file=sys.stderr)
-            print(traceback.format_exc(), file=sys.stderr)
-    else:
-        if shared.loaded_hypernetwork is not None:
-            print(f"Unloading hypernetwork")
-
-        shared.loaded_hypernetwork = None
-
-
-def apply_hypernetwork(hypernetwork, context, layer=None):
-    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
-
-    if hypernetwork_layers is None:
-        return context, context
-
-    if layer is not None:
-        layer.hyper_k = hypernetwork_layers[0]
-        layer.hyper_v = hypernetwork_layers[1]
-
-    context_k = hypernetwork_layers[0](context)
-    context_v = hypernetwork_layers[1](context)
-    return context_k, context_v
-
-
-def attention_CrossAttention_forward(self, x, context=None, mask=None):
-    h = self.heads
-
-    q = self.to_q(x)
-    context = default(context, x)
-
-    context_k, context_v = apply_hypernetwork(shared.loaded_hypernetwork, context, self)
-    k = self.to_k(context_k)
-    v = self.to_v(context_v)
-
-    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)
-
-
-def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory, steps, create_image_every, save_hypernetwork_every, template_file, preview_image_prompt):
-    assert hypernetwork_name, 'embedding not selected'
-
-    path = shared.hypernetworks.get(hypernetwork_name, None)
-    shared.loaded_hypernetwork = Hypernetwork()
-    shared.loaded_hypernetwork.load(path)
-
-    shared.state.textinfo = "Initializing hypernetwork training..."
-    shared.state.job_count = steps
-
-    filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt')
-
-    log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), hypernetwork_name)
-
-    if save_hypernetwork_every > 0:
-        hypernetwork_dir = os.path.join(log_directory, "hypernetworks")
-        os.makedirs(hypernetwork_dir, exist_ok=True)
-    else:
-        hypernetwork_dir = None
-
-    if create_image_every > 0:
-        images_dir = os.path.join(log_directory, "images")
-        os.makedirs(images_dir, exist_ok=True)
-    else:
-        images_dir = None
-
-    cond_model = shared.sd_model.cond_stage_model
-
-    shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
-    with torch.autocast("cuda"):
-        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=512, height=512, repeats=1, placeholder_token=hypernetwork_name, model=shared.sd_model, device=devices.device, template_file=template_file)
-
-    hypernetwork = shared.loaded_hypernetwork
-    weights = hypernetwork.weights()
-    for weight in weights:
-        weight.requires_grad = True
-
-    optimizer = torch.optim.AdamW(weights, lr=learn_rate)
-
-    losses = torch.zeros((32,))
-
-    last_saved_file = "<none>"
-    last_saved_image = "<none>"
-
-    ititial_step = hypernetwork.step or 0
-    if ititial_step > steps:
-        return hypernetwork, filename
-
-    pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)
-    for i, (x, text) in pbar:
-        hypernetwork.step = i + ititial_step
-
-        if hypernetwork.step > steps:
-            break
-
-        if shared.state.interrupted:
-            break
-
-        with torch.autocast("cuda"):
-            c = cond_model([text])
-
-            x = x.to(devices.device)
-            loss = shared.sd_model(x.unsqueeze(0), c)[0]
-            del x
-
-            losses[hypernetwork.step % losses.shape[0]] = loss.item()
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-        pbar.set_description(f"loss: {losses.mean():.7f}")
-
-        if hypernetwork.step > 0 and hypernetwork_dir is not None and hypernetwork.step % save_hypernetwork_every == 0:
-            last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork_name}-{hypernetwork.step}.pt')
-            hypernetwork.save(last_saved_file)
-
-        if hypernetwork.step > 0 and images_dir is not None and hypernetwork.step % create_image_every == 0:
-            last_saved_image = os.path.join(images_dir, f'{hypernetwork_name}-{hypernetwork.step}.png')
-
-            preview_text = text if preview_image_prompt == "" else preview_image_prompt
-
-            p = processing.StableDiffusionProcessingTxt2Img(
-                sd_model=shared.sd_model,
-                prompt=preview_text,
-                steps=20,
-                do_not_save_grid=True,
-                do_not_save_samples=True,
-            )
-
-            processed = processing.process_images(p)
-            image = processed.images[0]
-
-            shared.state.current_image = image
-            image.save(last_saved_image)
-
-            last_saved_image += f", prompt: {preview_text}"
-
-        shared.state.job_no = hypernetwork.step
-
-        shared.state.textinfo = f"""
-<p>
-Loss: {losses.mean():.7f}<br/>
-Step: {hypernetwork.step}<br/>
-Last prompt: {html.escape(text)}<br/>
-Last saved embedding: {html.escape(last_saved_file)}<br/>
-Last saved image: {html.escape(last_saved_image)}<br/>
-</p>
-"""
-
-    checkpoint = sd_models.select_checkpoint()
-
-    hypernetwork.sd_checkpoint = checkpoint.hash
-    hypernetwork.sd_checkpoint_name = checkpoint.model_name
-    hypernetwork.save(filename)
-
-    return hypernetwork, filename
-
-
diff --git a/modules/hypernetwork/ui.py b/modules/hypernetwork/ui.py
deleted file mode 100644
index f6d1d0a3..00000000
--- a/modules/hypernetwork/ui.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import html
-import os
-
-import gradio as gr
-
-import modules.textual_inversion.textual_inversion
-import modules.textual_inversion.preprocess
-from modules import sd_hijack, shared
-from modules.hypernetwork import hypernetwork
-
-
-def create_hypernetwork(name):
-    fn = os.path.join(shared.cmd_opts.hypernetwork_dir, f"{name}.pt")
-    assert not os.path.exists(fn), f"file {fn} already exists"
-
-    hypernet = modules.hypernetwork.hypernetwork.Hypernetwork(name=name)
-    hypernet.save(fn)
-
-    shared.reload_hypernetworks()
-
-    return gr.Dropdown.update(choices=sorted([x for x in shared.hypernetworks.keys()])), f"Created: {fn}", ""
-
-
-def train_hypernetwork(*args):
-
-    initial_hypernetwork = shared.loaded_hypernetwork
-
-    try:
-        sd_hijack.undo_optimizations()
-
-        hypernetwork, filename = modules.hypernetwork.hypernetwork.train_hypernetwork(*args)
-
-        res = f"""
-Training {'interrupted' if shared.state.interrupted else 'finished'} at {hypernetwork.step} steps.
-Hypernetwork saved to {html.escape(filename)}
-"""
-        return res, ""
-    except Exception:
-        raise
-    finally:
-        shared.loaded_hypernetwork = initial_hypernetwork
-        sd_hijack.apply_optimizations()
-
diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
new file mode 100644
index 00000000..aa701bda
--- /dev/null
+++ b/modules/hypernetworks/hypernetwork.py
@@ -0,0 +1,283 @@
+import datetime
+import glob
+import html
+import os
+import sys
+import traceback
+import tqdm
+
+import torch
+
+from ldm.util import default
+from modules import devices, shared, processing, sd_models
+import torch
+from torch import einsum
+from einops import rearrange, repeat
+import modules.textual_inversion.dataset
+
+
+class HypernetworkModule(torch.nn.Module):
+    def __init__(self, dim, state_dict=None):
+        super().__init__()
+
+        self.linear1 = torch.nn.Linear(dim, dim * 2)
+        self.linear2 = torch.nn.Linear(dim * 2, dim)
+
+        if state_dict is not None:
+            self.load_state_dict(state_dict, strict=True)
+        else:
+
+            self.linear1.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear1.bias.data.zero_()
+            self.linear2.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear2.bias.data.zero_()
+
+        self.to(devices.device)
+
+    def forward(self, x):
+        return x + (self.linear2(self.linear1(x)))
+
+
+class Hypernetwork:
+    filename = None
+    name = None
+
+    def __init__(self, name=None):
+        self.filename = None
+        self.name = name
+        self.layers = {}
+        self.step = 0
+        self.sd_checkpoint = None
+        self.sd_checkpoint_name = None
+
+        for size in [320, 640, 768, 1280]:
+            self.layers[size] = (HypernetworkModule(size), HypernetworkModule(size))
+
+    def weights(self):
+        res = []
+
+        for k, layers in self.layers.items():
+            for layer in layers:
+                layer.train()
+                res += [layer.linear1.weight, layer.linear1.bias, layer.linear2.weight, layer.linear2.bias]
+
+        return res
+
+    def save(self, filename):
+        state_dict = {}
+
+        for k, v in self.layers.items():
+            state_dict[k] = (v[0].state_dict(), v[1].state_dict())
+
+        state_dict['step'] = self.step
+        state_dict['name'] = self.name
+        state_dict['sd_checkpoint'] = self.sd_checkpoint
+        state_dict['sd_checkpoint_name'] = self.sd_checkpoint_name
+
+        torch.save(state_dict, filename)
+
+    def load(self, filename):
+        self.filename = filename
+        if self.name is None:
+            self.name = os.path.splitext(os.path.basename(filename))[0]
+
+        state_dict = torch.load(filename, map_location='cpu')
+
+        for size, sd in state_dict.items():
+            if type(size) == int:
+                self.layers[size] = (HypernetworkModule(size, sd[0]), HypernetworkModule(size, sd[1]))
+
+        self.name = state_dict.get('name', self.name)
+        self.step = state_dict.get('step', 0)
+        self.sd_checkpoint = state_dict.get('sd_checkpoint', None)
+        self.sd_checkpoint_name = state_dict.get('sd_checkpoint_name', None)
+
+
+def list_hypernetworks(path):
+    res = {}
+    for filename in glob.iglob(os.path.join(path, '**/*.pt'), recursive=True):
+        name = os.path.splitext(os.path.basename(filename))[0]
+        res[name] = filename
+    return res
+
+
+def load_hypernetwork(filename):
+    path = shared.hypernetworks.get(filename, None)
+    if path is not None:
+        print(f"Loading hypernetwork {filename}")
+        try:
+            shared.loaded_hypernetwork = Hypernetwork()
+            shared.loaded_hypernetwork.load(path)
+
+        except Exception:
+            print(f"Error loading hypernetwork {path}", file=sys.stderr)
+            print(traceback.format_exc(), file=sys.stderr)
+    else:
+        if shared.loaded_hypernetwork is not None:
+            print(f"Unloading hypernetwork")
+
+        shared.loaded_hypernetwork = None
+
+
+def apply_hypernetwork(hypernetwork, context, layer=None):
+    hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context.shape[2], None)
+
+    if hypernetwork_layers is None:
+        return context, context
+
+    if layer is not None:
+        layer.hyper_k = hypernetwork_layers[0]
+        layer.hyper_v = hypernetwork_layers[1]
+
+    context_k = hypernetwork_layers[0](context)
+    context_v = hypernetwork_layers[1](context)
+    return context_k, context_v
+
+
+def attention_CrossAttention_forward(self, x, context=None, mask=None):
+    h = self.heads
+
+    q = self.to_q(x)
+    context = default(context, x)
+
+    context_k, context_v = apply_hypernetwork(shared.loaded_hypernetwork, context, self)
+    k = self.to_k(context_k)
+    v = self.to_v(context_v)
+
+    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)
+
+
+def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory, steps, create_image_every, save_hypernetwork_every, template_file, preview_image_prompt):
+    assert hypernetwork_name, 'embedding not selected'
+
+    path = shared.hypernetworks.get(hypernetwork_name, None)
+    shared.loaded_hypernetwork = Hypernetwork()
+    shared.loaded_hypernetwork.load(path)
+
+    shared.state.textinfo = "Initializing hypernetwork training..."
+    shared.state.job_count = steps
+
+    filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt')
+
+    log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), hypernetwork_name)
+
+    if save_hypernetwork_every > 0:
+        hypernetwork_dir = os.path.join(log_directory, "hypernetworks")
+        os.makedirs(hypernetwork_dir, exist_ok=True)
+    else:
+        hypernetwork_dir = None
+
+    if create_image_every > 0:
+        images_dir = os.path.join(log_directory, "images")
+        os.makedirs(images_dir, exist_ok=True)
+    else:
+        images_dir = None
+
+    cond_model = shared.sd_model.cond_stage_model
+
+    shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
+    with torch.autocast("cuda"):
+        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=512, height=512, repeats=1, placeholder_token=hypernetwork_name, model=shared.sd_model, device=devices.device, template_file=template_file)
+
+    hypernetwork = shared.loaded_hypernetwork
+    weights = hypernetwork.weights()
+    for weight in weights:
+        weight.requires_grad = True
+
+    optimizer = torch.optim.AdamW(weights, lr=learn_rate)
+
+    losses = torch.zeros((32,))
+
+    last_saved_file = "<none>"
+    last_saved_image = "<none>"
+
+    ititial_step = hypernetwork.step or 0
+    if ititial_step > steps:
+        return hypernetwork, filename
+
+    pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)
+    for i, (x, text) in pbar:
+        hypernetwork.step = i + ititial_step
+
+        if hypernetwork.step > steps:
+            break
+
+        if shared.state.interrupted:
+            break
+
+        with torch.autocast("cuda"):
+            c = cond_model([text])
+
+            x = x.to(devices.device)
+            loss = shared.sd_model(x.unsqueeze(0), c)[0]
+            del x
+
+            losses[hypernetwork.step % losses.shape[0]] = loss.item()
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        pbar.set_description(f"loss: {losses.mean():.7f}")
+
+        if hypernetwork.step > 0 and hypernetwork_dir is not None and hypernetwork.step % save_hypernetwork_every == 0:
+            last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork_name}-{hypernetwork.step}.pt')
+            hypernetwork.save(last_saved_file)
+
+        if hypernetwork.step > 0 and images_dir is not None and hypernetwork.step % create_image_every == 0:
+            last_saved_image = os.path.join(images_dir, f'{hypernetwork_name}-{hypernetwork.step}.png')
+
+            preview_text = text if preview_image_prompt == "" else preview_image_prompt
+
+            p = processing.StableDiffusionProcessingTxt2Img(
+                sd_model=shared.sd_model,
+                prompt=preview_text,
+                steps=20,
+                do_not_save_grid=True,
+                do_not_save_samples=True,
+            )
+
+            processed = processing.process_images(p)
+            image = processed.images[0]
+
+            shared.state.current_image = image
+            image.save(last_saved_image)
+
+            last_saved_image += f", prompt: {preview_text}"
+
+        shared.state.job_no = hypernetwork.step
+
+        shared.state.textinfo = f"""
+<p>
+Loss: {losses.mean():.7f}<br/>
+Step: {hypernetwork.step}<br/>
+Last prompt: {html.escape(text)}<br/>
+Last saved embedding: {html.escape(last_saved_file)}<br/>
+Last saved image: {html.escape(last_saved_image)}<br/>
+</p>
+"""
+
+    checkpoint = sd_models.select_checkpoint()
+
+    hypernetwork.sd_checkpoint = checkpoint.hash
+    hypernetwork.sd_checkpoint_name = checkpoint.model_name
+    hypernetwork.save(filename)
+
+    return hypernetwork, filename
+
+
diff --git a/modules/hypernetworks/ui.py b/modules/hypernetworks/ui.py
new file mode 100644
index 00000000..811bc31e
--- /dev/null
+++ b/modules/hypernetworks/ui.py
@@ -0,0 +1,43 @@
+import html
+import os
+
+import gradio as gr
+
+import modules.textual_inversion.textual_inversion
+import modules.textual_inversion.preprocess
+from modules import sd_hijack, shared
+from modules.hypernetworks import hypernetwork
+
+
+def create_hypernetwork(name):
+    fn = os.path.join(shared.cmd_opts.hypernetwork_dir, f"{name}.pt")
+    assert not os.path.exists(fn), f"file {fn} already exists"
+
+    hypernet = modules.hypernetwork.hypernetwork.Hypernetwork(name=name)
+    hypernet.save(fn)
+
+    shared.reload_hypernetworks()
+
+    return gr.Dropdown.update(choices=sorted([x for x in shared.hypernetworks.keys()])), f"Created: {fn}", ""
+
+
+def train_hypernetwork(*args):
+
+    initial_hypernetwork = shared.loaded_hypernetwork
+
+    try:
+        sd_hijack.undo_optimizations()
+
+        hypernetwork, filename = modules.hypernetwork.hypernetwork.train_hypernetwork(*args)
+
+        res = f"""
+Training {'interrupted' if shared.state.interrupted else 'finished'} at {hypernetwork.step} steps.
+Hypernetwork saved to {html.escape(filename)}
+"""
+        return res, ""
+    except Exception:
+        raise
+    finally:
+        shared.loaded_hypernetwork = initial_hypernetwork
+        sd_hijack.apply_optimizations()
+
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index f873049a..f07ec041 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -37,7 +37,7 @@ def apply_optimizations():
 
 
 def undo_optimizations():
-    from modules.hypernetwork import hypernetwork
+    from modules.hypernetworks import hypernetwork
 
     ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
     ldm.modules.diffusionmodules.model.nonlinearity = diffusionmodules_model_nonlinearity
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:
diff --git a/modules/shared.py b/modules/shared.py
index 375e3afb..1dc2ccf2 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -14,7 +14,7 @@ import modules.sd_models
 import modules.styles
 import modules.devices as devices
 from modules import sd_samplers
-from modules.hypernetwork import hypernetwork
+from modules.hypernetworks import hypernetwork
 from modules.paths import models_path, script_path, sd_path
 
 sd_model_file = os.path.join(script_path, 'model.ckpt')
diff --git a/modules/ui.py b/modules/ui.py
index f57f32db..42e5d866 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -39,7 +39,7 @@ import modules.generation_parameters_copypaste
 from modules import prompt_parser
 from modules.images import save_image
 import modules.textual_inversion.ui
-import modules.hypernetwork.ui
+import modules.hypernetworks.ui
 
 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the browser will not show any UI
 mimetypes.init()
diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index 16918c99..cddb192a 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -11,7 +11,7 @@ import modules.scripts as scripts
 import gradio as gr
 
 from modules import images
-from modules.hypernetwork import hypernetwork
+from modules.hypernetworks import hypernetwork
 from modules.processing import process_images, Processed, get_correct_sampler
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
diff --git a/webui.py b/webui.py
index ba2156c8..faa38a0d 100644
--- a/webui.py
+++ b/webui.py
@@ -29,7 +29,7 @@ from modules import devices
 from modules import modelloader
 from modules.paths import script_path
 from modules.shared import cmd_opts
-import modules.hypernetwork.hypernetwork
+import modules.hypernetworks.hypernetwork
 
 modelloader.cleanup_models()
 modules.sd_models.setup_model()
-- 
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.py               |  5 ++-
 modules/sd_hijack_optimizations.py | 79 ++++++++++++++++++++++++++++++++++++++
 modules/shared.py                  |  5 ++-
 3 files changed, 86 insertions(+), 3 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index f07ec041..5a1b167f 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -30,8 +30,11 @@ def apply_optimizations():
     elif cmd_opts.opt_split_attention_v1:
         print("Applying v1 cross attention optimization.")
         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1
+    elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention_invokeai or not torch.cuda.is_available()):
+        print("Applying cross attention optimization (InvokeAI).")
+        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_invokeAI
     elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention or torch.cuda.is_available()):
-        print("Applying cross attention optimization.")
+        print("Applying cross attention optimization (Doggettx).")
         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward
         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.cross_attention_attnblock_forward
 
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)
diff --git a/modules/shared.py b/modules/shared.py
index 1dc2ccf2..20b45f23 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -50,9 +50,10 @@ parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with
 parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")
 parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")
 parser.add_argument("--deepdanbooru", action='store_true', help="enable deepdanbooru interrogator")
-parser.add_argument("--opt-split-attention", action='store_true', help="force-enables cross-attention layer optimization. By default, it's on for torch.cuda and off for other torch devices.")
-parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
+parser.add_argument("--opt-split-attention", action='store_true', help="force-enables Doggettx's cross-attention layer optimization. By default, it's on for torch cuda.")
+parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")
 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
+parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
 parser.add_argument("--use-cpu", nargs='+',choices=['SD', 'GFPGAN', 'BSRGAN', 'ESRGAN', 'SCUNet', 'CodeFormer'], help="use CPU as torch device for specified modules", default=[])
 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
 parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
-- 
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.py               | 10 ++++++++--
 modules/sd_hijack_optimizations.py | 19 +++++++++++++++----
 2 files changed, 23 insertions(+), 6 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 5a1b167f..ac70f876 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -10,6 +10,7 @@ from torch.nn.functional import silu
 import modules.textual_inversion.textual_inversion
 from modules import prompt_parser, devices, sd_hijack_optimizations, shared
 from modules.shared import opts, device, cmd_opts
+from modules.sd_hijack_optimizations import invokeAI_mps_available
 
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
@@ -31,8 +32,13 @@ def apply_optimizations():
         print("Applying v1 cross attention optimization.")
         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1
     elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention_invokeai or not torch.cuda.is_available()):
-        print("Applying cross attention optimization (InvokeAI).")
-        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_invokeAI
+        if not invokeAI_mps_available and shared.device.type == 'mps':
+            print("The InvokeAI cross attention optimization for MPS requires the psutil package which is not installed.")
+            print("Applying v1 cross attention optimization.")
+            ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1
+        else:
+            print("Applying cross attention optimization (InvokeAI).")
+            ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_invokeAI
     elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention or torch.cuda.is_available()):
         print("Applying cross attention optimization (Doggettx).")
         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward
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

---
 README.md                          |  1 +
 modules/sd_hijack_optimizations.py | 13 +++++++++++++
 2 files changed, 14 insertions(+)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/README.md b/README.md
index a10faa01..859a91b6 100644
--- a/README.md
+++ b/README.md
@@ -123,6 +123,7 @@ The documentation was moved from this README over to the project's [wiki](https:
 - LDSR - https://github.com/Hafiidz/latent-diffusion
 - Ideas for optimizations - https://github.com/basujindal/stable-diffusion
 - Doggettx - Cross Attention layer optimization - https://github.com/Doggettx/stable-diffusion, original idea for prompt editing.
+- InvokeAI, lstein - Cross Attention layer optimization - https://github.com/invoke-ai/InvokeAI (originally http://github.com/lstein/stable-diffusion)
 - Rinon Gal - Textual Inversion - https://github.com/rinongal/textual_inversion (we're not using his code, but we are using his ideas).
 - Idea for SD upscale - https://github.com/jquesnelle/txt2imghd
 - Noise generation for outpainting mk2 - https://github.com/parlance-zz/g-diffuser-bot
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.py               |  2 ++
 modules/sd_hijack_optimizations.py | 31 +++++++++++++++++++++++++++++++
 modules/ui.py                      | 24 ++++++++++++------------
 3 files changed, 45 insertions(+), 12 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index eaedac13..6141f705 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -16,6 +16,7 @@ import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
 import ldm.models.diffusion.ddim
 import ldm.models.diffusion.plms
+import ldm.modules.diffusionmodules.openaimodel
 
 attention_CrossAttention_forward = ldm.modules.attention.CrossAttention.forward
 diffusionmodules_model_nonlinearity = ldm.modules.diffusionmodules.model.nonlinearity
@@ -26,6 +27,7 @@ def apply_optimizations():
     undo_optimizations()
 
     ldm.modules.diffusionmodules.model.nonlinearity = silu
+    ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_hijack_optimizations.patched_unet_forward
 
     if cmd_opts.force_enable_xformers or (cmd_opts.xformers and shared.xformers_available and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (9, 0)):
         print("Applying xformers cross attention optimization.")
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)
diff --git a/modules/ui.py b/modules/ui.py
index e6da1b2a..85e531af 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -380,8 +380,8 @@ def create_seed_inputs():
 
     with gr.Row(visible=False) as seed_extra_row_2:
         seed_extras.append(seed_extra_row_2)
-        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from width", value=0)
-        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from height", value=0)
+        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=1, label="Resize seed from width", value=0)
+        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=1, label="Resize seed from height", value=0)
 
     random_seed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[seed])
     random_subseed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[subseed])
@@ -715,8 +715,8 @@ def create_ui(wrap_gradio_gpu_call):
                 sampler_index = gr.Radio(label='Sampling method', elem_id="txt2img_sampling", choices=[x.name for x in samplers], value=samplers[0].name, type="index")
 
                 with gr.Group():
-                    width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)
-                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
+                    width = gr.Slider(minimum=64, maximum=2048, step=1, label="Width", value=512)
+                    height = gr.Slider(minimum=64, maximum=2048, step=1, label="Height", value=512)
 
                 with gr.Row():
                     restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)
@@ -724,8 +724,8 @@ def create_ui(wrap_gradio_gpu_call):
                     enable_hr = gr.Checkbox(label='Highres. fix', value=False)
 
                 with gr.Row(visible=False) as hr_options:
-                    firstphase_width = gr.Slider(minimum=0, maximum=1024, step=64, label="Firstpass width", value=0)
-                    firstphase_height = gr.Slider(minimum=0, maximum=1024, step=64, label="Firstpass height", value=0)
+                    firstphase_width = gr.Slider(minimum=0, maximum=1024, step=1, label="Firstpass width", value=0)
+                    firstphase_height = gr.Slider(minimum=0, maximum=1024, step=1, label="Firstpass height", value=0)
                     denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.7)
 
                 with gr.Row(equal_height=True):
@@ -901,8 +901,8 @@ def create_ui(wrap_gradio_gpu_call):
                 sampler_index = gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index")
 
                 with gr.Group():
-                    width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512, elem_id="img2img_width")
-                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512, elem_id="img2img_height")
+                    width = gr.Slider(minimum=64, maximum=2048, step=1, label="Width", value=512, elem_id="img2img_width")
+                    height = gr.Slider(minimum=64, maximum=2048, step=1, label="Height", value=512, elem_id="img2img_height")
 
                 with gr.Row():
                     restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)
@@ -1231,8 +1231,8 @@ def create_ui(wrap_gradio_gpu_call):
                 with gr.Tab(label="Preprocess images"):
                     process_src = gr.Textbox(label='Source directory')
                     process_dst = gr.Textbox(label='Destination directory')
-                    process_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)
-                    process_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
+                    process_width = gr.Slider(minimum=64, maximum=2048, step=1, label="Width", value=512)
+                    process_height = gr.Slider(minimum=64, maximum=2048, step=1, label="Height", value=512)
                     preprocess_txt_action = gr.Dropdown(label='Existing Caption txt Action', value="ignore", choices=["ignore", "copy", "prepend", "append"])
 
                     with gr.Row():
@@ -1289,8 +1289,8 @@ def create_ui(wrap_gradio_gpu_call):
                     dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
                     log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion")
                     template_file = gr.Textbox(label='Prompt template file', value=os.path.join(script_path, "textual_inversion_templates", "style_filewords.txt"))
-                    training_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)
-                    training_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
+                    training_width = gr.Slider(minimum=64, maximum=2048, step=1, label="Width", value=512)
+                    training_height = gr.Slider(minimum=64, maximum=2048, step=1, label="Height", value=512)
                     steps = gr.Number(label='Max steps', value=100000, precision=0)
                     create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0)
                     save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0)
-- 
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.py               |  5 +++--
 modules/sd_hijack_optimizations.py | 28 ----------------------------
 modules/sd_hijack_unet.py          | 30 ++++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 30 deletions(-)
 create mode 100644 modules/sd_hijack_unet.py

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 92874a79..47dbc1b7 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -11,7 +11,7 @@ import modules.textual_inversion.textual_inversion
 from modules import prompt_parser, devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint
 from modules.hypernetworks import hypernetwork
 from modules.shared import opts, device, cmd_opts
-from modules import sd_hijack_clip, sd_hijack_open_clip
+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet
 
 from modules.sd_hijack_optimizations import invokeAI_mps_available
 
@@ -35,11 +35,12 @@ ldm.modules.attention.BasicTransformerBlock.ATTENTION_MODES["softmax-xformers"]
 ldm.modules.attention.print = lambda *args: None
 ldm.modules.diffusionmodules.model.print = lambda *args: None
 
+
 def apply_optimizations():
     undo_optimizations()
 
     ldm.modules.diffusionmodules.model.nonlinearity = silu
-    ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_hijack_optimizations.patched_unet_forward
+    ldm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th
 
     if cmd_opts.force_enable_xformers or (cmd_opts.xformers and shared.xformers_available and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (9, 0)):
         print("Applying xformers cross attention optimization.")
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)
diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py
new file mode 100644
index 00000000..1b9d7757
--- /dev/null
+++ b/modules/sd_hijack_unet.py
@@ -0,0 +1,30 @@
+import torch
+
+
+class TorchHijackForUnet:
+    """
+    This is torch, but with cat that resizes tensors to appropriate dimensions if they do not match;
+    this makes it possible to create pictures with dimensions that are muliples of 8 rather than 64
+    """
+
+    def __getattr__(self, item):
+        if item == 'cat':
+            return self.cat
+
+        if hasattr(torch, item):
+            return getattr(torch, item)
+
+        raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, item))
+
+    def cat(self, tensors, *args, **kwargs):
+        if len(tensors) == 2:
+            a, b = tensors
+            if a.shape[-2:] != b.shape[-2:]:
+                a = torch.nn.functional.interpolate(a, b.shape[-2:], mode="nearest")
+
+            tensors = (a, b)
+
+        return torch.cat(tensors, *args, **kwargs)
+
+
+th = TorchHijackForUnet()
-- 
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.py               | 10 ++--------
 modules/sd_hijack_optimizations.py |  3 ---
 2 files changed, 2 insertions(+), 11 deletions(-)

(limited to 'modules/sd_hijack_optimizations.py')

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 47dbc1b7..690a9ec2 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -1,16 +1,10 @@
-import math
-import os
-import sys
-import traceback
 import torch
-import numpy as np
-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, sd_hijack_checkpoint
+from modules import devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint
 from modules.hypernetworks import hypernetwork
-from modules.shared import opts, device, cmd_opts
+from modules.shared import cmd_opts
 from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet
 
 from modules.sd_hijack_optimizations import invokeAI_mps_available
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