modules/sd_hijack_optimizations.py

# from __future__ import annotations
# import math
# import psutil
# import platform
#
# import torch
# from torch import einsum
#
# from ldm.util import default
# from einops import rearrange
#
# from modules import shared, errors, devices, sub_quadratic_attention
# from modules.hypernetworks import hypernetwork
#
# import ldm.modules.attention
# import ldm.modules.diffusionmodules.model
#
# import sgm.modules.attention
# import sgm.modules.diffusionmodules.model
#
# diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.AttnBlock.forward
# sgm_diffusionmodules_model_AttnBlock_forward = sgm.modules.diffusionmodules.model.AttnBlock.forward
#
#
# class SdOptimization:
#     name: str = None
#     label: str | None = None
#     cmd_opt: str | None = None
#     priority: int = 0
#
#     def title(self):
#         if self.label is None:
#             return self.name
#
#         return f"{self.name} - {self.label}"
#
#     def is_available(self):
#         return True
#
#     def apply(self):
#         pass
#
#     def undo(self):
#         ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
#         ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward
#
#         sgm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
#         sgm.modules.diffusionmodules.model.AttnBlock.forward = sgm_diffusionmodules_model_AttnBlock_forward
#
#
# class SdOptimizationXformers(SdOptimization):
#     name = "xformers"
#     cmd_opt = "xformers"
#     priority = 100
#
#     def is_available(self):
#         return shared.cmd_opts.force_enable_xformers or (shared.xformers_available and torch.cuda.is_available() and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (9, 0))
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = xformers_attention_forward
#         ldm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward
#         sgm.modules.attention.CrossAttention.forward = xformers_attention_forward
#         sgm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward
#
#
# class SdOptimizationSdpNoMem(SdOptimization):
#     name = "sdp-no-mem"
#     label = "scaled dot product without memory efficient attention"
#     cmd_opt = "opt_sdp_no_mem_attention"
#     priority = 80
#
#     def is_available(self):
#         return hasattr(torch.nn.functional, "scaled_dot_product_attention") and callable(torch.nn.functional.scaled_dot_product_attention)
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = scaled_dot_product_no_mem_attention_forward
#         ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_no_mem_attnblock_forward
#         sgm.modules.attention.CrossAttention.forward = scaled_dot_product_no_mem_attention_forward
#         sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_no_mem_attnblock_forward
#
#
# class SdOptimizationSdp(SdOptimizationSdpNoMem):
#     name = "sdp"
#     label = "scaled dot product"
#     cmd_opt = "opt_sdp_attention"
#     priority = 70
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward
#         ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward
#         sgm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward
#         sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward
#
#
# class SdOptimizationSubQuad(SdOptimization):
#     name = "sub-quadratic"
#     cmd_opt = "opt_sub_quad_attention"
#
#     @property
#     def priority(self):
#         return 1000 if shared.device.type == 'mps' else 10
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = sub_quad_attention_forward
#         ldm.modules.diffusionmodules.model.AttnBlock.forward = sub_quad_attnblock_forward
#         sgm.modules.attention.CrossAttention.forward = sub_quad_attention_forward
#         sgm.modules.diffusionmodules.model.AttnBlock.forward = sub_quad_attnblock_forward
#
#
# class SdOptimizationV1(SdOptimization):
#     name = "V1"
#     label = "original v1"
#     cmd_opt = "opt_split_attention_v1"
#     priority = 10
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1
#         sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1
#
#
# class SdOptimizationInvokeAI(SdOptimization):
#     name = "InvokeAI"
#     cmd_opt = "opt_split_attention_invokeai"
#
#     @property
#     def priority(self):
#         return 1000 if shared.device.type != 'mps' and not torch.cuda.is_available() else 10
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeAI
#         sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeAI
#
#
# class SdOptimizationDoggettx(SdOptimization):
#     name = "Doggettx"
#     cmd_opt = "opt_split_attention"
#     priority = 90
#
#     def apply(self):
#         ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward
#         ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
#         sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward
#         sgm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
#
#
# def list_optimizers(res):
#     res.extend([
#         SdOptimizationXformers(),
#         SdOptimizationSdpNoMem(),
#         SdOptimizationSdp(),
#         SdOptimizationSubQuad(),
#         SdOptimizationV1(),
#         SdOptimizationInvokeAI(),
#         SdOptimizationDoggettx(),
#     ])
#
#
# if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers:
#     try:
#         import xformers.ops
#         shared.xformers_available = True
#     except Exception:
#         errors.report("Cannot import xformers", exc_info=True)
#
#
# def get_available_vram():
#     if shared.device.type == 'cuda':
#         stats = torch.cuda.memory_stats(shared.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
#         return mem_free_total
#     else:
#         return psutil.virtual_memory().available
#
#
# # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
# def split_cross_attention_forward_v1(self, x, context=None, mask=None, **kwargs):
#     h = self.heads
#
#     q_in = self.to_q(x)
#     context = default(context, x)
#
#     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
#     k_in = self.to_k(context_k)
#     v_in = self.to_v(context_v)
#     del context, context_k, context_v, x
#
#     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
#     if shared.opts.upcast_attn:
#         q, k, v = q.float(), k.float(), v.float()
#
#     with devices.without_autocast(disable=not shared.opts.upcast_attn):
#         r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
#         for i in range(0, q.shape[0], 2):
#             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
#
#     r1 = r1.to(dtype)
#
#     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
#     del r1
#
#     return self.to_out(r2)
#
#
# # taken from https://github.com/Doggettx/stable-diffusion and modified
# def split_cross_attention_forward(self, x, context=None, mask=None, **kwargs):
#     h = self.heads
#
#     q_in = self.to_q(x)
#     context = default(context, x)
#
#     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
#     k_in = self.to_k(context_k)
#     v_in = self.to_v(context_v)
#
#     dtype = q_in.dtype
#     if shared.opts.upcast_attn:
#         q_in, k_in, v_in = q_in.float(), k_in.float(), v_in if v_in.device.type == 'mps' else v_in.float()
#
#     with devices.without_autocast(disable=not shared.opts.upcast_attn):
#         k_in = k_in * self.scale
#
#         del context, x
#
#         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
#         for i in range(0, q.shape[1], slice_size):
#             end = min(i + slice_size, q.shape[1])
#             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)
#
#     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
#     del r1
#
#     return self.to_out(r2)
#
#
# # -- Taken from https://github.com/invoke-ai/InvokeAI and modified --
# mem_total_gb = psutil.virtual_memory().total // (1 << 30)
#
#
# def einsum_op_compvis(q, k, v):
#     s = einsum('b i d, b j d -> b i j', q, k)
#     s = s.softmax(dim=-1, dtype=s.dtype)
#     return einsum('b i j, b j d -> b i d', s, v)
#
#
# def einsum_op_slice_0(q, k, v, slice_size):
#     r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
#     for i in range(0, q.shape[0], slice_size):
#         end = i + slice_size
#         r[i:end] = einsum_op_compvis(q[i:end], k[i:end], v[i:end])
#     return r
#
#
# def einsum_op_slice_1(q, k, v, slice_size):
#     r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
#     for i in range(0, q.shape[1], slice_size):
#         end = i + slice_size
#         r[:, i:end] = einsum_op_compvis(q[:, i:end], k, v)
#     return r
#
#
# def einsum_op_mps_v1(q, k, v):
#     if q.shape[0] * q.shape[1] <= 2**16: # (512x512) max q.shape[1]: 4096
#         return einsum_op_compvis(q, k, v)
#     else:
#         slice_size = math.floor(2**30 / (q.shape[0] * q.shape[1]))
#         if slice_size % 4096 == 0:
#             slice_size -= 1
#         return einsum_op_slice_1(q, k, v, slice_size)
#
#
# def einsum_op_mps_v2(q, k, v):
#     if mem_total_gb > 8 and q.shape[0] * q.shape[1] <= 2**16:
#         return einsum_op_compvis(q, k, v)
#     else:
#         return einsum_op_slice_0(q, k, v, 1)
#
#
# def einsum_op_tensor_mem(q, k, v, max_tensor_mb):
#     size_mb = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() // (1 << 20)
#     if size_mb <= max_tensor_mb:
#         return einsum_op_compvis(q, k, v)
#     div = 1 << int((size_mb - 1) / max_tensor_mb).bit_length()
#     if div <= q.shape[0]:
#         return einsum_op_slice_0(q, k, v, q.shape[0] // div)
#     return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1))
#
#
# def einsum_op_cuda(q, k, v):
#     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(q.device)
#     mem_free_torch = mem_reserved - mem_active
#     mem_free_total = mem_free_cuda + mem_free_torch
#     # Divide factor of safety as there's copying and fragmentation
#     return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
#
#
# def einsum_op(q, k, v):
#     if q.device.type == 'cuda':
#         return einsum_op_cuda(q, k, v)
#
#     if q.device.type == 'mps':
#         if mem_total_gb >= 32 and q.shape[0] % 32 != 0 and q.shape[0] * q.shape[1] < 2**18:
#             return einsum_op_mps_v1(q, k, v)
#         return einsum_op_mps_v2(q, k, v)
#
#     # Smaller slices are faster due to L2/L3/SLC caches.
#     # Tested on i7 with 8MB L3 cache.
#     return einsum_op_tensor_mem(q, k, v, 32)
#
#
# def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None, **kwargs):
#     h = self.heads
#
#     q = self.to_q(x)
#     context = default(context, x)
#
#     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
#     k = self.to_k(context_k)
#     v = self.to_v(context_v)
#     del context, context_k, context_v, x
#
#     dtype = q.dtype
#     if shared.opts.upcast_attn:
#         q, k, v = q.float(), k.float(), v if v.device.type == 'mps' else v.float()
#
#     with devices.without_autocast(disable=not shared.opts.upcast_attn):
#         k = k * self.scale
#
#         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))
#
# # -- End of code from https://github.com/invoke-ai/InvokeAI --
#
#
# # Based on Birch-san's modified implementation of sub-quadratic attention from https://github.com/Birch-san/diffusers/pull/1
# # The sub_quad_attention_forward function is under the MIT License listed under Memory Efficient Attention in the Licenses section of the web UI interface
# def sub_quad_attention_forward(self, x, context=None, mask=None, **kwargs):
#     assert mask is None, "attention-mask not currently implemented for SubQuadraticCrossAttnProcessor."
#
#     h = self.heads
#
#     q = self.to_q(x)
#     context = default(context, x)
#
#     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
#     k = self.to_k(context_k)
#     v = self.to_v(context_v)
#     del context, context_k, context_v, x
#
#     q = q.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
#     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()
#
#     x = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training)
#
#     x = x.to(dtype)
#
#     x = x.unflatten(0, (-1, h)).transpose(1,2).flatten(start_dim=2)
#
#     out_proj, dropout = self.to_out
#     x = out_proj(x)
#     x = dropout(x)
#
#     return x
#
#
# def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_size_min=None, chunk_threshold=None, use_checkpoint=True):
#     bytes_per_token = torch.finfo(q.dtype).bits//8
#     batch_x_heads, q_tokens, _ = q.shape
#     _, k_tokens, _ = k.shape
#     qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens
#
#     if chunk_threshold is None:
#         if q.device.type == 'mps':
#             chunk_threshold_bytes = 268435456 * (2 if platform.processor() == 'i386' else bytes_per_token)
#         else:
#             chunk_threshold_bytes = int(get_available_vram() * 0.7)
#     elif chunk_threshold == 0:
#         chunk_threshold_bytes = None
#     else:
#         chunk_threshold_bytes = int(0.01 * chunk_threshold * get_available_vram())
#
#     if kv_chunk_size_min is None and chunk_threshold_bytes is not None:
#         kv_chunk_size_min = chunk_threshold_bytes // (batch_x_heads * bytes_per_token * (k.shape[2] + v.shape[2]))
#     elif kv_chunk_size_min == 0:
#         kv_chunk_size_min = None
#
#     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
#         kv_chunk_size = k_tokens
#
#     with devices.without_autocast(disable=q.dtype == v.dtype):
#         return sub_quadratic_attention.efficient_dot_product_attention(
#             q,
#             k,
#             v,
#             query_chunk_size=q_chunk_size,
#             kv_chunk_size=kv_chunk_size,
#             kv_chunk_size_min = kv_chunk_size_min,
#             use_checkpoint=use_checkpoint,
#         )
#
#
# def get_xformers_flash_attention_op(q, k, v):
#     if not shared.cmd_opts.xformers_flash_attention:
#         return None
#
#     try:
#         flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp
#         fw, bw = flash_attention_op
#         if fw.supports(xformers.ops.fmha.Inputs(query=q, key=k, value=v, attn_bias=None)):
#             return flash_attention_op
#     except Exception as e:
#         errors.display_once(e, "enabling flash attention")
#
#     return None
#
#
# def xformers_attention_forward(self, x, context=None, mask=None, **kwargs):
#     h = self.heads
#     q_in = self.to_q(x)
#     context = default(context, x)
#
#     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
#     k_in = self.to_k(context_k)
#     v_in = self.to_v(context_v)
#
#     q, k, v = (t.reshape(t.shape[0], t.shape[1], h, -1) for t in (q_in, k_in, v_in))
#
#     del q_in, k_in, v_in
#
#     dtype = q.dtype
#     if shared.opts.upcast_attn:
#         q, k, v = q.float(), k.float(), v.float()
#
#     out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=get_xformers_flash_attention_op(q, k, v))
#
#     out = out.to(dtype)
#
#     b, n, h, d = out.shape
#     out = out.reshape(b, n, h * d)
#     return self.to_out(out)
#
#
# # Based on Diffusers usage of scaled dot product attention from https://github.com/huggingface/diffusers/blob/c7da8fd23359a22d0df2741688b5b4f33c26df21/src/diffusers/models/cross_attention.py
# # The scaled_dot_product_attention_forward function contains parts of code under Apache-2.0 license listed under Scaled Dot Product Attention in the Licenses section of the web UI interface
# def scaled_dot_product_attention_forward(self, x, context=None, mask=None, **kwargs):
#     batch_size, sequence_length, inner_dim = x.shape
#
#     if mask is not None:
#         mask = self.prepare_attention_mask(mask, sequence_length, batch_size)
#         mask = mask.view(batch_size, self.heads, -1, mask.shape[-1])
#
#     h = self.heads
#     q_in = self.to_q(x)
#     context = default(context, x)
#
#     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
#     k_in = self.to_k(context_k)
#     v_in = self.to_v(context_v)
#
#     head_dim = inner_dim // h
#     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
#     if shared.opts.upcast_attn:
#         q, k, v = q.float(), k.float(), v.float()
#
#     # the output of sdp = (batch, num_heads, seq_len, head_dim)
#     hidden_states = torch.nn.functional.scaled_dot_product_attention(
#         q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
#     )
#
#     hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, h * head_dim)
#     hidden_states = hidden_states.to(dtype)
#
#     # linear proj
#     hidden_states = self.to_out[0](hidden_states)
#     # dropout
#     hidden_states = self.to_out[1](hidden_states)
#     return hidden_states
#
#
# def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None, **kwargs):
#     with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
#         return scaled_dot_product_attention_forward(self, x, context, mask)
#
#
# def cross_attention_attnblock_forward(self, x):
#         h_ = x
#         h_ = self.norm(h_)
#         q1 = self.q(h_)
#         k1 = self.k(h_)
#         v = self.v(h_)
#
#         # compute attention
#         b, c, h, w = q1.shape
#
#         q2 = q1.reshape(b, c, h*w)
#         del q1
#
#         q = q2.permute(0, 2, 1)   # b,hw,c
#         del q2
#
#         k = k1.reshape(b, c, h*w) # b,c,hw
#         del k1
#
#         h_ = torch.zeros_like(k, device=q.device)
#
#         mem_free_total = get_available_vram()
#
#         tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
#         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)))
#
#         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
#
#             w1 = torch.bmm(q[:, i:end], k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
#             w2 = w1 * (int(c)**(-0.5))
#             del w1
#             w3 = torch.nn.functional.softmax(w2, dim=2, dtype=q.dtype)
#             del w2
#
#             # attend to values
#             v1 = v.reshape(b, c, h*w)
#             w4 = w3.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
#             del w3
#
#             h_[:, :, i:end] = torch.bmm(v1, w4)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
#             del v1, w4
#
#         h2 = h_.reshape(b, c, h, w)
#         del h_
#
#         h3 = self.proj_out(h2)
#         del h2
#
#         h3 += x
#
#         return h3
#
#
# def xformers_attnblock_forward(self, x):
#     try:
#         h_ = x
#         h_ = self.norm(h_)
#         q = self.q(h_)
#         k = self.k(h_)
#         v = self.v(h_)
#         b, c, h, w = q.shape
#         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()
#         q = q.contiguous()
#         k = k.contiguous()
#         v = v.contiguous()
#         out = xformers.ops.memory_efficient_attention(q, k, v, op=get_xformers_flash_attention_op(q, k, v))
#         out = out.to(dtype)
#         out = rearrange(out, 'b (h w) c -> b c h w', h=h)
#         out = self.proj_out(out)
#         return x + out
#     except NotImplementedError:
#         return cross_attention_attnblock_forward(self, x)
#
#
# def sdp_attnblock_forward(self, x):
#     h_ = x
#     h_ = self.norm(h_)
#     q = self.q(h_)
#     k = self.k(h_)
#     v = self.v(h_)
#     b, c, h, w = q.shape
#     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, v = q.float(), k.float(), v.float()
#     q = q.contiguous()
#     k = k.contiguous()
#     v = v.contiguous()
#     out = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
#     out = out.to(dtype)
#     out = rearrange(out, 'b (h w) c -> b c h w', h=h)
#     out = self.proj_out(out)
#     return x + out
#
#
# def sdp_no_mem_attnblock_forward(self, x):
#     with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
#         return sdp_attnblock_forward(self, x)
#
#
# def sub_quad_attnblock_forward(self, x):
#     h_ = x
#     h_ = self.norm(h_)
#     q = self.q(h_)
#     k = self.k(h_)
#     v = self.v(h_)
#     b, c, h, w = q.shape
#     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()
#     out = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training)
#     out = rearrange(out, 'b (h w) c -> b c h w', h=h)
#     out = self.proj_out(out)
#     return x + out