diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2022-10-18 11:04:17 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-10-18 11:04:17 +0000 |
| commit | f6c758d055ca8e181d5bfc056cb85ff5a2df7f66 (patch) | |
| tree | dd1a7276d6d43f950227e54dbae28fbf6801d35b /modules/sd_hijack_optimizations.py | |
| parent | 665beebc0825a6fad410c8252f27f6f6f0bd900b (diff) | |
| parent | c71008c74156635558bb2e877d1628913f6f781e (diff) | |
| download | stable-diffusion-webui-gfx803-f6c758d055ca8e181d5bfc056cb85ff5a2df7f66.tar.gz stable-diffusion-webui-gfx803-f6c758d055ca8e181d5bfc056cb85ff5a2df7f66.tar.bz2 stable-diffusion-webui-gfx803-f6c758d055ca8e181d5bfc056cb85ff5a2df7f66.zip | |
Merge branch 'master' into dark-mode
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:
|