michael/stable-diffusion-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

add xformers attnblock and hypernetwork support

Browse source
This commit is contained in:
C43H66N12O12S2 2022-10-08 11:55:02 +03:00 committed by GitHub
parent b70eaeb200
commit 5d54f35c58
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
modules/sd_hijack_optimizations.py
View file

@ -98,8 +98,14 @@ def xformers_attention_forward(self, x, context=None, mask=None):
h = self.heads h = self.heads
q_in = self.to_q(x) q_in = self.to_q(x)
context = default(context, x) context = default(context, x)
k_in = self.to_k(context) hypernetwork = shared.selected_hypernetwork()
v_in = self.to_v(context) 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)) 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 del q_in, k_in, v_in
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None) 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 h3 += x
return h3 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