-
-
Notifications
You must be signed in to change notification settings - Fork 98
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add support for BS-RoFormer and Mel-Band RoFormer models (#72)
* First pass at adding bs and mel-roformer models to MDXC, not yet tested * Got some things kinda working * Added working implementation of Mel-Roformer * Made mel-roformer the default model
- Loading branch information
Showing
11 changed files
with
1,638 additions
and
224 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
268 changes: 189 additions & 79 deletions
268
audio_separator/separator/architectures/mdxc_separator.py
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| 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.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.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 |
Oops, something went wrong.