mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
478 lines
15 KiB
Python
478 lines
15 KiB
Python
# https://github.com/Comfy-Org/ComfyUI/blob/v0.7.0/comfy/ldm/modules/attention.py
|
|
|
|
import logging
|
|
import math
|
|
|
|
import torch
|
|
from einops import rearrange, repeat
|
|
from torch import einsum
|
|
|
|
from backend import memory_management, operations
|
|
from backend.args import args
|
|
from backend.logging import setup_logger
|
|
|
|
logger = logging.getLogger("attention")
|
|
setup_logger(logger)
|
|
|
|
|
|
if memory_management.xformers_enabled() or memory_management.xformers_enabled_vae():
|
|
import xformers
|
|
import xformers.ops
|
|
|
|
|
|
if memory_management.sage_enabled():
|
|
import importlib.metadata
|
|
|
|
IS_SAGE_3 = False
|
|
|
|
if importlib.metadata.version("sageattention").startswith("1"):
|
|
IS_SAGE_1 = True
|
|
from sageattention import sageattn
|
|
else:
|
|
IS_SAGE_1 = False
|
|
from functools import partial
|
|
|
|
import sageattention
|
|
|
|
from backend.args import SageAttentionFuncs
|
|
|
|
match args.sage_function:
|
|
case SageAttentionFuncs.auto:
|
|
sageattn = sageattention.sageattn
|
|
case SageAttentionFuncs.fp16_triton:
|
|
sageattn = sageattention.sageattn_qk_int8_pv_fp16_triton
|
|
case SageAttentionFuncs.fp16_cuda:
|
|
sageattn = partial(sageattention.sageattn_qk_int8_pv_fp16_cuda, pv_accum_dtype="fp32")
|
|
case SageAttentionFuncs.fp8_cuda:
|
|
sageattn = partial(sageattention.sageattn_qk_int8_pv_fp8_cuda, pv_accum_dtype="fp32+fp32")
|
|
case SageAttentionFuncs.fp8_cuda_pp:
|
|
sageattn = partial(sageattention.sageattn_qk_int8_pv_fp8_cuda, pv_accum_dtype="fp32+fp16")
|
|
case SageAttentionFuncs.sageattn3:
|
|
from sageattn3 import sageattn3_blackwell
|
|
|
|
IS_SAGE_3 = True
|
|
|
|
def sageattn(q, k, v, attn_mask=None, is_causal=False, tensor_layout="NHD"):
|
|
q, k, v = [x.transpose(1, 2) if tensor_layout == "NHD" else x for x in (q, k, v)]
|
|
out = sageattn3_blackwell(q, k, v, is_causal=is_causal, attn_mask=attn_mask)
|
|
return out.transpose(1, 2) if tensor_layout == "NHD" else out
|
|
|
|
|
|
if memory_management.flash_enabled():
|
|
from flash_attn import flash_attn_func
|
|
|
|
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
|
|
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
|
|
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
|
|
|
|
@flash_attn_wrapper.register_fake
|
|
def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False):
|
|
return q.new_empty(q.shape)
|
|
|
|
|
|
def get_attn_precision(attn_precision: torch.dtype, current_dtype: torch.dtype) -> torch.dtype:
|
|
memory_management.force_upcast_attention_dtype().get(current_dtype, attn_precision)
|
|
|
|
|
|
def exists(val) -> bool:
|
|
return val is not None
|
|
|
|
|
|
if memory_management.is_nvidia():
|
|
SDP_BATCH_LIMIT = 2**15
|
|
else:
|
|
SDP_BATCH_LIMIT = 2**31
|
|
|
|
|
|
# region Attentions
|
|
|
|
|
|
def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
attn_precision = get_attn_precision(attn_precision, q.dtype)
|
|
|
|
if skip_reshape:
|
|
b, _, _, dim_head = q.shape
|
|
else:
|
|
b, _, dim_head = q.shape
|
|
dim_head //= heads
|
|
|
|
scale = dim_head**-0.5
|
|
|
|
h = heads
|
|
if skip_reshape:
|
|
q, k, v = map(
|
|
lambda t: t.reshape(b * heads, -1, dim_head),
|
|
(q, k, v),
|
|
)
|
|
else:
|
|
q, k, v = map(
|
|
lambda t: t.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(),
|
|
(q, k, v),
|
|
)
|
|
|
|
if attn_precision == torch.float32:
|
|
sim = einsum("b i d, b j d -> b i j", q.float(), k.float()) * scale
|
|
else:
|
|
sim = einsum("b i d, b j d -> b i j", q, k) * scale
|
|
|
|
del q, k
|
|
|
|
if exists(mask):
|
|
if mask.dtype == torch.bool:
|
|
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)
|
|
else:
|
|
if len(mask.shape) == 2:
|
|
bs = 1
|
|
else:
|
|
bs = mask.shape[0]
|
|
mask = mask.reshape(bs, -1, mask.shape[-2], mask.shape[-1]).expand(b, heads, -1, -1).reshape(-1, mask.shape[-2], mask.shape[-1])
|
|
sim.add_(mask)
|
|
|
|
sim = sim.softmax(dim=-1)
|
|
|
|
out = einsum("b i j, b j d -> b i d", sim.to(v.dtype), v)
|
|
|
|
if skip_output_reshape:
|
|
out = out.unsqueeze(0).reshape(b, heads, -1, dim_head)
|
|
else:
|
|
out = out.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head)
|
|
|
|
return out
|
|
|
|
|
|
def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
b = q.shape[0]
|
|
dim_head = q.shape[-1]
|
|
|
|
if torch.jit.is_tracing() or torch.jit.is_scripting():
|
|
return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs)
|
|
|
|
if skip_reshape:
|
|
q, k, v = map(
|
|
lambda t: t.permute(0, 2, 1, 3),
|
|
(q, k, v),
|
|
)
|
|
else:
|
|
dim_head //= heads
|
|
q, k, v = map(
|
|
lambda t: t.reshape(b, -1, heads, dim_head),
|
|
(q, k, v),
|
|
)
|
|
|
|
if mask is not None:
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
if mask.ndim == 3:
|
|
mask = mask.unsqueeze(1)
|
|
pad = 8 - mask.shape[-1] % 8
|
|
mask_out = torch.empty([mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device)
|
|
mask_out[..., : mask.shape[-1]] = mask
|
|
mask = mask_out[..., : mask.shape[-1]]
|
|
mask = mask.expand(b, heads, -1, -1)
|
|
|
|
try:
|
|
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
|
|
_fallback = False
|
|
except Exception as e:
|
|
if "(too new)" in str(e):
|
|
logger.warning("xformers does not work on RTX 50s")
|
|
else:
|
|
logger.error(f"Error running xformers: {e}")
|
|
_fallback = True
|
|
|
|
if _fallback:
|
|
if not skip_reshape:
|
|
q, k, v = map(
|
|
lambda t: t.transpose(1, 2),
|
|
(q, k, v),
|
|
)
|
|
return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, **kwargs)
|
|
|
|
if skip_output_reshape:
|
|
out = out.permute(0, 2, 1, 3)
|
|
else:
|
|
out = out.reshape(b, -1, heads * dim_head)
|
|
|
|
return out
|
|
|
|
|
|
def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
if skip_reshape:
|
|
b, _, _, dim_head = q.shape
|
|
else:
|
|
b, _, dim_head = q.shape
|
|
dim_head //= heads
|
|
q, k, v = map(
|
|
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
|
|
(q, k, v),
|
|
)
|
|
|
|
if mask is not None:
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
if mask.ndim == 3:
|
|
mask = mask.unsqueeze(1)
|
|
|
|
if SDP_BATCH_LIMIT >= b:
|
|
out = operations.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
|
if not skip_output_reshape:
|
|
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
|
else:
|
|
out = torch.empty((b, q.shape[2], heads * dim_head), dtype=q.dtype, layout=q.layout, device=q.device)
|
|
for i in range(0, b, SDP_BATCH_LIMIT):
|
|
m = mask
|
|
if mask is not None:
|
|
if mask.shape[0] > 1:
|
|
m = mask[i : i + SDP_BATCH_LIMIT]
|
|
|
|
out[i : i + SDP_BATCH_LIMIT] = operations.scaled_dot_product_attention(q[i : i + SDP_BATCH_LIMIT], k[i : i + SDP_BATCH_LIMIT], v[i : i + SDP_BATCH_LIMIT], attn_mask=m, dropout_p=0.0, is_causal=False).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head)
|
|
|
|
return out
|
|
|
|
|
|
@torch.compiler.disable
|
|
def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
in_dtype = v.dtype
|
|
if torch.float32 in (q.dtype, k.dtype, v.dtype):
|
|
q, k, v = q.to(torch.float16), k.to(torch.float16), v.to(torch.float16)
|
|
|
|
if skip_reshape:
|
|
b, _, _, dim_head = q.shape
|
|
tensor_layout = "HND"
|
|
else:
|
|
b, _, dim_head = q.shape
|
|
dim_head //= heads
|
|
q, k, v = map(
|
|
lambda t: t.view(b, -1, heads, dim_head),
|
|
(q, k, v),
|
|
)
|
|
tensor_layout = "NHD"
|
|
|
|
if mask is not None:
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
if mask.ndim == 3:
|
|
mask = mask.unsqueeze(1)
|
|
|
|
_fallback: bool = ((not IS_SAGE_1) and dim_head > 128) or (IS_SAGE_1 and (dim_head not in (64, 96, 128)))
|
|
|
|
try:
|
|
if not _fallback:
|
|
out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout).to(in_dtype)
|
|
except Exception as e:
|
|
logger.error(f"Error running sageattn: {e}")
|
|
_fallback = True
|
|
|
|
if _fallback:
|
|
if tensor_layout == "NHD":
|
|
q, k, v = map(
|
|
lambda t: t.transpose(1, 2),
|
|
(q, k, v),
|
|
)
|
|
return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape, **kwargs)
|
|
|
|
if tensor_layout == "HND":
|
|
if not skip_output_reshape:
|
|
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
|
else:
|
|
if skip_output_reshape:
|
|
out = out.transpose(1, 2)
|
|
else:
|
|
out = out.reshape(b, -1, heads * dim_head)
|
|
|
|
return out
|
|
|
|
|
|
def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
if skip_reshape:
|
|
b, _, _, dim_head = q.shape
|
|
else:
|
|
b, _, dim_head = q.shape
|
|
dim_head //= heads
|
|
q, k, v = map(
|
|
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
|
|
(q, k, v),
|
|
)
|
|
|
|
if mask is not None:
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
if mask.ndim == 3:
|
|
mask = mask.unsqueeze(1)
|
|
|
|
try:
|
|
assert mask is None
|
|
out = flash_attn_wrapper(
|
|
q.transpose(1, 2),
|
|
k.transpose(1, 2),
|
|
v.transpose(1, 2),
|
|
dropout_p=0.0,
|
|
causal=False,
|
|
).transpose(1, 2)
|
|
_fallback = False
|
|
except Exception as e:
|
|
logger.error(f"Error running flash_attn: {e}")
|
|
_fallback = True
|
|
|
|
if _fallback:
|
|
out = operations.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
|
|
|
if not skip_output_reshape:
|
|
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
|
|
|
return out
|
|
|
|
|
|
if memory_management.sage_enabled():
|
|
attention_function = attention_sage
|
|
if IS_SAGE_1:
|
|
logger.info("Using SageAttention")
|
|
elif IS_SAGE_3:
|
|
logger.info("Using SageAttention 3")
|
|
else:
|
|
match args.sage_function:
|
|
case SageAttentionFuncs.auto:
|
|
logger.info("Using SageAttention 2")
|
|
case SageAttentionFuncs.fp16_triton:
|
|
logger.info("Using SageAttention 2 (fp16 Triton)")
|
|
case SageAttentionFuncs.fp16_cuda:
|
|
logger.info("Using SageAttention 2 (fp16 CUDA)")
|
|
case SageAttentionFuncs.fp8_cuda:
|
|
logger.info("Using SageAttention 2 (fp8 CUDA)")
|
|
case SageAttentionFuncs.fp8_cuda_pp:
|
|
logger.info("Using SageAttention 2 (fp8 CUDA ++)")
|
|
|
|
elif memory_management.flash_enabled():
|
|
logger.info("Using FlashAttention")
|
|
attention_function = attention_flash
|
|
elif memory_management.xformers_enabled():
|
|
logger.info("Using xformers Cross Attention")
|
|
attention_function = attention_xformers
|
|
elif memory_management.pytorch_attention_enabled():
|
|
logger.info("Using PyTorch Cross Attention")
|
|
attention_function = attention_pytorch
|
|
else:
|
|
logger.info("Using Basic Cross Attention")
|
|
attention_function = attention_basic
|
|
|
|
|
|
# region VAE
|
|
|
|
|
|
def slice_attention_vae(q, k, v):
|
|
r1 = torch.zeros_like(k, device=q.device)
|
|
scale = int(q.shape[-1]) ** (-0.5)
|
|
|
|
mem_free_total = memory_management.get_free_memory(q.device)
|
|
|
|
tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * 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)))
|
|
|
|
while True:
|
|
try:
|
|
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 = torch.bmm(q[:, i:end], k) * scale
|
|
|
|
s2 = torch.nn.functional.softmax(s1, dim=2).permute(0, 2, 1)
|
|
del s1
|
|
|
|
r1[:, :, i:end] = torch.bmm(v, s2)
|
|
del s2
|
|
break
|
|
except Exception as e:
|
|
if not memory_management.is_oom(e):
|
|
raise e
|
|
if steps > 128:
|
|
raise e
|
|
|
|
logger.warning("Out of Memory Error; retrying with higher steps...")
|
|
memory_management.soft_empty_cache()
|
|
steps *= 2
|
|
|
|
return r1
|
|
|
|
|
|
def normal_attention_vae(q, k, v):
|
|
orig_shape = q.shape
|
|
b = orig_shape[0]
|
|
c = orig_shape[1]
|
|
|
|
q = q.reshape(b, c, -1)
|
|
q = q.permute(0, 2, 1)
|
|
k = k.reshape(b, c, -1)
|
|
v = v.reshape(b, c, -1)
|
|
|
|
r1 = slice_attention_vae(q, k, v)
|
|
h_ = r1.reshape(orig_shape)
|
|
del r1
|
|
return h_
|
|
|
|
|
|
def xformers_attention_vae(q, k, v):
|
|
orig_shape = q.shape
|
|
B = orig_shape[0]
|
|
C = orig_shape[1]
|
|
q, k, v = map(
|
|
lambda t: t.view(B, C, -1).transpose(1, 2).contiguous(),
|
|
(q, k, v),
|
|
)
|
|
|
|
try:
|
|
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
|
|
out = out.transpose(1, 2).reshape(orig_shape)
|
|
_fallback = False
|
|
except Exception:
|
|
_fallback = True
|
|
|
|
if _fallback:
|
|
out = slice_attention_vae(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(orig_shape)
|
|
|
|
return out
|
|
|
|
|
|
def pytorch_attention_vae(q, k, v):
|
|
orig_shape = q.shape
|
|
B = orig_shape[0]
|
|
C = orig_shape[1]
|
|
q, k, v = map(
|
|
lambda t: t.view(B, 1, C, -1).transpose(2, 3).contiguous(),
|
|
(q, k, v),
|
|
)
|
|
|
|
try:
|
|
out = operations.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
|
|
out = out.transpose(2, 3).reshape(orig_shape)
|
|
_fallback = False
|
|
except Exception as e:
|
|
if not memory_management.is_oom(e):
|
|
raise e
|
|
logger.warning("Out of Memory Error; retrying with Slice Attention")
|
|
_fallback = True
|
|
|
|
if _fallback:
|
|
memory_management.soft_empty_cache()
|
|
out = slice_attention_vae(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(orig_shape)
|
|
|
|
return out
|
|
|
|
|
|
if memory_management.xformers_enabled_vae():
|
|
logger.info("Using xformers Attention for VAE")
|
|
attention_function_vae = xformers_attention_vae
|
|
elif memory_management.pytorch_attention_enabled():
|
|
logger.info("Using PyTorch Attention for VAE")
|
|
attention_function_vae = pytorch_attention_vae
|
|
else:
|
|
logger.info("Using Slice Attention for VAE")
|
|
attention_function_vae = normal_attention_vae
|