| 70 | |
| 71 | # Stolen from https://github.com/basujindal/stable-diffusion/blob/main/optimizedSD/splitAttention.py |
| 72 | def _hacked_sliced_attentin_forward(self, x, context=None, mask=None): |
| 73 | h = self.heads |
| 74 | |
| 75 | q = self.to_q(x) |
| 76 | context = default(context, x) |
| 77 | k = self.to_k(context) |
| 78 | v = self.to_v(context) |
| 79 | del context, x |
| 80 | |
| 81 | q, k, v = map(lambda t: einops.rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) |
| 82 | |
| 83 | limit = k.shape[0] |
| 84 | att_step = 1 |
| 85 | q_chunks = list(torch.tensor_split(q, limit // att_step, dim=0)) |
| 86 | k_chunks = list(torch.tensor_split(k, limit // att_step, dim=0)) |
| 87 | v_chunks = list(torch.tensor_split(v, limit // att_step, dim=0)) |
| 88 | |
| 89 | q_chunks.reverse() |
| 90 | k_chunks.reverse() |
| 91 | v_chunks.reverse() |
| 92 | sim = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device) |
| 93 | del k, q, v |
| 94 | for i in range(0, limit, att_step): |
| 95 | q_buffer = q_chunks.pop() |
| 96 | k_buffer = k_chunks.pop() |
| 97 | v_buffer = v_chunks.pop() |
| 98 | sim_buffer = torch.einsum('b i d, b j d -> b i j', q_buffer, k_buffer) * self.scale |
| 99 | |
| 100 | del k_buffer, q_buffer |
| 101 | # attention, what we cannot get enough of, by chunks |
| 102 | |
| 103 | sim_buffer = sim_buffer.softmax(dim=-1) |
| 104 | |
| 105 | sim_buffer = torch.einsum('b i j, b j d -> b i d', sim_buffer, v_buffer) |
| 106 | del v_buffer |
| 107 | sim[i:i + att_step, :, :] = sim_buffer |
| 108 | |
| 109 | del sim_buffer |
| 110 | sim = einops.rearrange(sim, '(b h) n d -> b n (h d)', h=h) |
| 111 | return self.to_out(sim) |