cherry-pick @Any-Winter-4079's invoke-ai/InvokeAI#540. this is a collaboration incorporating a lot of people's contributions -- including for example @Doggettx and the original code from @neonsecret on which the Doggetx optimizations were based (see invoke-ai/InvokeAI#431, https://github.com/sd-webui/stable-diffusion-webui/pull/771\#issuecomment-1239716055). Takes exactly the same amount of time to run 8 steps as original CompVis code does (10.4 secs, ~1.25s/it).

Author: Birch-san

ldm/modules/attention.py

@@ -7,6 +7,8 @@
 
 from ldm.modules.diffusionmodules.util import checkpoint
 
+import psutil
+
 
 def exists(val):
     return val is not None
@@ -167,6 +169,80 @@ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.
             nn.Dropout(dropout)
         )
 
+        if torch.cuda.is_available():
+            self.einsum_op = self.einsum_op_cuda
+        else:
+            self.mem_total = psutil.virtual_memory().total / (1024**3)
+            self.einsum_op = self.einsum_op_mps_v1 if self.mem_total >= 32 else self.einsum_op_mps_v2
+    
+    def einsum_op_compvis(self, q, k, v, r1):
+        s1 = einsum('b i d, b j d -> b i j', q, k) * self.scale # faster
+        s2 = s1.softmax(dim=-1, dtype=q.dtype)
+        del s1
+        r1 = einsum('b i j, b j d -> b i d', s2, v)
+        del s2
+        return r1
+
+    def einsum_op_mps_v1(self, q, k, v, r1):
+        if q.shape[1] <= 4096: # (512x512) max q.shape[1]: 4096
+            r1 = self.einsum_op_compvis(q, k, v, r1)
+        else:
+            slice_size = math.floor(2**30 / (q.shape[0] * q.shape[1]))
+            for i in range(0, q.shape[1], slice_size):
+                end = i + slice_size
+                s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
+                s2 = s1.softmax(dim=-1, dtype=r1.dtype)
+                del s1  
+                r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
+                del s2
+        return r1
+
+    def einsum_op_mps_v2(self, q, k, v, r1):
+        if self.mem_total >= 8 and q.shape[1] <= 4096:
+                r1 = self.einsum_op_compvis(q, k, v, r1)
+        else:
+            slice_size = 1
+            for i in range(0, q.shape[0], slice_size):
+                end = min(q.shape[0], i + slice_size)
+                s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])
+                s1 *= self.scale
+                s2 = s1.softmax(dim=-1, dtype=r1.dtype)
+                del s1
+                r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])
+                del s2
+        return r1
+
+    def einsum_op_cuda(self, q, k, v, r1):
+        stats = torch.cuda.memory_stats(q.device)
+        mem_active = stats['active_bytes.all.current']
+        mem_reserved = stats['reserved_bytes.all.current']
+        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
+        mem_free_torch = mem_reserved - mem_active
+        mem_free_total = mem_free_cuda + mem_free_torch
+
+        gb = 1024 ** 3
+        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * 4
+        mem_required = tensor_size * 2.5
+        steps = 1
+
+        if mem_required > mem_free_total:
+            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
+
+        if steps > 64:
+            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
+            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
+                            f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')
+
+        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]  
+        for i in range(0, q.shape[1], slice_size):
+            end = min(q.shape[1], i + slice_size)
+            s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
+            s2 = s1.softmax(dim=-1, dtype=r1.dtype)
+            del s1
+            r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
+            del s2 
+        return r1
+
     def forward(self, x, context=None, mask=None):
         h = self.heads
 
@@ -179,25 +255,12 @@ def forward(self, x, context=None, mask=None):
 
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
 
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
-        del q, k
-
-        if exists(mask):
-            mask = rearrange(mask, 'b ... -> b (...)')
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
-            sim.masked_fill_(~mask, max_neg_value)
-            del mask
-
-        # attention, what we cannot get enough of
-        attn = sim.softmax(dim=-1)
-        del sim
-
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        del attn, v
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
-        del h
-        return self.to_out(out)
+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
+        r1 = self.einsum_op(q, k, v, r1)
+        del q, k, v
+        r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
+        del r1
+        return self.to_out(r2)
 
 
 class BasicTransformerBlock(nn.Module):