diff options
author | C43H66N12O12S2 <36072735+C43H66N12O12S2@users.noreply.github.com> | 2022-10-08 08:55:02 +0000 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-10-08 08:55:02 +0000 |
commit | 5d54f35c583bd5a3b0ee271a862827f1ca81ef09 (patch) | |
tree | 2c0cfa8f120ecb71fa7781eb4f45ee2abd01d3bb | |
parent | b70eaeb2005a5a9593119e7fd32b8072c2a208d5 (diff) | |
download | stable-diffusion-webui-gfx803-5d54f35c583bd5a3b0ee271a862827f1ca81ef09.tar.gz stable-diffusion-webui-gfx803-5d54f35c583bd5a3b0ee271a862827f1ca81ef09.tar.bz2 stable-diffusion-webui-gfx803-5d54f35c583bd5a3b0ee271a862827f1ca81ef09.zip |
add xformers attnblock and hypernetwork support
-rw-r--r-- | modules/sd_hijack_optimizations.py | 20 |
1 files changed, 18 insertions, 2 deletions
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
|