diff options
| author | MalumaDev <piano.lu92@gmail.com> | 2022-10-18 17:27:30 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-10-18 17:27:30 +0200 |
| commit | c2765c9bcd264f5a8922348cd03521cb2ff306b3 (patch) | |
| tree | 475fb1048477fac8f0b9afc741b6a8eefeb339d7 /modules/sd_hijack_optimizations.py | |
| parent | 1997ccff13fc75af223f571e8c927c3d77273dd9 (diff) | |
| parent | c1093b8051606f0ac90506b7114c4b55d0447c70 (diff) | |
Merge branch 'master' into test_resolve_conflicts
Diffstat (limited to 'modules/sd_hijack_optimizations.py')
| -rw-r--r-- | modules/sd_hijack_optimizations.py | 14 |
1 files changed, 10 insertions, 4 deletions
diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 79405525..a3345bb9 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -296,10 +296,16 @@ 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)
+ 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))
+ 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)
return x + out
except NotImplementedError:
|