mirror of
https://github.com/Anjok07/ultimatevocalremovergui.git
synced 2024-11-28 01:10:56 +01:00
110 lines
2.9 KiB
Python
110 lines
2.9 KiB
Python
from functools import wraps
|
|
from packaging import version
|
|
from collections import namedtuple
|
|
|
|
import torch
|
|
from torch import nn, einsum
|
|
import torch.nn.functional as F
|
|
|
|
from einops import rearrange, reduce
|
|
|
|
# constants
|
|
|
|
FlashAttentionConfig = namedtuple('FlashAttentionConfig', ['enable_flash', 'enable_math', 'enable_mem_efficient'])
|
|
|
|
# helpers
|
|
|
|
def exists(val):
|
|
return val is not None
|
|
|
|
def once(fn):
|
|
called = False
|
|
@wraps(fn)
|
|
def inner(x):
|
|
nonlocal called
|
|
if called:
|
|
return
|
|
called = True
|
|
return fn(x)
|
|
return inner
|
|
|
|
print_once = once(print)
|
|
|
|
# main class
|
|
|
|
class Attend(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dropout = 0.,
|
|
flash = False
|
|
):
|
|
super().__init__()
|
|
self.dropout = dropout
|
|
self.attn_dropout = nn.Dropout(dropout)
|
|
|
|
self.flash = flash
|
|
assert not (flash and version.parse(torch.__version__) < version.parse('2.0.0')), 'in order to use flash attention, you must be using pytorch 2.0 or above'
|
|
|
|
# determine efficient attention configs for cuda and cpu
|
|
|
|
self.cpu_config = FlashAttentionConfig(True, True, True)
|
|
self.cuda_config = None
|
|
|
|
if not torch.cuda.is_available() or not flash:
|
|
return
|
|
|
|
device_properties = torch.cuda.get_device_properties(torch.device('cuda'))
|
|
|
|
if device_properties.major == 8 and device_properties.minor == 0:
|
|
print_once('A100 GPU detected, using flash attention if input tensor is on cuda')
|
|
self.cuda_config = FlashAttentionConfig(True, False, False)
|
|
else:
|
|
self.cuda_config = FlashAttentionConfig(False, True, True)
|
|
|
|
def flash_attn(self, q, k, v):
|
|
_, heads, q_len, _, k_len, is_cuda, device = *q.shape, k.shape[-2], q.is_cuda, q.device
|
|
|
|
# Check if there is a compatible device for flash attention
|
|
|
|
config = self.cuda_config if is_cuda else self.cpu_config
|
|
|
|
# pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale
|
|
|
|
with torch.backends.cuda.sdp_kernel(**config._asdict()):
|
|
out = F.scaled_dot_product_attention(
|
|
q, k, v,
|
|
dropout_p = self.dropout if self.training else 0.
|
|
)
|
|
|
|
return out
|
|
|
|
def forward(self, q, k, v):
|
|
"""
|
|
einstein notation
|
|
b - batch
|
|
h - heads
|
|
n, i, j - sequence length (base sequence length, source, target)
|
|
d - feature dimension
|
|
"""
|
|
|
|
q_len, k_len, device = q.shape[-2], k.shape[-2], q.device
|
|
|
|
scale = q.shape[-1] ** -0.5
|
|
|
|
if self.flash:
|
|
return self.flash_attn(q, k, v)
|
|
|
|
# similarity
|
|
|
|
sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale
|
|
|
|
# attention
|
|
|
|
attn = sim.softmax(dim=-1)
|
|
attn = self.attn_dropout(attn)
|
|
|
|
# aggregate values
|
|
|
|
out = einsum(f"b h i j, b h j d -> b h i d", attn, v)
|
|
|
|
return out |