Merge branch 'dev' of https://github.com/AUTOMATIC1111/stable-diffusion-webui into tomesd

1 files changed, 26 insertions, 24 deletions

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 372555ff..f00fe55c 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -49,7 +49,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     v_in = self.to_v(context_v)

     del context, context_k, context_v, x

 

-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))

+    q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q_in, k_in, v_in))

     del q_in, k_in, v_in

 

     dtype = q.dtype

@@ -62,10 +62,10 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
             end = i + 2

             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)

             del s1

-    

+

             r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])

             del s2

         del q, k, v

@@ -95,43 +95,43 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
 

     with devices.without_autocast(disable=not shared.opts.upcast_attn):

         k_in = k_in * self.scale

-    

+

         del context, x

-    

-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))

+

+        q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q_in, k_in, v_in))

         del q_in, k_in, v_in

-    

+

         r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)

-    

+

         mem_free_total = get_available_vram()

-    

+

         gb = 1024 ** 3

         tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()

         modifier = 3 if q.element_size() == 2 else 2.5

         mem_required = tensor_size * modifier

         steps = 1

-    

+

         if mem_required > mem_free_total:

             steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))

             # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "

             #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")

-    

+

         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 = i + slice_size

             s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)

-    

+

             s2 = s1.softmax(dim=-1, dtype=q.dtype)

             del s1

-    

+

             r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)

             del s2

-    

+

         del q, k, v

 

     r1 = r1.to(dtype)

@@ -228,8 +228,8 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
 

     with devices.without_autocast(disable=not shared.opts.upcast_attn):

         k = k * self.scale

-    

-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

+

+        q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q, k, v))

         r = einsum_op(q, k, v)

     r = r.to(dtype)

     return self.to_out(rearrange(r, '(b h) n d -> b n (h d)', h=h))

@@ -256,6 +256,9 @@ def sub_quad_attention_forward(self, x, context=None, mask=None):
     k = k.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)

     v = v.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)

 

+    if q.device.type == 'mps':

+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()

+

     dtype = q.dtype

     if shared.opts.upcast_attn:

         q, k = q.float(), k.float()

@@ -293,7 +296,6 @@ def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_
     if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes:

         # the big matmul fits into our memory limit; do everything in 1 chunk,

         # i.e. send it down the unchunked fast-path

-        query_chunk_size = q_tokens

         kv_chunk_size = k_tokens

 

     with devices.without_autocast(disable=q.dtype == v.dtype):

@@ -332,7 +334,7 @@ def xformers_attention_forward(self, x, context=None, mask=None):
     k_in = self.to_k(context_k)

     v_in = self.to_v(context_v)

 

-    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 = (rearrange(t, 'b n (h d) -> b n h d', h=h) for t in (q_in, k_in, v_in))

     del q_in, k_in, v_in

 

     dtype = q.dtype

@@ -367,7 +369,7 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None):
     q = q_in.view(batch_size, -1, h, head_dim).transpose(1, 2)

     k = k_in.view(batch_size, -1, h, head_dim).transpose(1, 2)

     v = v_in.view(batch_size, -1, h, head_dim).transpose(1, 2)

-    

+

     del q_in, k_in, v_in

 

     dtype = q.dtype

@@ -449,7 +451,7 @@ def cross_attention_attnblock_forward(self, x):
         h3 += x

 

         return h3

-    

+

 def xformers_attnblock_forward(self, x):

     try:

         h_ = x

@@ -458,7 +460,7 @@ def xformers_attnblock_forward(self, x):
         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, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v))

         dtype = q.dtype

         if shared.opts.upcast_attn:

             q, k = q.float(), k.float()

@@ -480,7 +482,7 @@ def sdp_attnblock_forward(self, x):
     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, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v))

     dtype = q.dtype

     if shared.opts.upcast_attn:

         q, k = q.float(), k.float()

@@ -504,7 +506,7 @@ def sub_quad_attnblock_forward(self, x):
     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, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v))

     q = q.contiguous()

     k = k.contiguous()

     v = v.contiguous()