Import RVC 20260716 Nvidia 50x0 v2bb

tools/file_io.py

@@ -0,0 +1,15 @@
+def read_text(path, errors="strict", newline=None):
+    last_error = None
+    for encoding in (None, "utf8", "gbk"):
+        try:
+            kwargs = {"errors": "strict", "newline": newline}
+            if encoding is not None:
+                kwargs["encoding"] = encoding
+            with open(path, "r", **kwargs) as file:
+                return file.read()
+        except UnicodeDecodeError as error:
+            last_error = error
+    if errors != "strict":
+        with open(path, "r", encoding="gbk", errors=errors, newline=newline) as file:
+            return file.read()
+    raise last_error

tools/progress.py

@@ -0,0 +1,46 @@
+import math
+
+from i18n.i18n import I18nAuto
+
+
+i18n = I18nAuto()
+
+
+def should_report(index, total, max_updates=12):
+    if total <= 0:
+        return False
+    if total <= max_updates:
+        return True
+    interval = max(1, math.ceil(total / max_updates))
+    return index == 0 or index + 1 == total or (index + 1) % interval == 0
+
+
+def batch_status(title, current, total, success, failed, latest="", failures=None):
+    if total <= 0:
+        state = i18n("等待输入")
+    elif current >= total:
+        state = i18n("已完成")
+    else:
+        state = i18n("处理中")
+    lines = [
+        "【%s】" % title,
+        "%s：%s" % (i18n("状态"), state),
+        "%s：%s/%s | %s：%s | %s：%s"
+        % (
+            i18n("进度"),
+            current,
+            total,
+            i18n("成功"),
+            success,
+            i18n("失败"),
+            failed,
+        ),
+    ]
+    if latest:
+        lines.append("%s：%s" % (i18n("当前"), latest))
+    if failures:
+        lines.append("%s：" % i18n("失败记录"))
+        lines.extend(failures[-10:])
+        if len(failures) > 10:
+            lines.append(i18n("……仅显示最近10条失败记录"))
+    return "\n".join(lines)

tools/torchgate/__init__.py

@@ -0,0 +1,13 @@
+"""
+TorchGating is a PyTorch-based implementation of Spectral Gating
+================================================
+Author: Asaf Zorea
+
+Contents
+--------
+torchgate imports all the functions from PyTorch, and in addition provides:
+ TorchGating       --- A PyTorch module that applies a spectral gate to an input signal
+
+"""
+
+from .torchgate import TorchGate

tools/torchgate/torchgate.py

@@ -0,0 +1,280 @@
+import torch
+from infer.rmvpe import STFT
+from torch.nn.functional import conv1d, conv2d
+from typing import Union, Optional
+from .utils import linspace, temperature_sigmoid, amp_to_db
+
+
+class TorchGate(torch.nn.Module):
+    """
+    A PyTorch module that applies a spectral gate to an input signal.
+
+    Arguments:
+        sr {int} -- Sample rate of the input signal.
+        nonstationary {bool} -- Whether to use non-stationary or stationary masking (default: {False}).
+        n_std_thresh_stationary {float} -- Number of standard deviations above mean to threshold noise for
+                                           stationary masking (default: {1.5}).
+        n_thresh_nonstationary {float} -- Number of multiplies above smoothed magnitude spectrogram. for
+                                        non-stationary masking (default: {1.3}).
+        temp_coeff_nonstationary {float} -- Temperature coefficient for non-stationary masking (default: {0.1}).
+        n_movemean_nonstationary {int} -- Number of samples for moving average smoothing in non-stationary masking
+                                          (default: {20}).
+        prop_decrease {float} -- Proportion to decrease signal by where the mask is zero (default: {1.0}).
+        n_fft {int} -- Size of FFT for STFT (default: {1024}).
+        win_length {[int]} -- Window length for STFT. If None, defaults to `n_fft` (default: {None}).
+        hop_length {[int]} -- Hop length for STFT. If None, defaults to `win_length` // 4 (default: {None}).
+        freq_mask_smooth_hz {float} -- Frequency smoothing width for mask (in Hz). If None, no smoothing is applied
+                                     (default: {500}).
+        time_mask_smooth_ms {float} -- Time smoothing width for mask (in ms). If None, no smoothing is applied
+                                     (default: {50}).
+    """
+
+    @torch.no_grad()
+    def __init__(
+        self,
+        sr,
+        nonstationary = False,
+        n_std_thresh_stationary = 1.5,
+        n_thresh_nonstationary = 1.3,
+        temp_coeff_nonstationary = 0.1,
+        n_movemean_nonstationary = 20,
+        prop_decrease = 1.0,
+        n_fft = 1024,
+        win_length = None,
+        hop_length = None,
+        freq_mask_smooth_hz = 500,
+        time_mask_smooth_ms = 50,
+    ):
+        super().__init__()
+
+        # General Params
+        self.sr = sr
+        self.nonstationary = nonstationary
+        assert 0.0 <= prop_decrease <= 1.0
+        self.prop_decrease = prop_decrease
+
+        # STFT Params
+        self.n_fft = n_fft
+        self.win_length = self.n_fft if win_length is None else win_length
+        self.hop_length = self.win_length // 4 if hop_length is None else hop_length
+
+        # Stationary Params
+        self.n_std_thresh_stationary = n_std_thresh_stationary
+
+        # Non-Stationary Params
+        self.temp_coeff_nonstationary = temp_coeff_nonstationary
+        self.n_movemean_nonstationary = n_movemean_nonstationary
+        self.n_thresh_nonstationary = n_thresh_nonstationary
+
+        # Smooth Mask Params
+        self.freq_mask_smooth_hz = freq_mask_smooth_hz
+        self.time_mask_smooth_ms = time_mask_smooth_ms
+        self.register_buffer("smoothing_filter", self._generate_mask_smoothing_filter())
+
+    @torch.no_grad()
+    def _generate_mask_smoothing_filter(self) :
+        """
+        A PyTorch module that applies a spectral gate to an input signal using the STFT.
+
+        Returns:
+            smoothing_filter (torch.Tensor): a 2D tensor representing the smoothing filter,
+            with shape (n_grad_freq, n_grad_time), where n_grad_freq is the number of frequency
+            bins to smooth and n_grad_time is the number of time frames to smooth.
+            If both self.freq_mask_smooth_hz and self.time_mask_smooth_ms are None, returns None.
+        """
+        if self.freq_mask_smooth_hz is None and self.time_mask_smooth_ms is None:
+            return None
+
+        n_grad_freq = (
+            1
+            if self.freq_mask_smooth_hz is None
+            else int(self.freq_mask_smooth_hz / (self.sr / (self.n_fft / 2)))
+        )
+        if n_grad_freq < 1:
+            raise ValueError(
+                f"freq_mask_smooth_hz needs to be at least {int((self.sr / (self._n_fft / 2)))} Hz"
+            )
+
+        n_grad_time = (
+            1
+            if self.time_mask_smooth_ms is None
+            else int(self.time_mask_smooth_ms / ((self.hop_length / self.sr) * 1000))
+        )
+        if n_grad_time < 1:
+            raise ValueError(
+                f"time_mask_smooth_ms needs to be at least {int((self.hop_length / self.sr) * 1000)} ms"
+            )
+
+        if n_grad_time == 1 and n_grad_freq == 1:
+            return None
+
+        v_f = torch.cat(
+            [
+                linspace(0, 1, n_grad_freq + 1, endpoint=False),
+                linspace(1, 0, n_grad_freq + 2),
+            ]
+        )[1:-1]
+        v_t = torch.cat(
+            [
+                linspace(0, 1, n_grad_time + 1, endpoint=False),
+                linspace(1, 0, n_grad_time + 2),
+            ]
+        )[1:-1]
+        smoothing_filter = torch.outer(v_f, v_t).unsqueeze(0).unsqueeze(0)
+
+        return smoothing_filter / smoothing_filter.sum()
+
+    @torch.no_grad()
+    def _stationary_mask(
+        self, X_db, xn = None
+    ) :
+        """
+        Computes a stationary binary mask to filter out noise in a log-magnitude spectrogram.
+
+        Arguments:
+            X_db (torch.Tensor): 2D tensor of shape (frames, freq_bins) containing the log-magnitude spectrogram.
+            xn (torch.Tensor): 1D tensor containing the audio signal corresponding to X_db.
+
+        Returns:
+            sig_mask (torch.Tensor): Binary mask of the same shape as X_db, where values greater than the threshold
+            are set to 1, and the rest are set to 0.
+        """
+        if xn is not None:
+            if "privateuseone" in str(xn.device):
+                if not hasattr(self, "stft"):
+                    self.stft = STFT(
+                        filter_length=self.n_fft,
+                        hop_length=self.hop_length,
+                        win_length=self.win_length,
+                        window="hann",
+                    ).to(xn.device)
+                XN = self.stft.transform(xn)
+            else:
+                XN = torch.stft(
+                    xn,
+                    n_fft=self.n_fft,
+                    hop_length=self.hop_length,
+                    win_length=self.win_length,
+                    return_complex=True,
+                    pad_mode="constant",
+                    center=True,
+                    window=torch.hann_window(self.win_length).to(xn.device),
+                )
+            XN_db = amp_to_db(XN).to(dtype=X_db.dtype)
+        else:
+            XN_db = X_db
+
+        # calculate mean and standard deviation along the frequency axis
+        std_freq_noise, mean_freq_noise = torch.std_mean(XN_db, dim=-1)
+
+        # compute noise threshold
+        noise_thresh = mean_freq_noise + std_freq_noise * self.n_std_thresh_stationary
+
+        # create binary mask by thresholding the spectrogram
+        sig_mask = X_db > noise_thresh.unsqueeze(2)
+        return sig_mask
+
+    @torch.no_grad()
+    def _nonstationary_mask(self, X_abs) :
+        """
+        Computes a non-stationary binary mask to filter out noise in a log-magnitude spectrogram.
+
+        Arguments:
+            X_abs (torch.Tensor): 2D tensor of shape (frames, freq_bins) containing the magnitude spectrogram.
+
+        Returns:
+            sig_mask (torch.Tensor): Binary mask of the same shape as X_abs, where values greater than the threshold
+            are set to 1, and the rest are set to 0.
+        """
+        X_smoothed = (
+            conv1d(
+                X_abs.reshape(-1, 1, X_abs.shape[-1]),
+                torch.ones(
+                    self.n_movemean_nonstationary,
+                    dtype=X_abs.dtype,
+                    device=X_abs.device,
+                ).view(1, 1, -1),
+                padding="same",
+            ).view(X_abs.shape)
+            / self.n_movemean_nonstationary
+        )
+
+        # Compute slowness ratio and apply temperature sigmoid
+        slowness_ratio = (X_abs - X_smoothed) / (X_smoothed + 1e-6)
+        sig_mask = temperature_sigmoid(
+            slowness_ratio, self.n_thresh_nonstationary, self.temp_coeff_nonstationary
+        )
+
+        return sig_mask
+
+    def forward(
+        self, x, xn = None
+    ) :
+        """
+        Apply the proposed algorithm to the input signal.
+
+        Arguments:
+            x (torch.Tensor): The input audio signal, with shape (batch_size, signal_length).
+            xn (Optional[torch.Tensor]): The noise signal used for stationary noise reduction. If `None`, the input
+                                         signal is used as the noise signal. Default: `None`.
+
+        Returns:
+            torch.Tensor: The denoised audio signal, with the same shape as the input signal.
+        """
+
+        # Compute short-time Fourier transform (STFT)
+        if "privateuseone" in str(x.device):
+            if not hasattr(self, "stft"):
+                self.stft = STFT(
+                    filter_length=self.n_fft,
+                    hop_length=self.hop_length,
+                    win_length=self.win_length,
+                    window="hann",
+                ).to(x.device)
+            X, phase = self.stft.transform(x, return_phase=True)
+        else:
+            X = torch.stft(
+                x,
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                return_complex=True,
+                pad_mode="constant",
+                center=True,
+                window=torch.hann_window(self.win_length).to(x.device),
+            )
+
+        # Compute signal mask based on stationary or nonstationary assumptions
+        if self.nonstationary:
+            sig_mask = self._nonstationary_mask(X.abs())
+        else:
+            sig_mask = self._stationary_mask(amp_to_db(X), xn)
+
+        # Propagate decrease in signal power
+        sig_mask = self.prop_decrease * (sig_mask.float() - 1.0) + 1.0
+
+        # Smooth signal mask with 2D convolution
+        if self.smoothing_filter is not None:
+            sig_mask = conv2d(
+                sig_mask.unsqueeze(1),
+                self.smoothing_filter.to(sig_mask.dtype),
+                padding="same",
+            )
+
+        # Apply signal mask to STFT magnitude and phase components
+        Y = X * sig_mask.squeeze(1)
+
+        # Inverse STFT to obtain time-domain signal
+        if "privateuseone" in str(Y.device):
+            y = self.stft.inverse(Y, phase)
+        else:
+            y = torch.istft(
+                Y,
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                center=True,
+                window=torch.hann_window(self.win_length).to(Y.device),
+            )
+
+        return y.to(dtype=x.dtype)

tools/torchgate/utils.py

@@ -0,0 +1,70 @@
+import torch
+from torch.types import Number
+
+
+@torch.no_grad()
+def amp_to_db(
+    x, eps=torch.finfo(torch.float64).eps, top_db=40
+) :
+    """
+    Convert the input tensor from amplitude to decibel scale.
+
+    Arguments:
+        x {[torch.Tensor]} -- [Input tensor.]
+
+    Keyword Arguments:
+        eps {[float]} -- [Small value to avoid numerical instability.]
+                          (default: {torch.finfo(torch.float64).eps})
+        top_db {[float]} -- [threshold the output at ``top_db`` below the peak]
+            `             (default: {40})
+
+    Returns:
+        [torch.Tensor] -- [Output tensor in decibel scale.]
+    """
+    x_db = 20 * torch.log10(x.abs() + eps)
+    return torch.max(x_db, (x_db.max(-1).values - top_db).unsqueeze(-1))
+
+
+@torch.no_grad()
+def temperature_sigmoid(x, x0, temp_coeff) :
+    """
+    Apply a sigmoid function with temperature scaling.
+
+    Arguments:
+        x {[torch.Tensor]} -- [Input tensor.]
+        x0 {[float]} -- [Parameter that controls the threshold of the sigmoid.]
+        temp_coeff {[float]} -- [Parameter that controls the slope of the sigmoid.]
+
+    Returns:
+        [torch.Tensor] -- [Output tensor after applying the sigmoid with temperature scaling.]
+    """
+    return torch.sigmoid((x - x0) / temp_coeff)
+
+
+@torch.no_grad()
+def linspace(
+    start, stop, num = 50, endpoint = True, **kwargs
+) :
+    """
+    Generate a linearly spaced 1-D tensor.
+
+    Arguments:
+        start {[Number]} -- [The starting value of the sequence.]
+        stop {[Number]} -- [The end value of the sequence, unless `endpoint` is set to False.
+                            In that case, the sequence consists of all but the last of ``num + 1``
+                            evenly spaced samples, so that `stop` is excluded. Note that the step
+                            size changes when `endpoint` is False.]
+
+    Keyword Arguments:
+        num {[int]} -- [Number of samples to generate. Default is 50. Must be non-negative.]
+        endpoint {[bool]} -- [If True, `stop` is the last sample. Otherwise, it is not included.
+                              Default is True.]
+        **kwargs -- [Additional arguments to be passed to the underlying PyTorch `linspace` function.]
+
+    Returns:
+        [torch.Tensor] -- [1-D tensor of `num` equally spaced samples from `start` to `stop`.]
+    """
+    if endpoint:
+        return torch.linspace(start, stop, num, **kwargs)
+    else:
+        return torch.linspace(start, stop, num + 1, **kwargs)[:-1]

tools/uvr5/bs_roformer/attend.py

@@ -0,0 +1,70 @@
+from packaging import version
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+
+def exists(val):
+    return val is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0, flash=False, scale=None):
+        super().__init__()
+        self.scale = scale
+        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"
+        )
+
+    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
+
+        if exists(self.scale):
+            default_scale = q.shape[-1] ** -0.5
+            q = q * (self.scale / default_scale)
+
+        # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale
+        # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
+        return F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout if self.training else 0.0)
+
+    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 = default(self.scale, q.shape[-1] ** -0.5)
+
+        # DirectML does not expose PyTorch's SDPA kernels.  Keep the existing
+        # SDPA path for CUDA/CPU/MPS and use the mathematically equivalent
+        # einsum implementation below for PrivateUse1 tensors.
+        if self.flash and q.device.type != "privateuseone":
+            return self.flash_attn(q, k, v)
+
+        # similarity
+
+        sim = einsum("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("b h i j, b h j d -> b h i d", attn, v)
+
+        return out

tools/uvr5/bs_roformer/bs_roformer.py

@@ -0,0 +1,347 @@
+from functools import partial
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+from tools.uvr5.bs_roformer.attend import Attend
+from torch.utils.checkpoint import checkpoint
+from typing import Tuple, Optional, Callable
+from tools.uvr5.rotary_embedding_torch import RotaryEmbedding
+from einops import rearrange, pack, unpack
+from einops.layers.torch import Rearrange
+
+def exists(val):
+    return val is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+def l2norm(t):
+    return F.normalize(t, dim=-1, p=2)
+
+class RMSNorm(Module):
+
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim=-1) * self.scale * self.gamma
+
+class FeedForward(Module):
+
+    def __init__(self, dim, mult=4, dropout=0.0):
+        super().__init__()
+        dim_inner = int(dim * mult)
+        self.net = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim_inner, dim), nn.Dropout(dropout))
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head ** (-0.5)
+        dim_inner = heads * dim_head
+        self.rotary_embed = rotary_embed
+        self.attend = Attend(flash=flash, dropout=dropout)
+        self.norm = RMSNorm(dim)
+        self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False)
+        self.to_gates = nn.Linear(dim, heads)
+        self.to_out = nn.Sequential(nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout))
+
+    def forward(self, x):
+        x = self.norm(x)
+        (q, k, v) = rearrange(self.to_qkv(x), 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.heads)
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+        out = self.attend(q, k, v)
+        gates = self.to_gates(x)
+        out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid()
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class LinearAttention(Module):
+    """
+    this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al.
+    """
+
+    def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.norm = RMSNorm(dim)
+        self.to_qkv = nn.Sequential(nn.Linear(dim, dim_inner * 3, bias=False), Rearrange('b n (qkv h d) -> qkv b h d n', qkv=3, h=heads))
+        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
+        self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
+        self.to_out = nn.Sequential(Rearrange('b h d n -> b n (h d)'), nn.Linear(dim_inner, dim, bias=False))
+
+    def forward(self, x):
+        x = self.norm(x)
+        (q, k, v) = self.to_qkv(x)
+        (q, k) = map(l2norm, (q, k))
+        q = q * self.temperature.exp()
+        out = self.attend(q, k, v)
+        return self.to_out(out)
+
+class Transformer(Module):
+
+    def __init__(self, *, dim, depth, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, ff_mult=4, norm_output=True, rotary_embed=None, flash_attn=True, linear_attn=False):
+        super().__init__()
+        self.layers = ModuleList([])
+        for _ in range(depth):
+            if linear_attn:
+                attn = LinearAttention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, flash=flash_attn)
+            else:
+                attn = Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed, flash=flash_attn)
+            self.layers.append(ModuleList([attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)]))
+        self.norm = RMSNorm(dim) if norm_output else nn.Identity()
+
+    def forward(self, x):
+        for (attn, ff) in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class BandSplit(Module):
+
+    def __init__(self, dim, dim_inputs):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_features = ModuleList([])
+        for dim_in in dim_inputs:
+            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
+            self.to_features.append(net)
+
+    def forward(self, x):
+        x = x.split(self.dim_inputs, dim=-1)
+        outs = []
+        for (split_input, to_feature) in zip(x, self.to_features):
+            split_output = to_feature(split_input)
+            outs.append(split_output)
+        return torch.stack(outs, dim=-2)
+
+def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
+    dim_hidden = default(dim_hidden, dim_in)
+    net = []
+    dims = (dim_in, *(dim_hidden,) * (depth - 1), dim_out)
+    for (ind, (layer_dim_in, layer_dim_out)) in enumerate(zip(dims[:-1], dims[1:])):
+        is_last = ind == len(dims) - 2
+        net.append(nn.Linear(layer_dim_in, layer_dim_out))
+        if is_last:
+            continue
+        net.append(activation())
+    return nn.Sequential(*net)
+
+class MaskEstimator(Module):
+
+    def __init__(self, dim, dim_inputs, depth, mlp_expansion_factor=4):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_freqs = ModuleList([])
+        dim_hidden = dim * mlp_expansion_factor
+        for dim_in in dim_inputs:
+            net = []
+            mlp = nn.Sequential(MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1))
+            self.to_freqs.append(mlp)
+
+    def forward(self, x):
+        x = x.unbind(dim=-2)
+        outs = []
+        for (band_features, mlp) in zip(x, self.to_freqs):
+            freq_out = mlp(band_features)
+            outs.append(freq_out)
+        return torch.cat(outs, dim=-1)
+DEFAULT_FREQS_PER_BANDS = (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 12, 12, 12, 12, 12, 12, 12, 12, 24, 24, 24, 24, 24, 24, 24, 24, 48, 48, 48, 48, 48, 48, 48, 48, 128, 129)
+
+class BSRoformer(Module):
+
+    def __init__(self, dim, *, depth, stereo=False, num_stems=1, time_transformer_depth=2, freq_transformer_depth=2, linear_transformer_depth=0, freqs_per_bands=DEFAULT_FREQS_PER_BANDS, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, flash_attn=True, dim_freqs_in=1025, stft_n_fft=2048, stft_hop_length=512, stft_win_length=2048, stft_normalized=False, stft_window_fn=None, mask_estimator_depth=2, multi_stft_resolution_loss_weight=1.0, multi_stft_resolutions_window_sizes=(4096, 2048, 1024, 512, 256), multi_stft_hop_size=147, multi_stft_normalized=False, multi_stft_window_fn=torch.hann_window, mlp_expansion_factor=4, use_torch_checkpoint=False, skip_connection=False):
+        super().__init__()
+        self.stereo = stereo
+        self.audio_channels = 2 if stereo else 1
+        self.num_stems = num_stems
+        self.use_torch_checkpoint = use_torch_checkpoint
+        self.skip_connection = skip_connection
+        self.layers = ModuleList([])
+        transformer_kwargs = dict(dim=dim, heads=heads, dim_head=dim_head, attn_dropout=attn_dropout, ff_dropout=ff_dropout, flash_attn=flash_attn, norm_output=False)
+        time_rotary_embed = RotaryEmbedding(dim=dim_head)
+        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
+        for _ in range(depth):
+            tran_modules = []
+            if linear_transformer_depth > 0:
+                tran_modules.append(Transformer(depth=linear_transformer_depth, linear_attn=True, **transformer_kwargs))
+            tran_modules.append(Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs))
+            tran_modules.append(Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs))
+            self.layers.append(nn.ModuleList(tran_modules))
+        self.final_norm = RMSNorm(dim)
+        self.stft_kwargs = dict(n_fft=stft_n_fft, hop_length=stft_hop_length, win_length=stft_win_length, normalized=stft_normalized)
+        self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length)
+        freqs = torch.stft(torch.randn(1, 4096), **self.stft_kwargs, window=torch.ones(stft_win_length), return_complex=True).shape[1]
+        assert len(freqs_per_bands) > 1
+        assert sum(freqs_per_bands) == freqs, f'the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}'
+        freqs_per_bands_with_complex = tuple((2 * f * self.audio_channels for f in freqs_per_bands))
+        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
+        self.mask_estimators = nn.ModuleList([])
+        for _ in range(num_stems):
+            mask_estimator = MaskEstimator(dim=dim, dim_inputs=freqs_per_bands_with_complex, depth=mask_estimator_depth, mlp_expansion_factor=mlp_expansion_factor)
+            self.mask_estimators.append(mask_estimator)
+        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
+        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
+        self.multi_stft_n_fft = stft_n_fft
+        self.multi_stft_window_fn = multi_stft_window_fn
+        self.multi_stft_kwargs = dict(hop_length=multi_stft_hop_size, normalized=multi_stft_normalized)
+
+    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
+        """
+        einops
+
+        b - batch
+        f - freq
+        t - time
+        s - audio channel (1 for mono, 2 for stereo)
+        n - number of 'stems'
+        c - complex (2)
+        d - feature dimension
+        """
+        device = raw_audio.device
+        x_is_dml = device.type == 'privateuseone'
+        x_is_mps = True if device.type == 'mps' else False
+        if raw_audio.ndim == 2:
+            raw_audio = rearrange(raw_audio, 'b t -> b 1 t')
+        channels = raw_audio.shape[1]
+        assert not self.stereo and channels == 1 or (self.stereo and channels == 2), 'stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)'
+        (raw_audio, batch_audio_channel_packed_shape) = pack_one(raw_audio, '* t')
+        if x_is_dml:
+            # DirectML has no complex/STFT kernels.  Keep only the spectral
+            # boundary on CPU and move its real representation to DirectML.
+            stft_window = self.stft_window_fn(device='cpu')
+            stft_complex = torch.stft(
+                raw_audio.cpu(),
+                **self.stft_kwargs,
+                window=stft_window,
+                return_complex=True,
+            )
+            stft_repr_cpu = torch.view_as_real(stft_complex)
+            stft_repr_cpu = unpack_one(
+                stft_repr_cpu, batch_audio_channel_packed_shape, '* f t c'
+            )
+            stft_repr_cpu = rearrange(
+                stft_repr_cpu, 'b s f t c -> b (f s) t c'
+            )
+            stft_repr = stft_repr_cpu.to(device)
+        else:
+            stft_window = self.stft_window_fn(device=device)
+            try:
+                stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
+            except:
+                stft_repr = torch.stft(raw_audio.cpu() if x_is_mps else raw_audio, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=True).to(device)
+            stft_repr = torch.view_as_real(stft_repr)
+            stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c')
+            stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
+        x = rearrange(stft_repr, 'b f t c -> b t (f c)')
+        if self.use_torch_checkpoint:
+            x = checkpoint(self.band_split, x, use_reentrant=False)
+        else:
+            x = self.band_split(x)
+        store = [None] * len(self.layers)
+        for (i, transformer_block) in enumerate(self.layers):
+            if len(transformer_block) == 3:
+                (linear_transformer, time_transformer, freq_transformer) = transformer_block
+                (x, ft_ps) = pack([x], 'b * d')
+                if self.use_torch_checkpoint:
+                    x = checkpoint(linear_transformer, x, use_reentrant=False)
+                else:
+                    x = linear_transformer(x)
+                (x,) = unpack(x, ft_ps, 'b * d')
+            else:
+                (time_transformer, freq_transformer) = transformer_block
+            if self.skip_connection:
+                for j in range(i):
+                    x = x + store[j]
+            x = rearrange(x, 'b t f d -> b f t d')
+            (x, ps) = pack([x], '* t d')
+            if self.use_torch_checkpoint:
+                x = checkpoint(time_transformer, x, use_reentrant=False)
+            else:
+                x = time_transformer(x)
+            (x,) = unpack(x, ps, '* t d')
+            x = rearrange(x, 'b f t d -> b t f d')
+            (x, ps) = pack([x], '* f d')
+            if self.use_torch_checkpoint:
+                x = checkpoint(freq_transformer, x, use_reentrant=False)
+            else:
+                x = freq_transformer(x)
+            (x,) = unpack(x, ps, '* f d')
+            if self.skip_connection:
+                store[i] = x
+        x = self.final_norm(x)
+        num_stems = len(self.mask_estimators)
+        if self.use_torch_checkpoint:
+            mask = torch.stack([checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], dim=1)
+        else:
+            mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
+        mask = rearrange(mask, 'b n t (f c) -> b n f t c', c=2)
+        if x_is_dml:
+            # Complex masking and ISTFT stay on CPU; all learned real-valued
+            # layers above remain on DirectML.
+            stft_repr = rearrange(stft_repr_cpu, 'b f t c -> b 1 f t c')
+            stft_repr = torch.view_as_complex(stft_repr.contiguous())
+            mask = torch.view_as_complex(mask.float().cpu().contiguous())
+            stft_repr = stft_repr * mask
+            stft_repr = rearrange(
+                stft_repr,
+                'b n (f s) t -> (b n s) f t',
+                s=self.audio_channels,
+            )
+            recon_audio = torch.istft(
+                stft_repr,
+                **self.stft_kwargs,
+                window=stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            ).to(device)
+        else:
+            stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c')
+            stft_repr = torch.view_as_complex(stft_repr)
+            mask = torch.view_as_complex(mask)
+            stft_repr = stft_repr * mask
+            stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels)
+            try:
+                recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, length=raw_audio.shape[-1])
+            except:
+                recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=False, length=raw_audio.shape[-1]).to(device)
+        recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', s=self.audio_channels, n=num_stems)
+        if num_stems == 1:
+            recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t')
+        if not exists(target):
+            return recon_audio
+        if self.num_stems > 1:
+            assert target.ndim == 4 and target.shape[1] == self.num_stems
+        if target.ndim == 2:
+            target = rearrange(target, '... t -> ... 1 t')
+        target = target[..., :recon_audio.shape[-1]]
+        loss_audio = recon_audio.cpu() if x_is_dml else recon_audio
+        loss_target = target.cpu() if x_is_dml else target
+        loss = F.l1_loss(loss_audio, loss_target)
+        multi_stft_resolution_loss = 0.0
+        for window_size in self.multi_stft_resolutions_window_sizes:
+            spectral_device = 'cpu' if x_is_dml else device
+            res_stft_kwargs = dict(n_fft=max(window_size, self.multi_stft_n_fft), win_length=window_size, return_complex=True, window=self.multi_stft_window_fn(window_size, device=spectral_device), **self.multi_stft_kwargs)
+            recon_Y = torch.stft(rearrange(loss_audio, '... s t -> (... s) t'), **res_stft_kwargs)
+            target_Y = torch.stft(rearrange(loss_target, '... s t -> (... s) t'), **res_stft_kwargs)
+            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(recon_Y, target_Y)
+        weighted_multi_resolution_loss = multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
+        total_loss = loss + weighted_multi_resolution_loss
+        if not return_loss_breakdown:
+            return total_loss
+        return (total_loss, (loss, multi_stft_resolution_loss))

tools/uvr5/bs_roformer/mel_band_roformer.py

@@ -0,0 +1,354 @@
+from functools import partial
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+from tools.uvr5.bs_roformer.attend import Attend
+from torch.utils.checkpoint import checkpoint
+from typing import Tuple, Optional, Callable
+from tools.uvr5.rotary_embedding_torch import RotaryEmbedding
+from einops import rearrange, pack, unpack, reduce, repeat
+from einops.layers.torch import Rearrange
+from librosa import filters
+
+def exists(val):
+    return val is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+def pad_at_dim(t, pad, dim=-1, value=0.0):
+    dims_from_right = -dim - 1 if dim < 0 else t.ndim - dim - 1
+    zeros = (0, 0) * dims_from_right
+    return F.pad(t, (*zeros, *pad), value=value)
+
+def l2norm(t):
+    return F.normalize(t, dim=-1, p=2)
+
+class RMSNorm(Module):
+
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim=-1) * self.scale * self.gamma
+
+class FeedForward(Module):
+
+    def __init__(self, dim, mult=4, dropout=0.0):
+        super().__init__()
+        dim_inner = int(dim * mult)
+        self.net = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim_inner, dim), nn.Dropout(dropout))
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head ** (-0.5)
+        dim_inner = heads * dim_head
+        self.rotary_embed = rotary_embed
+        self.attend = Attend(flash=flash, dropout=dropout)
+        self.norm = RMSNorm(dim)
+        self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False)
+        self.to_gates = nn.Linear(dim, heads)
+        self.to_out = nn.Sequential(nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout))
+
+    def forward(self, x):
+        x = self.norm(x)
+        (q, k, v) = rearrange(self.to_qkv(x), 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.heads)
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+        out = self.attend(q, k, v)
+        gates = self.to_gates(x)
+        out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid()
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class LinearAttention(Module):
+    """
+    this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al.
+    """
+
+    def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.norm = RMSNorm(dim)
+        self.to_qkv = nn.Sequential(nn.Linear(dim, dim_inner * 3, bias=False), Rearrange('b n (qkv h d) -> qkv b h d n', qkv=3, h=heads))
+        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
+        self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
+        self.to_out = nn.Sequential(Rearrange('b h d n -> b n (h d)'), nn.Linear(dim_inner, dim, bias=False))
+
+    def forward(self, x):
+        x = self.norm(x)
+        (q, k, v) = self.to_qkv(x)
+        (q, k) = map(l2norm, (q, k))
+        q = q * self.temperature.exp()
+        out = self.attend(q, k, v)
+        return self.to_out(out)
+
+class Transformer(Module):
+
+    def __init__(self, *, dim, depth, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, ff_mult=4, norm_output=True, rotary_embed=None, flash_attn=True, linear_attn=False):
+        super().__init__()
+        self.layers = ModuleList([])
+        for _ in range(depth):
+            if linear_attn:
+                attn = LinearAttention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, flash=flash_attn)
+            else:
+                attn = Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed, flash=flash_attn)
+            self.layers.append(ModuleList([attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)]))
+        self.norm = RMSNorm(dim) if norm_output else nn.Identity()
+
+    def forward(self, x):
+        for (attn, ff) in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class BandSplit(Module):
+
+    def __init__(self, dim, dim_inputs):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_features = ModuleList([])
+        for dim_in in dim_inputs:
+            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
+            self.to_features.append(net)
+
+    def forward(self, x):
+        x = x.split(self.dim_inputs, dim=-1)
+        outs = []
+        for (split_input, to_feature) in zip(x, self.to_features):
+            split_output = to_feature(split_input)
+            outs.append(split_output)
+        return torch.stack(outs, dim=-2)
+
+def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
+    dim_hidden = default(dim_hidden, dim_in)
+    net = []
+    dims = (dim_in, *(dim_hidden,) * depth, dim_out)
+    for (ind, (layer_dim_in, layer_dim_out)) in enumerate(zip(dims[:-1], dims[1:])):
+        is_last = ind == len(dims) - 2
+        net.append(nn.Linear(layer_dim_in, layer_dim_out))
+        if is_last:
+            continue
+        net.append(activation())
+    return nn.Sequential(*net)
+
+class MaskEstimator(Module):
+
+    def __init__(self, dim, dim_inputs, depth, mlp_expansion_factor=4):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_freqs = ModuleList([])
+        dim_hidden = dim * mlp_expansion_factor
+        for dim_in in dim_inputs:
+            net = []
+            mlp = nn.Sequential(MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1))
+            self.to_freqs.append(mlp)
+
+    def forward(self, x):
+        x = x.unbind(dim=-2)
+        outs = []
+        for (band_features, mlp) in zip(x, self.to_freqs):
+            freq_out = mlp(band_features)
+            outs.append(freq_out)
+        return torch.cat(outs, dim=-1)
+
+class MelBandRoformer(Module):
+
+    def __init__(self, dim, *, depth, stereo=False, num_stems=1, time_transformer_depth=2, freq_transformer_depth=2, linear_transformer_depth=0, num_bands=60, dim_head=64, heads=8, attn_dropout=0.1, ff_dropout=0.1, flash_attn=True, dim_freqs_in=1025, sample_rate=44100, stft_n_fft=2048, stft_hop_length=512, stft_win_length=2048, stft_normalized=False, stft_window_fn=None, mask_estimator_depth=1, multi_stft_resolution_loss_weight=1.0, multi_stft_resolutions_window_sizes=(4096, 2048, 1024, 512, 256), multi_stft_hop_size=147, multi_stft_normalized=False, multi_stft_window_fn=torch.hann_window, match_input_audio_length=False, mlp_expansion_factor=4, use_torch_checkpoint=False, skip_connection=False):
+        super().__init__()
+        self.stereo = stereo
+        self.audio_channels = 2 if stereo else 1
+        self.num_stems = num_stems
+        self.use_torch_checkpoint = use_torch_checkpoint
+        self.skip_connection = skip_connection
+        self.layers = ModuleList([])
+        transformer_kwargs = dict(dim=dim, heads=heads, dim_head=dim_head, attn_dropout=attn_dropout, ff_dropout=ff_dropout, flash_attn=flash_attn)
+        time_rotary_embed = RotaryEmbedding(dim=dim_head)
+        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
+        for _ in range(depth):
+            tran_modules = []
+            if linear_transformer_depth > 0:
+                tran_modules.append(Transformer(depth=linear_transformer_depth, linear_attn=True, **transformer_kwargs))
+            tran_modules.append(Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs))
+            tran_modules.append(Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs))
+            self.layers.append(nn.ModuleList(tran_modules))
+        self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length)
+        self.stft_kwargs = dict(n_fft=stft_n_fft, hop_length=stft_hop_length, win_length=stft_win_length, normalized=stft_normalized)
+        freqs = torch.stft(torch.randn(1, 4096), **self.stft_kwargs, window=torch.ones(stft_n_fft), return_complex=True).shape[1]
+        mel_filter_bank_numpy = filters.mel(sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands)
+        mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy)
+        mel_filter_bank[0][0] = 1.0
+        mel_filter_bank[-1, -1] = 1.0
+        freqs_per_band = mel_filter_bank > 0
+        assert freqs_per_band.any(dim=0).all(), 'all frequencies need to be covered by all bands for now'
+        repeated_freq_indices = repeat(torch.arange(freqs), 'f -> b f', b=num_bands)
+        freq_indices = repeated_freq_indices[freqs_per_band]
+        if stereo:
+            freq_indices = repeat(freq_indices, 'f -> f s', s=2)
+            freq_indices = freq_indices * 2 + torch.arange(2)
+            freq_indices = rearrange(freq_indices, 'f s -> (f s)')
+        self.register_buffer('freq_indices', freq_indices, persistent=False)
+        self.register_buffer('freqs_per_band', freqs_per_band, persistent=False)
+        num_freqs_per_band = reduce(freqs_per_band, 'b f -> b', 'sum')
+        num_bands_per_freq = reduce(freqs_per_band, 'b f -> f', 'sum')
+        self.register_buffer('num_freqs_per_band', num_freqs_per_band, persistent=False)
+        self.register_buffer('num_bands_per_freq', num_bands_per_freq, persistent=False)
+        freqs_per_bands_with_complex = tuple((2 * f * self.audio_channels for f in num_freqs_per_band.tolist()))
+        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
+        self.mask_estimators = nn.ModuleList([])
+        for _ in range(num_stems):
+            mask_estimator = MaskEstimator(dim=dim, dim_inputs=freqs_per_bands_with_complex, depth=mask_estimator_depth, mlp_expansion_factor=mlp_expansion_factor)
+            self.mask_estimators.append(mask_estimator)
+        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
+        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
+        self.multi_stft_n_fft = stft_n_fft
+        self.multi_stft_window_fn = multi_stft_window_fn
+        self.multi_stft_kwargs = dict(hop_length=multi_stft_hop_size, normalized=multi_stft_normalized)
+        self.match_input_audio_length = match_input_audio_length
+
+    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
+        """
+        einops
+
+        b - batch
+        f - freq
+        t - time
+        s - audio channel (1 for mono, 2 for stereo)
+        n - number of 'stems'
+        c - complex (2)
+        d - feature dimension
+        """
+        device = raw_audio.device
+        x_is_dml = device.type == 'privateuseone'
+        if raw_audio.ndim == 2:
+            raw_audio = rearrange(raw_audio, 'b t -> b 1 t')
+        (batch, channels, raw_audio_length) = raw_audio.shape
+        istft_length = raw_audio_length if self.match_input_audio_length else None
+        assert not self.stereo and channels == 1 or (self.stereo and channels == 2), 'stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)'
+        (raw_audio, batch_audio_channel_packed_shape) = pack_one(raw_audio, '* t')
+        if x_is_dml:
+            # DirectML has no STFT or complex tensor support.  Build the real
+            # spectral features on CPU, then run the learned network on DML.
+            stft_window = self.stft_window_fn(device='cpu')
+            stft_complex = torch.stft(
+                raw_audio.cpu(),
+                **self.stft_kwargs,
+                window=stft_window,
+                return_complex=True,
+            )
+            stft_repr = torch.view_as_real(stft_complex)
+            stft_repr = unpack_one(
+                stft_repr, batch_audio_channel_packed_shape, '* f t c'
+            )
+            stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
+            batch_arange = torch.arange(batch)[..., None]
+            x = stft_repr[batch_arange, self.freq_indices.cpu()].to(device)
+        else:
+            stft_window = self.stft_window_fn(device=device)
+            stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
+            stft_repr = torch.view_as_real(stft_repr)
+            stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c')
+            stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
+            batch_arange = torch.arange(batch, device=device)[..., None]
+            x = stft_repr[batch_arange, self.freq_indices]
+        x = rearrange(x, 'b f t c -> b t (f c)')
+        if self.use_torch_checkpoint:
+            x = checkpoint(self.band_split, x, use_reentrant=False)
+        else:
+            x = self.band_split(x)
+        store = [None] * len(self.layers)
+        for (i, transformer_block) in enumerate(self.layers):
+            if len(transformer_block) == 3:
+                (linear_transformer, time_transformer, freq_transformer) = transformer_block
+                (x, ft_ps) = pack([x], 'b * d')
+                if self.use_torch_checkpoint:
+                    x = checkpoint(linear_transformer, x, use_reentrant=False)
+                else:
+                    x = linear_transformer(x)
+                (x,) = unpack(x, ft_ps, 'b * d')
+            else:
+                (time_transformer, freq_transformer) = transformer_block
+            if self.skip_connection:
+                for j in range(i):
+                    x = x + store[j]
+            x = rearrange(x, 'b t f d -> b f t d')
+            (x, ps) = pack([x], '* t d')
+            if self.use_torch_checkpoint:
+                x = checkpoint(time_transformer, x, use_reentrant=False)
+            else:
+                x = time_transformer(x)
+            (x,) = unpack(x, ps, '* t d')
+            x = rearrange(x, 'b f t d -> b t f d')
+            (x, ps) = pack([x], '* f d')
+            if self.use_torch_checkpoint:
+                x = checkpoint(freq_transformer, x, use_reentrant=False)
+            else:
+                x = freq_transformer(x)
+            (x,) = unpack(x, ps, '* f d')
+            if self.skip_connection:
+                store[i] = x
+        num_stems = len(self.mask_estimators)
+        if self.use_torch_checkpoint:
+            masks = torch.stack([checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], dim=1)
+        else:
+            masks = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
+        masks = rearrange(masks, 'b n t (f c) -> b n f t c', c=2)
+        if x_is_dml:
+            masks = masks.float().cpu()
+        stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c')
+        stft_repr = torch.view_as_complex(stft_repr.contiguous())
+        masks = torch.view_as_complex(masks.contiguous())
+        masks = masks.type(stft_repr.dtype)
+        freq_indices = self.freq_indices.cpu() if x_is_dml else self.freq_indices
+        scatter_indices = repeat(freq_indices, 'f -> b n f t', b=batch, n=num_stems, t=stft_repr.shape[-1])
+        stft_repr_expanded_stems = repeat(stft_repr, 'b 1 ... -> b n ...', n=num_stems)
+        masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_(2, scatter_indices, masks)
+        num_bands_per_freq = self.num_bands_per_freq.cpu() if x_is_dml else self.num_bands_per_freq
+        denom = repeat(num_bands_per_freq, 'f -> (f r) 1', r=channels)
+        masks_averaged = masks_summed / denom.clamp(min=1e-08)
+        stft_repr = stft_repr * masks_averaged
+        stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels)
+        recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, length=istft_length)
+        if x_is_dml:
+            recon_audio = recon_audio.to(device)
+        recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', b=batch, s=self.audio_channels, n=num_stems)
+        if num_stems == 1:
+            recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t')
+        if not exists(target):
+            return recon_audio
+        if self.num_stems > 1:
+            assert target.ndim == 4 and target.shape[1] == self.num_stems
+        if target.ndim == 2:
+            target = rearrange(target, '... t -> ... 1 t')
+        target = target[..., :recon_audio.shape[-1]]
+        loss_audio = recon_audio.cpu() if x_is_dml else recon_audio
+        loss_target = target.cpu() if x_is_dml else target
+        loss = F.l1_loss(loss_audio, loss_target)
+        multi_stft_resolution_loss = 0.0
+        for window_size in self.multi_stft_resolutions_window_sizes:
+            spectral_device = 'cpu' if x_is_dml else device
+            res_stft_kwargs = dict(n_fft=max(window_size, self.multi_stft_n_fft), win_length=window_size, return_complex=True, window=self.multi_stft_window_fn(window_size, device=spectral_device), **self.multi_stft_kwargs)
+            recon_Y = torch.stft(rearrange(loss_audio, '... s t -> (... s) t'), **res_stft_kwargs)
+            target_Y = torch.stft(rearrange(loss_target, '... s t -> (... s) t'), **res_stft_kwargs)
+            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(recon_Y, target_Y)
+        weighted_multi_resolution_loss = multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
+        total_loss = loss + weighted_multi_resolution_loss
+        if not return_loss_breakdown:
+            return total_loss
+        return (total_loss, (loss, multi_stft_resolution_loss))

tools/uvr5/bsroformer.py

@@ -0,0 +1,317 @@
+# This code is modified from https://github.com/ZFTurbo/
+import os
+import warnings
+from contextlib import nullcontext
+
+import librosa
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn as nn
+import yaml
+from tqdm import tqdm
+from tools.file_io import read_text
+from i18n.i18n import I18nAuto
+
+warnings.filterwarnings("ignore")
+i18n = I18nAuto()
+
+
+class Roformer_Loader:
+    def get_config(self, config_path):
+        return yaml.load(read_text(config_path), Loader=yaml.FullLoader)
+
+    def get_default_config(self):
+        default_config = None
+        if self.model_type == "bs_roformer":
+            # Use model_bs_roformer_ep_368_sdr_12.9628.yaml and model_bs_roformer_ep_317_sdr_12.9755.yaml as default configuration files
+            # Other BS_Roformer models may not be compatible
+            # fmt: off
+            default_config = {
+                "audio": {"chunk_size": 352800, "sample_rate": 44100},
+                "model": {
+                    "dim": 512,
+                    "depth": 12,
+                    "stereo": True,
+                    "num_stems": 1,
+                    "time_transformer_depth": 1,
+                    "freq_transformer_depth": 1,
+                    "linear_transformer_depth": 0,
+                    "freqs_per_bands": (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 12, 12, 12, 12, 12, 12, 12, 12, 24, 24, 24, 24, 24, 24, 24, 24, 48, 48, 48, 48, 48, 48, 48, 48, 128, 129),
+                    "dim_head": 64,
+                    "heads": 8,
+                    "attn_dropout": 0.1,
+                    "ff_dropout": 0.1,
+                    "flash_attn": True,
+                    "dim_freqs_in": 1025,
+                    "stft_n_fft": 2048,
+                    "stft_hop_length": 441,
+                    "stft_win_length": 2048,
+                    "stft_normalized": False,
+                    "mask_estimator_depth": 2,
+                    "multi_stft_resolution_loss_weight": 1.0,
+                    "multi_stft_resolutions_window_sizes": (4096, 2048, 1024, 512, 256),
+                    "multi_stft_hop_size": 147,
+                    "multi_stft_normalized": False,
+                },
+                "training": {"instruments": ["vocals", "other"], "target_instrument": "vocals"},
+                "inference": {"batch_size": 2, "num_overlap": 2},
+            }
+            # fmt: on
+        elif self.model_type == "mel_band_roformer":
+            # Use model_mel_band_roformer_ep_3005_sdr_11.4360.yaml as default configuration files
+            # Other Mel_Band_Roformer models may not be compatible
+            default_config = {
+                "audio": {"chunk_size": 352800, "sample_rate": 44100},
+                "model": {
+                    "dim": 384,
+                    "depth": 12,
+                    "stereo": True,
+                    "num_stems": 1,
+                    "time_transformer_depth": 1,
+                    "freq_transformer_depth": 1,
+                    "linear_transformer_depth": 0,
+                    "num_bands": 60,
+                    "dim_head": 64,
+                    "heads": 8,
+                    "attn_dropout": 0.1,
+                    "ff_dropout": 0.1,
+                    "flash_attn": True,
+                    "dim_freqs_in": 1025,
+                    "sample_rate": 44100,
+                    "stft_n_fft": 2048,
+                    "stft_hop_length": 441,
+                    "stft_win_length": 2048,
+                    "stft_normalized": False,
+                    "mask_estimator_depth": 2,
+                    "multi_stft_resolution_loss_weight": 1.0,
+                    "multi_stft_resolutions_window_sizes": (4096, 2048, 1024, 512, 256),
+                    "multi_stft_hop_size": 147,
+                    "multi_stft_normalized": False,
+                },
+                "training": {"instruments": ["vocals", "other"], "target_instrument": "vocals"},
+                "inference": {"batch_size": 2, "num_overlap": 2},
+            }
+
+        return default_config
+
+    def get_model_from_config(self):
+        if self.model_type == "bs_roformer":
+            from tools.uvr5.bs_roformer.bs_roformer import BSRoformer
+
+            model = BSRoformer(**dict(self.config["model"]))
+        elif self.model_type == "mel_band_roformer":
+            from tools.uvr5.bs_roformer.mel_band_roformer import MelBandRoformer
+
+            model = MelBandRoformer(**dict(self.config["model"]))
+        else:
+            print(i18n("错误：未知模型：%s") % self.model_type)
+            model = None
+        return model
+
+    def demix_track(self, model, mix, device):
+        C = self.config["audio"]["chunk_size"]  # chunk_size
+        N = self.config["inference"]["num_overlap"]
+        fade_size = C // 10
+        step = int(C // N)
+        border = C - step
+        batch_size = self.config["inference"]["batch_size"]
+
+        length_init = mix.shape[-1]
+        progress_bar = tqdm(total=length_init // step + 1, desc="Processing", leave=False)
+
+        # Do pad from the beginning and end to account floating window results better
+        if length_init > 2 * border and (border > 0):
+            mix = nn.functional.pad(mix, (border, border), mode="reflect")
+
+        # Prepare windows arrays (do 1 time for speed up). This trick repairs click problems on the edges of segment
+        window_size = C
+        fadein = torch.linspace(0, 1, fade_size)
+        fadeout = torch.linspace(1, 0, fade_size)
+        window_start = torch.ones(window_size)
+        window_middle = torch.ones(window_size)
+        window_finish = torch.ones(window_size)
+        window_start[-fade_size:] *= fadeout  # First audio chunk, no fadein
+        window_finish[:fade_size] *= fadein  # Last audio chunk, no fadeout
+        window_middle[-fade_size:] *= fadeout
+        window_middle[:fade_size] *= fadein
+
+        device_type = device.type if isinstance(device, torch.device) else torch.device(device).type
+        amp_context = (
+            torch.amp.autocast("cuda") if device_type == "cuda" else nullcontext()
+        )
+        grad_context = (
+            torch.no_grad()
+            if device_type == "privateuseone"
+            else torch.inference_mode()
+        )
+        with amp_context:
+            # DirectML updates version counters in several linear kernels and
+            # therefore needs no_grad rather than inference_mode. CUDA and CPU
+            # retain the existing inference-mode path.
+            with grad_context:
+                if self.config["training"]["target_instrument"] is None:
+                    req_shape = (len(self.config["training"]["instruments"]),) + tuple(mix.shape)
+                else:
+                    req_shape = (1,) + tuple(mix.shape)
+
+                result = torch.zeros(req_shape, dtype=torch.float32)
+                counter = torch.zeros(req_shape, dtype=torch.float32)
+                i = 0
+                batch_data = []
+                batch_locations = []
+                while i < mix.shape[1]:
+                    part = mix[:, i : i + C].to(device)
+                    length = part.shape[-1]
+                    if length < C:
+                        if length > C // 2 + 1:
+                            part = nn.functional.pad(input=part, pad=(0, C - length), mode="reflect")
+                        else:
+                            part = nn.functional.pad(input=part, pad=(0, C - length, 0, 0), mode="constant", value=0)
+                    if self.is_half:
+                        part = part.half()
+                    batch_data.append(part)
+                    batch_locations.append((i, length))
+                    i += step
+                    progress_bar.update(1)
+
+                    if len(batch_data) >= batch_size or (i >= mix.shape[1]):
+                        arr = torch.stack(batch_data, dim=0)
+                        # print(23333333,arr.dtype)
+                        x = model(arr)
+
+                        window = window_middle
+                        if i - step == 0:  # First audio chunk, no fadein
+                            window = window_start
+                        elif i >= mix.shape[1]:  # Last audio chunk, no fadeout
+                            window = window_finish
+
+                        for j in range(len(batch_locations)):
+                            start, l = batch_locations[j]
+                            result[..., start : start + l] += x[j][..., :l].cpu() * window[..., :l]
+                            counter[..., start : start + l] += window[..., :l]
+
+                        batch_data = []
+                        batch_locations = []
+
+                estimated_sources = result / counter
+                estimated_sources = estimated_sources.cpu().numpy()
+                np.nan_to_num(estimated_sources, copy=False, nan=0.0)
+
+                if length_init > 2 * border and (border > 0):
+                    # Remove pad
+                    estimated_sources = estimated_sources[..., border:-border]
+
+        progress_bar.close()
+
+        if self.config["training"]["target_instrument"] is None:
+            return {k: v for k, v in zip(self.config["training"]["instruments"], estimated_sources)}
+        else:
+            return {k: v for k, v in zip([self.config["training"]["target_instrument"]], estimated_sources)}
+
+    def run_folder(self, input, vocal_root, others_root, format):
+        self.model.eval()
+        path = input
+        os.makedirs(vocal_root, exist_ok=True)
+        os.makedirs(others_root, exist_ok=True)
+        file_base_name = os.path.splitext(os.path.basename(path))[0]
+
+        sample_rate = 44100
+        if "sample_rate" in self.config["audio"]:
+            sample_rate = self.config["audio"]["sample_rate"]
+
+        try:
+            mix, sr = librosa.load(path, sr=sample_rate, mono=False)
+        except Exception as e:
+            print(i18n("无法读取音频：%s") % path)
+            print(i18n("错误信息：%s") % str(e))
+            return
+
+        # in case if model only supports mono tracks
+        isstereo = self.config["model"].get("stereo", True)
+        if not isstereo and len(mix.shape) != 1:
+            mix = np.mean(mix, axis=0)  # if more than 2 channels, take mean
+            print(i18n("音频包含多个声道，但模型仅支持单声道，将对所有声道取平均值"))
+
+        mix_orig = mix.copy()
+
+        mixture = torch.tensor(mix, dtype=torch.float32)
+        res = self.demix_track(self.model, mixture, self.device)
+
+        if self.config["training"]["target_instrument"] is not None:
+            # if target instrument is specified, save target instrument as vocal and other instruments as others
+            # other instruments are caculated by subtracting target instrument from mixture
+            target_instrument = self.config["training"]["target_instrument"]
+            other_instruments = [i for i in self.config["training"]["instruments"] if i != target_instrument]
+            other = mix_orig - res[target_instrument]  # caculate other instruments
+
+            path_vocal = "{}/{}_{}.wav".format(vocal_root, file_base_name, target_instrument)
+            path_other = "{}/{}_{}.wav".format(others_root, file_base_name, other_instruments[0])
+            self.save_audio(path_vocal, res[target_instrument].T, sr, format)
+            self.save_audio(path_other, other.T, sr, format)
+        else:
+            # if target instrument is not specified, save the first instrument as vocal and the rest as others
+            vocal_inst = self.config["training"]["instruments"][0]
+            path_vocal = "{}/{}_{}.wav".format(vocal_root, file_base_name, vocal_inst)
+            self.save_audio(path_vocal, res[vocal_inst].T, sr, format)
+            for other in self.config["training"]["instruments"][1:]:  # save other instruments
+                path_other = "{}/{}_{}.wav".format(others_root, file_base_name, other)
+                self.save_audio(path_other, res[other].T, sr, format)
+
+    def save_audio(self, path, data, sr, format):
+        # input path should be endwith '.wav'
+        if format in ["wav", "flac"]:
+            if format == "flac":
+                path = path[:-3] + "flac"
+            sf.write(path, data, sr)
+        else:
+            sf.write(path, data, sr)
+            os.system('ffmpeg -i "{}" -vn "{}" -q:a 2 -y'.format(path, path[:-3] + format))
+            try:
+                os.remove(path)
+            except:
+                pass
+
+    def __init__(self, model_path, config_path, device, is_half):
+        self.device = device
+        self.is_half = is_half
+        self.model_type = None
+        self.config = None
+
+        # get model_type, first try:
+        if "bs_roformer" in model_path.lower() or "bsroformer" in model_path.lower():
+            self.model_type = "bs_roformer"
+        elif "mel_band_roformer" in model_path.lower() or "melbandroformer" in model_path.lower():
+            self.model_type = "mel_band_roformer"
+
+        if not os.path.exists(config_path):
+            if self.model_type is None:
+                # if model_type is still None, raise an error
+                raise ValueError(
+                    "Error: Unknown model type. If you are using a model without a configuration file, Ensure that your model name includes 'bs_roformer', 'bsroformer', 'mel_band_roformer', or 'melbandroformer'. Otherwise, you can manually place the model configuration file into 'tools/uvr5/uvr5w_weights' and ensure that the configuration file is named as '<model_name>.yaml' then try it again."
+                )
+            self.config = self.get_default_config()
+        else:
+            # if there is a configuration file
+            self.config = self.get_config(config_path)
+            if self.model_type is None:
+                # if model_type is still None, second try, get model_type from the configuration file
+                if "freqs_per_bands" in self.config["model"]:
+                    # if freqs_per_bands in config, it's a bs_roformer model
+                    self.model_type = "bs_roformer"
+                else:
+                    # else it's a mel_band_roformer model
+                    self.model_type = "mel_band_roformer"
+
+        print(i18n("检测到模型类型：%s") % self.model_type)
+        model = self.get_model_from_config()
+        state_dict = torch.load(model_path, map_location="cpu")
+        model.load_state_dict(state_dict)
+
+        if is_half == False:
+            self.model = model.to(device)
+        else:
+            self.model = model.half().to(device)
+
+    def _path_audio_(self, input, others_root, vocal_root, format, is_hp3=False):
+        self.run_folder(input, vocal_root, others_root, format)

tools/uvr5/lib/lib_v5/layers_123821KB.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from . import spec_utils
+
+
+class Conv2DBNActiv(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
+        super(Conv2DBNActiv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                nin,
+                nout,
+                kernel_size=ksize,
+                stride=stride,
+                padding=pad,
+                dilation=dilation,
+                bias=False,
+            ),
+            nn.BatchNorm2d(nout),
+            activ(),
+        )
+
+    def __call__(self, x):
+        return self.conv(x)
+
+
+class SeperableConv2DBNActiv(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
+        super(SeperableConv2DBNActiv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                nin,
+                nin,
+                kernel_size=ksize,
+                stride=stride,
+                padding=pad,
+                dilation=dilation,
+                groups=nin,
+                bias=False,
+            ),
+            nn.Conv2d(nin, nout, kernel_size=1, bias=False),
+            nn.BatchNorm2d(nout),
+            activ(),
+        )
+
+    def __call__(self, x):
+        return self.conv(x)
+
+
+class Encoder(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
+        super(Encoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
+
+    def __call__(self, x):
+        skip = self.conv1(x)
+        h = self.conv2(skip)
+
+        return h, skip
+
+
+class Decoder(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
+        super(Decoder, self).__init__()
+        self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def __call__(self, x, skip=None):
+        x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
+        if skip is not None:
+            skip = spec_utils.crop_center(skip, x)
+            x = torch.cat([x, skip], dim=1)
+        h = self.conv(x)
+
+        if self.dropout is not None:
+            h = self.dropout(h)
+
+        return h
+
+
+class ASPPModule(nn.Module):
+    def __init__(self, nin, nout, dilations=(4, 8, 16), activ=nn.ReLU):
+        super(ASPPModule, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, None)),
+            Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
+        )
+        self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
+        self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ)
+        self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ)
+        self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
+        self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 5, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1))
+
+    def forward(self, x):
+        _, _, h, w = x.size()
+        feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
+        feat2 = self.conv2(x)
+        feat3 = self.conv3(x)
+        feat4 = self.conv4(x)
+        feat5 = self.conv5(x)
+        out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        bottle = self.bottleneck(out)
+        return bottle

tools/uvr5/lib/lib_v5/layers_new.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from . import spec_utils
+
+
+class Conv2DBNActiv(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
+        super(Conv2DBNActiv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                nin,
+                nout,
+                kernel_size=ksize,
+                stride=stride,
+                padding=pad,
+                dilation=dilation,
+                bias=False,
+            ),
+            nn.BatchNorm2d(nout),
+            activ(),
+        )
+
+    def __call__(self, x):
+        return self.conv(x)
+
+
+class Encoder(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
+        super(Encoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ)
+        self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
+
+    def __call__(self, x):
+        h = self.conv1(x)
+        h = self.conv2(h)
+
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
+        super(Decoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        # self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def __call__(self, x, skip=None):
+        x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
+
+        if skip is not None:
+            skip = spec_utils.crop_center(skip, x)
+            x = torch.cat([x, skip], dim=1)
+
+        h = self.conv1(x)
+        # h = self.conv2(h)
+
+        if self.dropout is not None:
+            h = self.dropout(h)
+
+        return h
+
+
+class ASPPModule(nn.Module):
+    def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False):
+        super(ASPPModule, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, None)),
+            Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ),
+        )
+        self.conv2 = Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
+        self.conv3 = Conv2DBNActiv(nin, nout, 3, 1, dilations[0], dilations[0], activ=activ)
+        self.conv4 = Conv2DBNActiv(nin, nout, 3, 1, dilations[1], dilations[1], activ=activ)
+        self.conv5 = Conv2DBNActiv(nin, nout, 3, 1, dilations[2], dilations[2], activ=activ)
+        self.bottleneck = Conv2DBNActiv(nout * 5, nout, 1, 1, 0, activ=activ)
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def forward(self, x):
+        _, _, h, w = x.size()
+        feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
+        feat2 = self.conv2(x)
+        feat3 = self.conv3(x)
+        feat4 = self.conv4(x)
+        feat5 = self.conv5(x)
+        out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        out = self.bottleneck(out)
+
+        if self.dropout is not None:
+            out = self.dropout(out)
+
+        return out
+
+
+class LSTMModule(nn.Module):
+    def __init__(self, nin_conv, nin_lstm, nout_lstm):
+        super(LSTMModule, self).__init__()
+        self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
+        self.lstm = nn.LSTM(input_size=nin_lstm, hidden_size=nout_lstm // 2, bidirectional=True)
+        self.dense = nn.Sequential(nn.Linear(nout_lstm, nin_lstm), nn.BatchNorm1d(nin_lstm), nn.ReLU())
+
+    def forward(self, x):
+        N, _, nbins, nframes = x.size()
+        h = self.conv(x)[:, 0]  # N, nbins, nframes
+        h = h.permute(2, 0, 1)  # nframes, N, nbins
+        h, _ = self.lstm(h)
+        h = self.dense(h.reshape(-1, h.size()[-1]))  # nframes * N, nbins
+        h = h.reshape(nframes, N, 1, nbins)
+        h = h.permute(1, 2, 3, 0)
+
+        return h

tools/uvr5/lib/lib_v5/model_param_init.py

@@ -0,0 +1,68 @@
+import json
+import pathlib
+from tools.file_io import read_text
+
+default_param = {}
+default_param["bins"] = 768
+default_param["unstable_bins"] = 9  # training only
+default_param["reduction_bins"] = 762  # training only
+default_param["sr"] = 44100
+default_param["pre_filter_start"] = 757
+default_param["pre_filter_stop"] = 768
+default_param["band"] = {}
+
+
+default_param["band"][1] = {
+    "sr": 11025,
+    "hl": 128,
+    "n_fft": 960,
+    "crop_start": 0,
+    "crop_stop": 245,
+    "lpf_start": 61,  # inference only
+    "res_type": "polyphase",
+}
+
+default_param["band"][2] = {
+    "sr": 44100,
+    "hl": 512,
+    "n_fft": 1536,
+    "crop_start": 24,
+    "crop_stop": 547,
+    "hpf_start": 81,  # inference only
+    "res_type": "sinc_best",
+}
+
+
+def int_keys(d):
+    r = {}
+    for k, v in d:
+        if k.isdigit():
+            k = int(k)
+        r[k] = v
+    return r
+
+
+class ModelParameters(object):
+    def __init__(self, config_path=""):
+        if ".pth" == pathlib.Path(config_path).suffix:
+            import zipfile
+
+            with zipfile.ZipFile(config_path, "r") as zip:
+                self.param = json.loads(zip.read("param.json"), object_pairs_hook=int_keys)
+        elif ".json" == pathlib.Path(config_path).suffix:
+            self.param = json.loads(
+                read_text(config_path), object_pairs_hook=int_keys
+            )
+        else:
+            self.param = default_param
+
+        for k in [
+            "mid_side",
+            "mid_side_b",
+            "mid_side_b2",
+            "stereo_w",
+            "stereo_n",
+            "reverse",
+        ]:
+            if k not in self.param:
+                self.param[k] = False

tools/uvr5/lib/lib_v5/modelparams/4band_v2.json

@@ -0,0 +1,54 @@
+{
+	"bins": 672,
+	"unstable_bins": 8,
+	"reduction_bins": 637,
+	"band": {
+		"1": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 640,
+			"crop_start": 0,
+			"crop_stop": 85,
+			"lpf_start": 25,
+			"lpf_stop": 53,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 320,
+			"crop_start": 4,
+			"crop_stop": 87,
+			"hpf_start": 25,
+			"hpf_stop": 12,
+			"lpf_start": 31,
+			"lpf_stop": 62,
+			"res_type": "polyphase"
+		},		
+		"3": {
+			"sr": 14700,
+			"hl": 160,
+			"n_fft": 512,
+			"crop_start": 17,
+			"crop_stop": 216,
+			"hpf_start": 48,
+			"hpf_stop": 24,
+			"lpf_start": 139,
+			"lpf_stop": 210,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 480,
+			"n_fft": 960,
+			"crop_start": 78,
+			"crop_stop": 383,
+			"hpf_start": 130,
+			"hpf_stop": 86,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 668,
+	"pre_filter_stop": 672
+}

tools/uvr5/lib/lib_v5/modelparams/4band_v3.json

@@ -0,0 +1,54 @@
+{
+	"bins": 672,
+	"unstable_bins": 8,
+	"reduction_bins": 530,
+	"band": {
+		"1": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 640,
+			"crop_start": 0,
+			"crop_stop": 85,
+			"lpf_start": 25,
+			"lpf_stop": 53,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 320,
+			"crop_start": 4,
+			"crop_stop": 87,
+			"hpf_start": 25,
+			"hpf_stop": 12,
+			"lpf_start": 31,
+			"lpf_stop": 62,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 14700,
+			"hl": 160,
+			"n_fft": 512,
+			"crop_start": 17,
+			"crop_stop": 216,
+			"hpf_start": 48,
+			"hpf_stop": 24,
+			"lpf_start": 139,
+			"lpf_stop": 210,
+			"res_type": "polyphase"
+		},
+		"4": {
+			"sr": 44100,
+			"hl": 480,
+			"n_fft": 960,
+			"crop_start": 78,
+			"crop_stop": 383,
+			"hpf_start": 130,
+			"hpf_stop": 86,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 668,
+	"pre_filter_stop": 672
+}

tools/uvr5/lib/lib_v5/nets_61968KB.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from . import layers_123821KB as layers
+
+
+class BaseASPPNet(nn.Module):
+    def __init__(self, nin, ch, dilations=(4, 8, 16)):
+        super(BaseASPPNet, self).__init__()
+        self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
+        self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
+        self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
+        self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
+
+        self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
+
+        self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
+        self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
+        self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
+        self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
+
+    def __call__(self, x):
+        h, e1 = self.enc1(x)
+        h, e2 = self.enc2(h)
+        h, e3 = self.enc3(h)
+        h, e4 = self.enc4(h)
+
+        h = self.aspp(h)
+
+        h = self.dec4(h, e4)
+        h = self.dec3(h, e3)
+        h = self.dec2(h, e2)
+        h = self.dec1(h, e1)
+
+        return h
+
+
+class CascadedASPPNet(nn.Module):
+    def __init__(self, n_fft):
+        super(CascadedASPPNet, self).__init__()
+        self.stg1_low_band_net = BaseASPPNet(2, 32)
+        self.stg1_high_band_net = BaseASPPNet(2, 32)
+
+        self.stg2_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0)
+        self.stg2_full_band_net = BaseASPPNet(16, 32)
+
+        self.stg3_bridge = layers.Conv2DBNActiv(66, 32, 1, 1, 0)
+        self.stg3_full_band_net = BaseASPPNet(32, 64)
+
+        self.out = nn.Conv2d(64, 2, 1, bias=False)
+        self.aux1_out = nn.Conv2d(32, 2, 1, bias=False)
+        self.aux2_out = nn.Conv2d(32, 2, 1, bias=False)
+
+        self.max_bin = n_fft // 2
+        self.output_bin = n_fft // 2 + 1
+
+        self.offset = 128
+
+    def forward(self, x, aggressiveness=None):
+        mix = x.detach()
+        x = x.clone()
+
+        x = x[:, :, : self.max_bin]
+
+        bandw = x.size()[2] // 2
+        aux1 = torch.cat(
+            [
+                self.stg1_low_band_net(x[:, :, :bandw]),
+                self.stg1_high_band_net(x[:, :, bandw:]),
+            ],
+            dim=2,
+        )
+
+        h = torch.cat([x, aux1], dim=1)
+        aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
+
+        h = torch.cat([x, aux1, aux2], dim=1)
+        h = self.stg3_full_band_net(self.stg3_bridge(h))
+
+        mask = torch.sigmoid(self.out(h))
+        mask = F.pad(
+            input=mask,
+            pad=(0, 0, 0, self.output_bin - mask.size()[2]),
+            mode="replicate",
+        )
+
+        if self.training:
+            aux1 = torch.sigmoid(self.aux1_out(aux1))
+            aux1 = F.pad(
+                input=aux1,
+                pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
+                mode="replicate",
+            )
+            aux2 = torch.sigmoid(self.aux2_out(aux2))
+            aux2 = F.pad(
+                input=aux2,
+                pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
+                mode="replicate",
+            )
+            return mask * mix, aux1 * mix, aux2 * mix
+        else:
+            if aggressiveness:
+                mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
+                    mask[:, :, : aggressiveness["split_bin"]],
+                    1 + aggressiveness["value"] / 3,
+                )
+                mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
+                    mask[:, :, aggressiveness["split_bin"] :],
+                    1 + aggressiveness["value"],
+                )
+
+            return mask * mix
+
+    def predict(self, x_mag, aggressiveness=None):
+        h = self.forward(x_mag, aggressiveness)
+
+        if self.offset > 0:
+            h = h[:, :, :, self.offset : -self.offset]
+            assert h.size()[3] > 0
+
+        return h

tools/uvr5/lib/lib_v5/nets_new.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from . import layers_new
+
+
+class BaseNet(nn.Module):
+    def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6))):
+        super(BaseNet, self).__init__()
+        self.enc1 = layers_new.Conv2DBNActiv(nin, nout, 3, 1, 1)
+        self.enc2 = layers_new.Encoder(nout, nout * 2, 3, 2, 1)
+        self.enc3 = layers_new.Encoder(nout * 2, nout * 4, 3, 2, 1)
+        self.enc4 = layers_new.Encoder(nout * 4, nout * 6, 3, 2, 1)
+        self.enc5 = layers_new.Encoder(nout * 6, nout * 8, 3, 2, 1)
+
+        self.aspp = layers_new.ASPPModule(nout * 8, nout * 8, dilations, dropout=True)
+
+        self.dec4 = layers_new.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1)
+        self.dec3 = layers_new.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1)
+        self.dec2 = layers_new.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1)
+        self.lstm_dec2 = layers_new.LSTMModule(nout * 2, nin_lstm, nout_lstm)
+        self.dec1 = layers_new.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1)
+
+    def __call__(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(e1)
+        e3 = self.enc3(e2)
+        e4 = self.enc4(e3)
+        e5 = self.enc5(e4)
+
+        h = self.aspp(e5)
+
+        h = self.dec4(h, e4)
+        h = self.dec3(h, e3)
+        h = self.dec2(h, e2)
+        h = torch.cat([h, self.lstm_dec2(h)], dim=1)
+        h = self.dec1(h, e1)
+
+        return h
+
+
+class CascadedNet(nn.Module):
+    def __init__(self, n_fft, nout=32, nout_lstm=128):
+        super(CascadedNet, self).__init__()
+
+        self.max_bin = n_fft // 2
+        self.output_bin = n_fft // 2 + 1
+        self.nin_lstm = self.max_bin // 2
+        self.offset = 64
+
+        self.stg1_low_band_net = nn.Sequential(
+            BaseNet(2, nout // 2, self.nin_lstm // 2, nout_lstm),
+            layers_new.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0),
+        )
+
+        self.stg1_high_band_net = BaseNet(2, nout // 4, self.nin_lstm // 2, nout_lstm // 2)
+
+        self.stg2_low_band_net = nn.Sequential(
+            BaseNet(nout // 4 + 2, nout, self.nin_lstm // 2, nout_lstm),
+            layers_new.Conv2DBNActiv(nout, nout // 2, 1, 1, 0),
+        )
+        self.stg2_high_band_net = BaseNet(nout // 4 + 2, nout // 2, self.nin_lstm // 2, nout_lstm // 2)
+
+        self.stg3_full_band_net = BaseNet(3 * nout // 4 + 2, nout, self.nin_lstm, nout_lstm)
+
+        self.out = nn.Conv2d(nout, 2, 1, bias=False)
+        self.aux_out = nn.Conv2d(3 * nout // 4, 2, 1, bias=False)
+
+    def forward(self, x):
+        x = x[:, :, : self.max_bin]
+
+        bandw = x.size()[2] // 2
+        l1_in = x[:, :, :bandw]
+        h1_in = x[:, :, bandw:]
+        l1 = self.stg1_low_band_net(l1_in)
+        h1 = self.stg1_high_band_net(h1_in)
+        aux1 = torch.cat([l1, h1], dim=2)
+
+        l2_in = torch.cat([l1_in, l1], dim=1)
+        h2_in = torch.cat([h1_in, h1], dim=1)
+        l2 = self.stg2_low_band_net(l2_in)
+        h2 = self.stg2_high_band_net(h2_in)
+        aux2 = torch.cat([l2, h2], dim=2)
+
+        f3_in = torch.cat([x, aux1, aux2], dim=1)
+        f3 = self.stg3_full_band_net(f3_in)
+
+        mask = torch.sigmoid(self.out(f3))
+        mask = F.pad(
+            input=mask,
+            pad=(0, 0, 0, self.output_bin - mask.size()[2]),
+            mode="replicate",
+        )
+
+        if self.training:
+            aux = torch.cat([aux1, aux2], dim=1)
+            aux = torch.sigmoid(self.aux_out(aux))
+            aux = F.pad(
+                input=aux,
+                pad=(0, 0, 0, self.output_bin - aux.size()[2]),
+                mode="replicate",
+            )
+            return mask, aux
+        else:
+            return mask
+
+    def predict_mask(self, x):
+        mask = self.forward(x)
+
+        if self.offset > 0:
+            mask = mask[:, :, :, self.offset : -self.offset]
+            assert mask.size()[3] > 0
+
+        return mask
+
+    def predict(self, x, aggressiveness=None):
+        mask = self.forward(x)
+        pred_mag = x * mask
+
+        if self.offset > 0:
+            pred_mag = pred_mag[:, :, :, self.offset : -self.offset]
+            assert pred_mag.size()[3] > 0
+
+        return pred_mag

tools/uvr5/lib/lib_v5/spec_utils.py

@@ -0,0 +1,637 @@
+import hashlib
+import json
+import math
+import os
+
+import librosa
+import numpy as np
+import soundfile as sf
+from tqdm import tqdm
+
+
+def crop_center(h1, h2):
+    h1_shape = h1.size()
+    h2_shape = h2.size()
+
+    if h1_shape[3] == h2_shape[3]:
+        return h1
+    elif h1_shape[3] < h2_shape[3]:
+        raise ValueError("h1_shape[3] must be greater than h2_shape[3]")
+
+    # s_freq = (h2_shape[2] - h1_shape[2]) // 2
+    # e_freq = s_freq + h1_shape[2]
+    s_time = (h1_shape[3] - h2_shape[3]) // 2
+    e_time = s_time + h2_shape[3]
+    h1 = h1[:, :, :, s_time:e_time]
+
+    return h1
+
+
+def wave_to_spectrogram(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
+    if reverse:
+        wave_left = np.flip(np.asfortranarray(wave[0]))
+        wave_right = np.flip(np.asfortranarray(wave[1]))
+    elif mid_side:
+        wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2)
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1]))
+    elif mid_side_b2:
+        wave_left = np.asfortranarray(np.add(wave[1], wave[0] * 0.5))
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * 0.5))
+    else:
+        wave_left = np.asfortranarray(wave[0])
+        wave_right = np.asfortranarray(wave[1])
+
+    spec_left = librosa.stft(wave_left, n_fft=n_fft, hop_length=hop_length)
+    spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length)
+
+    spec = np.asfortranarray([spec_left, spec_right])
+
+    return spec
+
+
+def wave_to_spectrogram_mt(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
+    import threading
+
+    if reverse:
+        wave_left = np.flip(np.asfortranarray(wave[0]))
+        wave_right = np.flip(np.asfortranarray(wave[1]))
+    elif mid_side:
+        wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2)
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1]))
+    elif mid_side_b2:
+        wave_left = np.asfortranarray(np.add(wave[1], wave[0] * 0.5))
+        wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * 0.5))
+    else:
+        wave_left = np.asfortranarray(wave[0])
+        wave_right = np.asfortranarray(wave[1])
+
+    def run_thread(**kwargs):
+        global spec_left
+        spec_left = librosa.stft(**kwargs)
+
+    thread = threading.Thread(
+        target=run_thread,
+        kwargs={"y": wave_left, "n_fft": n_fft, "hop_length": hop_length},
+    )
+    thread.start()
+    spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length)
+    thread.join()
+
+    spec = np.asfortranarray([spec_left, spec_right])
+
+    return spec
+
+
+def combine_spectrograms(specs, mp):
+    l = min([specs[i].shape[2] for i in specs])
+    spec_c = np.zeros(shape=(2, mp.param["bins"] + 1, l), dtype=np.complex64)
+    offset = 0
+    bands_n = len(mp.param["band"])
+
+    for d in range(1, bands_n + 1):
+        h = mp.param["band"][d]["crop_stop"] - mp.param["band"][d]["crop_start"]
+        spec_c[:, offset : offset + h, :l] = specs[d][
+            :, mp.param["band"][d]["crop_start"] : mp.param["band"][d]["crop_stop"], :l
+        ]
+        offset += h
+
+    if offset > mp.param["bins"]:
+        raise ValueError("Too much bins")
+
+    # lowpass fiter
+    if mp.param["pre_filter_start"] > 0:  # and mp.param['band'][bands_n]['res_type'] in ['scipy', 'polyphase']:
+        if bands_n == 1:
+            spec_c = fft_lp_filter(spec_c, mp.param["pre_filter_start"], mp.param["pre_filter_stop"])
+        else:
+            gp = 1
+            for b in range(mp.param["pre_filter_start"] + 1, mp.param["pre_filter_stop"]):
+                g = math.pow(10, -(b - mp.param["pre_filter_start"]) * (3.5 - gp) / 20.0)
+                gp = g
+                spec_c[:, b, :] *= g
+
+    return np.asfortranarray(spec_c)
+
+
+def spectrogram_to_image(spec, mode="magnitude"):
+    if mode == "magnitude":
+        if np.iscomplexobj(spec):
+            y = np.abs(spec)
+        else:
+            y = spec
+        y = np.log10(y**2 + 1e-8)
+    elif mode == "phase":
+        if np.iscomplexobj(spec):
+            y = np.angle(spec)
+        else:
+            y = spec
+
+    y -= y.min()
+    y *= 255 / y.max()
+    img = np.uint8(y)
+
+    if y.ndim == 3:
+        img = img.transpose(1, 2, 0)
+        img = np.concatenate([np.max(img, axis=2, keepdims=True), img], axis=2)
+
+    return img
+
+
+def reduce_vocal_aggressively(X, y, softmask):
+    v = X - y
+    y_mag_tmp = np.abs(y)
+    v_mag_tmp = np.abs(v)
+
+    v_mask = v_mag_tmp > y_mag_tmp
+    y_mag = np.clip(y_mag_tmp - v_mag_tmp * v_mask * softmask, 0, np.inf)
+
+    return y_mag * np.exp(1.0j * np.angle(y))
+
+
+def mask_silence(mag, ref, thres=0.2, min_range=64, fade_size=32):
+    if min_range < fade_size * 2:
+        raise ValueError("min_range must be >= fade_area * 2")
+
+    mag = mag.copy()
+
+    idx = np.where(ref.mean(axis=(0, 1)) < thres)[0]
+    starts = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0])
+    ends = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1])
+    uninformative = np.where(ends - starts > min_range)[0]
+    if len(uninformative) > 0:
+        starts = starts[uninformative]
+        ends = ends[uninformative]
+        old_e = None
+        for s, e in zip(starts, ends):
+            if old_e is not None and s - old_e < fade_size:
+                s = old_e - fade_size * 2
+
+            if s != 0:
+                weight = np.linspace(0, 1, fade_size)
+                mag[:, :, s : s + fade_size] += weight * ref[:, :, s : s + fade_size]
+            else:
+                s -= fade_size
+
+            if e != mag.shape[2]:
+                weight = np.linspace(1, 0, fade_size)
+                mag[:, :, e - fade_size : e] += weight * ref[:, :, e - fade_size : e]
+            else:
+                e += fade_size
+
+            mag[:, :, s + fade_size : e - fade_size] += ref[:, :, s + fade_size : e - fade_size]
+            old_e = e
+
+    return mag
+
+
+def align_wave_head_and_tail(a, b):
+    l = min([a[0].size, b[0].size])
+
+    return a[:l, :l], b[:l, :l]
+
+
+def cache_or_load(mix_path, inst_path, mp):
+    mix_basename = os.path.splitext(os.path.basename(mix_path))[0]
+    inst_basename = os.path.splitext(os.path.basename(inst_path))[0]
+
+    cache_dir = "mph{}".format(hashlib.sha1(json.dumps(mp.param, sort_keys=True).encode("utf-8")).hexdigest())
+    mix_cache_dir = os.path.join("cache", cache_dir)
+    inst_cache_dir = os.path.join("cache", cache_dir)
+
+    os.makedirs(mix_cache_dir, exist_ok=True)
+    os.makedirs(inst_cache_dir, exist_ok=True)
+
+    mix_cache_path = os.path.join(mix_cache_dir, mix_basename + ".npy")
+    inst_cache_path = os.path.join(inst_cache_dir, inst_basename + ".npy")
+
+    if os.path.exists(mix_cache_path) and os.path.exists(inst_cache_path):
+        X_spec_m = np.load(mix_cache_path)
+        y_spec_m = np.load(inst_cache_path)
+    else:
+        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
+
+        for d in range(len(mp.param["band"]), 0, -1):
+            bp = mp.param["band"][d]
+
+            if d == len(mp.param["band"]):  # high-end band
+                X_wave[d], _ = librosa.load(
+                    mix_path, sr=bp["sr"], mono=False, dtype=np.float32, res_type=bp["res_type"]
+                )
+                y_wave[d], _ = librosa.load(
+                    inst_path,
+                    sr=bp["sr"],
+                    mono=False,
+                    dtype=np.float32,
+                    res_type=bp["res_type"],
+                )
+            else:  # lower bands
+                X_wave[d] = librosa.resample(
+                    X_wave[d + 1],
+                    orig_sr=mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+                y_wave[d] = librosa.resample(
+                    y_wave[d + 1],
+                    orig_sr=mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+
+            X_wave[d], y_wave[d] = align_wave_head_and_tail(X_wave[d], y_wave[d])
+
+            X_spec_s[d] = wave_to_spectrogram(
+                X_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                mp.param["mid_side"],
+                mp.param["mid_side_b2"],
+                mp.param["reverse"],
+            )
+            y_spec_s[d] = wave_to_spectrogram(
+                y_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                mp.param["mid_side"],
+                mp.param["mid_side_b2"],
+                mp.param["reverse"],
+            )
+
+        del X_wave, y_wave
+
+        X_spec_m = combine_spectrograms(X_spec_s, mp)
+        y_spec_m = combine_spectrograms(y_spec_s, mp)
+
+        if X_spec_m.shape != y_spec_m.shape:
+            raise ValueError("The combined spectrograms are different: " + mix_path)
+
+        _, ext = os.path.splitext(mix_path)
+
+        np.save(mix_cache_path, X_spec_m)
+        np.save(inst_cache_path, y_spec_m)
+
+    return X_spec_m, y_spec_m
+
+
+def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse):
+    spec_left = np.asfortranarray(spec[0])
+    spec_right = np.asfortranarray(spec[1])
+
+    wave_left = librosa.istft(spec_left, hop_length=hop_length)
+    wave_right = librosa.istft(spec_right, hop_length=hop_length)
+
+    if reverse:
+        return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)])
+    elif mid_side:
+        return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)])
+    elif mid_side_b2:
+        return np.asfortranarray(
+            [
+                np.add(wave_right / 1.25, 0.4 * wave_left),
+                np.subtract(wave_left / 1.25, 0.4 * wave_right),
+            ]
+        )
+    else:
+        return np.asfortranarray([wave_left, wave_right])
+
+
+def spectrogram_to_wave_mt(spec, hop_length, mid_side, reverse, mid_side_b2):
+    import threading
+
+    spec_left = np.asfortranarray(spec[0])
+    spec_right = np.asfortranarray(spec[1])
+
+    def run_thread(**kwargs):
+        global wave_left
+        wave_left = librosa.istft(**kwargs)
+
+    thread = threading.Thread(target=run_thread, kwargs={"stft_matrix": spec_left, "hop_length": hop_length})
+    thread.start()
+    wave_right = librosa.istft(spec_right, hop_length=hop_length)
+    thread.join()
+
+    if reverse:
+        return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)])
+    elif mid_side:
+        return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)])
+    elif mid_side_b2:
+        return np.asfortranarray(
+            [
+                np.add(wave_right / 1.25, 0.4 * wave_left),
+                np.subtract(wave_left / 1.25, 0.4 * wave_right),
+            ]
+        )
+    else:
+        return np.asfortranarray([wave_left, wave_right])
+
+
+def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None):
+    wave_band = {}
+    bands_n = len(mp.param["band"])
+    offset = 0
+
+    for d in range(1, bands_n + 1):
+        bp = mp.param["band"][d]
+        spec_s = np.ndarray(shape=(2, bp["n_fft"] // 2 + 1, spec_m.shape[2]), dtype=complex)
+        h = bp["crop_stop"] - bp["crop_start"]
+        spec_s[:, bp["crop_start"] : bp["crop_stop"], :] = spec_m[:, offset : offset + h, :]
+
+        offset += h
+        if d == bands_n:  # higher
+            if extra_bins_h:  # if --high_end_process bypass
+                max_bin = bp["n_fft"] // 2
+                spec_s[:, max_bin - extra_bins_h : max_bin, :] = extra_bins[:, :extra_bins_h, :]
+            if bp["hpf_start"] > 0:
+                spec_s = fft_hp_filter(spec_s, bp["hpf_start"], bp["hpf_stop"] - 1)
+            if bands_n == 1:
+                wave = spectrogram_to_wave(
+                    spec_s,
+                    bp["hl"],
+                    mp.param["mid_side"],
+                    mp.param["mid_side_b2"],
+                    mp.param["reverse"],
+                )
+            else:
+                wave = np.add(
+                    wave,
+                    spectrogram_to_wave(
+                        spec_s,
+                        bp["hl"],
+                        mp.param["mid_side"],
+                        mp.param["mid_side_b2"],
+                        mp.param["reverse"],
+                    ),
+                )
+        else:
+            sr = mp.param["band"][d + 1]["sr"]
+            if d == 1:  # lower
+                spec_s = fft_lp_filter(spec_s, bp["lpf_start"], bp["lpf_stop"])
+                wave = librosa.resample(
+                    spectrogram_to_wave(
+                        spec_s,
+                        bp["hl"],
+                        mp.param["mid_side"],
+                        mp.param["mid_side_b2"],
+                        mp.param["reverse"],
+                    ),
+                    orig_sr=bp["sr"],
+                    target_sr=sr,
+                    res_type="sinc_fastest",
+                )
+            else:  # mid
+                spec_s = fft_hp_filter(spec_s, bp["hpf_start"], bp["hpf_stop"] - 1)
+                spec_s = fft_lp_filter(spec_s, bp["lpf_start"], bp["lpf_stop"])
+                wave2 = np.add(
+                    wave,
+                    spectrogram_to_wave(
+                        spec_s,
+                        bp["hl"],
+                        mp.param["mid_side"],
+                        mp.param["mid_side_b2"],
+                        mp.param["reverse"],
+                    ),
+                )
+                # wave = librosa.core.resample(wave2, orig_sr=bp['sr'], target_sr=sr, res_type="sinc_fastest")
+                wave = librosa.core.resample(wave2, orig_sr=bp["sr"], target_sr=sr, res_type="scipy")
+
+    return wave.T
+
+
+def fft_lp_filter(spec, bin_start, bin_stop):
+    g = 1.0
+    for b in range(bin_start, bin_stop):
+        g -= 1 / (bin_stop - bin_start)
+        spec[:, b, :] = g * spec[:, b, :]
+
+    spec[:, bin_stop:, :] *= 0
+
+    return spec
+
+
+def fft_hp_filter(spec, bin_start, bin_stop):
+    g = 1.0
+    for b in range(bin_start, bin_stop, -1):
+        g -= 1 / (bin_start - bin_stop)
+        spec[:, b, :] = g * spec[:, b, :]
+
+    spec[:, 0 : bin_stop + 1, :] *= 0
+
+    return spec
+
+
+def mirroring(a, spec_m, input_high_end, mp):
+    if "mirroring" == a:
+        mirror = np.flip(
+            np.abs(
+                spec_m[
+                    :,
+                    mp.param["pre_filter_start"] - 10 - input_high_end.shape[1] : mp.param["pre_filter_start"] - 10,
+                    :,
+                ]
+            ),
+            1,
+        )
+        mirror = mirror * np.exp(1.0j * np.angle(input_high_end))
+
+        return np.where(np.abs(input_high_end) <= np.abs(mirror), input_high_end, mirror)
+
+    if "mirroring2" == a:
+        mirror = np.flip(
+            np.abs(
+                spec_m[
+                    :,
+                    mp.param["pre_filter_start"] - 10 - input_high_end.shape[1] : mp.param["pre_filter_start"] - 10,
+                    :,
+                ]
+            ),
+            1,
+        )
+        mi = np.multiply(mirror, input_high_end * 1.7)
+
+        return np.where(np.abs(input_high_end) <= np.abs(mi), input_high_end, mi)
+
+
+def ensembling(a, specs):
+    for i in range(1, len(specs)):
+        if i == 1:
+            spec = specs[0]
+
+        ln = min([spec.shape[2], specs[i].shape[2]])
+        spec = spec[:, :, :ln]
+        specs[i] = specs[i][:, :, :ln]
+
+        if "min_mag" == a:
+            spec = np.where(np.abs(specs[i]) <= np.abs(spec), specs[i], spec)
+        if "max_mag" == a:
+            spec = np.where(np.abs(specs[i]) >= np.abs(spec), specs[i], spec)
+
+    return spec
+
+
+def stft(wave, nfft, hl):
+    wave_left = np.asfortranarray(wave[0])
+    wave_right = np.asfortranarray(wave[1])
+    spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl)
+    spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl)
+    spec = np.asfortranarray([spec_left, spec_right])
+
+    return spec
+
+
+def istft(spec, hl):
+    spec_left = np.asfortranarray(spec[0])
+    spec_right = np.asfortranarray(spec[1])
+
+    wave_left = librosa.istft(spec_left, hop_length=hl)
+    wave_right = librosa.istft(spec_right, hop_length=hl)
+    wave = np.asfortranarray([wave_left, wave_right])
+
+    return wave
+
+
+if __name__ == "__main__":
+    import argparse
+    import time
+
+    import cv2
+    from model_param_init import ModelParameters
+
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--algorithm",
+        "-a",
+        type=str,
+        choices=["invert", "invert_p", "min_mag", "max_mag", "deep", "align"],
+        default="min_mag",
+    )
+    p.add_argument(
+        "--model_params",
+        "-m",
+        type=str,
+        default=os.path.join("modelparams", "1band_sr44100_hl512.json"),
+    )
+    p.add_argument("--output_name", "-o", type=str, default="output")
+    p.add_argument("--vocals_only", "-v", action="store_true")
+    p.add_argument("input", nargs="+")
+    args = p.parse_args()
+
+    start_time = time.time()
+
+    if args.algorithm.startswith("invert") and len(args.input) != 2:
+        raise ValueError("There should be two input files.")
+
+    if not args.algorithm.startswith("invert") and len(args.input) < 2:
+        raise ValueError("There must be at least two input files.")
+
+    wave, specs = {}, {}
+    mp = ModelParameters(args.model_params)
+
+    for i in range(len(args.input)):
+        spec = {}
+
+        for d in range(len(mp.param["band"]), 0, -1):
+            bp = mp.param["band"][d]
+
+            if d == len(mp.param["band"]):  # high-end band
+                wave[d], _ = librosa.load(
+                    args.input[i],
+                    sr=bp["sr"],
+                    mono=False,
+                    dtype=np.float32,
+                    res_type=bp["res_type"],
+                )
+
+                if len(wave[d].shape) == 1:  # mono to stereo
+                    wave[d] = np.array([wave[d], wave[d]])
+            else:  # lower bands
+                wave[d] = librosa.resample(
+                    wave[d + 1],
+                    orig_sr=mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+
+            spec[d] = wave_to_spectrogram(
+                wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                mp.param["mid_side"],
+                mp.param["mid_side_b2"],
+                mp.param["reverse"],
+            )
+
+        specs[i] = combine_spectrograms(spec, mp)
+
+    del wave
+
+    if args.algorithm == "deep":
+        d_spec = np.where(np.abs(specs[0]) <= np.abs(spec[1]), specs[0], spec[1])
+        v_spec = d_spec - specs[1]
+        sf.write(
+            os.path.join("{}.wav".format(args.output_name)),
+            cmb_spectrogram_to_wave(v_spec, mp),
+            mp.param["sr"],
+        )
+
+    if args.algorithm.startswith("invert"):
+        ln = min([specs[0].shape[2], specs[1].shape[2]])
+        specs[0] = specs[0][:, :, :ln]
+        specs[1] = specs[1][:, :, :ln]
+
+        if "invert_p" == args.algorithm:
+            X_mag = np.abs(specs[0])
+            y_mag = np.abs(specs[1])
+            max_mag = np.where(X_mag >= y_mag, X_mag, y_mag)
+            v_spec = specs[1] - max_mag * np.exp(1.0j * np.angle(specs[0]))
+        else:
+            specs[1] = reduce_vocal_aggressively(specs[0], specs[1], 0.2)
+            v_spec = specs[0] - specs[1]
+
+            if not args.vocals_only:
+                X_mag = np.abs(specs[0])
+                y_mag = np.abs(specs[1])
+                v_mag = np.abs(v_spec)
+
+                X_image = spectrogram_to_image(X_mag)
+                y_image = spectrogram_to_image(y_mag)
+                v_image = spectrogram_to_image(v_mag)
+
+                cv2.imwrite("{}_X.png".format(args.output_name), X_image)
+                cv2.imwrite("{}_y.png".format(args.output_name), y_image)
+                cv2.imwrite("{}_v.png".format(args.output_name), v_image)
+
+                sf.write(
+                    "{}_X.wav".format(args.output_name),
+                    cmb_spectrogram_to_wave(specs[0], mp),
+                    mp.param["sr"],
+                )
+                sf.write(
+                    "{}_y.wav".format(args.output_name),
+                    cmb_spectrogram_to_wave(specs[1], mp),
+                    mp.param["sr"],
+                )
+
+        sf.write(
+            "{}_v.wav".format(args.output_name),
+            cmb_spectrogram_to_wave(v_spec, mp),
+            mp.param["sr"],
+        )
+    else:
+        if not args.algorithm == "deep":
+            sf.write(
+                os.path.join("ensembled", "{}.wav".format(args.output_name)),
+                cmb_spectrogram_to_wave(ensembling(args.algorithm, specs), mp),
+                mp.param["sr"],
+            )
+
+    if args.algorithm == "align":
+        trackalignment = [
+            {
+                "file1": '"{}"'.format(args.input[0]),
+                "file2": '"{}"'.format(args.input[1]),
+            }
+        ]
+
+        for i, e in tqdm(enumerate(trackalignment), desc="Performing Alignment..."):
+            os.system(f"python lib/align_tracks.py {e['file1']} {e['file2']}")
+
+    # print('Total time: {0:.{1}f}s'.format(time.time() - start_time, 1))

tools/uvr5/lib/utils.py

@@ -0,0 +1,82 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+
+def make_padding(width, cropsize, offset):
+    left = offset
+    roi_size = cropsize - left * 2
+    if roi_size == 0:
+        roi_size = cropsize
+    right = roi_size - (width % roi_size) + left
+
+    return left, right, roi_size
+
+
+def inference(X_spec, device, model, aggressiveness, data):
+    """
+    data : dic configs
+    """
+
+    def _execute(X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half=True):
+        model.eval()
+        with torch.no_grad():
+            preds = []
+
+            iterations = [n_window]
+
+            total_iterations = sum(iterations)
+            for i in tqdm(range(n_window)):
+                start = i * roi_size
+                X_mag_window = X_mag_pad[None, :, :, start : start + data["window_size"]]
+                X_mag_window = torch.from_numpy(X_mag_window)
+                if is_half:
+                    X_mag_window = X_mag_window.half()
+                X_mag_window = X_mag_window.to(device)
+
+                pred = model.predict(X_mag_window, aggressiveness)
+
+                pred = pred.detach().cpu().numpy()
+                preds.append(pred[0])
+
+            pred = np.concatenate(preds, axis=2)
+        return pred
+
+    def preprocess(X_spec):
+        X_mag = np.abs(X_spec)
+        X_phase = np.angle(X_spec)
+
+        return X_mag, X_phase
+
+    X_mag, X_phase = preprocess(X_spec)
+
+    coef = X_mag.max()
+    X_mag_pre = X_mag / coef
+
+    n_frame = X_mag_pre.shape[2]
+    pad_l, pad_r, roi_size = make_padding(n_frame, data["window_size"], model.offset)
+    n_window = int(np.ceil(n_frame / roi_size))
+
+    X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")
+
+    if list(model.state_dict().values())[0].dtype == torch.float16:
+        is_half = True
+    else:
+        is_half = False
+    pred = _execute(X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half)
+    pred = pred[:, :, :n_frame]
+
+    if data["tta"]:
+        pad_l += roi_size // 2
+        pad_r += roi_size // 2
+        n_window += 1
+
+        X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")
+
+        pred_tta = _execute(X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half)
+        pred_tta = pred_tta[:, :, roi_size // 2 :]
+        pred_tta = pred_tta[:, :, :n_frame]
+
+        return (pred + pred_tta) * 0.5 * coef, X_mag, np.exp(1.0j * X_phase)
+    else:
+        return pred * coef, X_mag, np.exp(1.0j * X_phase)

tools/uvr5/mdxnet.py

@@ -0,0 +1,315 @@
+import os
+import logging
+import sysconfig
+
+logger = logging.getLogger(__name__)
+
+import librosa
+import numpy as np
+import soundfile as sf
+import torch
+from tqdm import tqdm
+
+
+_ORT_CUDA_DLL_HANDLES = []
+
+
+def _configure_ort_cuda_dll_paths():
+    """Expose pip-installed CUDA 11/cuDNN 8 DLLs to ONNX Runtime on Windows."""
+    if os.name != "nt":
+        return
+
+    site_packages = os.path.normpath(sysconfig.get_paths()["purelib"])
+    nvidia_root = os.path.join(site_packages, "nvidia")
+    dll_dirs = [
+        os.path.join(nvidia_root, "cuda_runtime", "bin"),
+        os.path.join(nvidia_root, "cublas", "bin"),
+        os.path.join(nvidia_root, "cufft", "bin"),
+        os.path.join(nvidia_root, "cudnn", "bin"),
+        os.path.join(nvidia_root, "cuda_nvrtc", "bin"),
+        os.path.join(os.path.dirname(torch.__file__), "lib"),
+    ]
+    dll_dirs = [path for path in dll_dirs if os.path.isdir(path)]
+    if not dll_dirs:
+        return
+
+    current_path = os.environ.get("PATH", "")
+    current_dirs = [path for path in current_path.split(os.pathsep) if path]
+    known_dirs = {os.path.normcase(os.path.normpath(path)) for path in current_dirs}
+    prepend_dirs = []
+    for path in dll_dirs:
+        normalized = os.path.normcase(os.path.normpath(path))
+        if normalized not in known_dirs:
+            prepend_dirs.append(path)
+            known_dirs.add(normalized)
+    if prepend_dirs:
+        os.environ["PATH"] = os.pathsep.join(prepend_dirs + current_dirs)
+
+    # Python 3.8+ restricts DLL lookup for extension modules. Keep the handles
+    # alive for the process lifetime in addition to updating PATH.
+    if hasattr(os, "add_dll_directory"):
+        for path in dll_dirs:
+            try:
+                _ORT_CUDA_DLL_HANDLES.append(os.add_dll_directory(path))
+            except OSError:
+                logger.warning("Unable to add ONNX Runtime DLL directory: %s", path)
+
+
+_configure_ort_cuda_dll_paths()
+
+cpu = torch.device("cpu")
+
+
+class ConvTDFNetTrim:
+    def __init__(self, device, model_name, target_name, L, dim_f, dim_t, n_fft, hop=1024):
+        super(ConvTDFNetTrim, self).__init__()
+
+        self.dim_f = dim_f
+        self.dim_t = 2**dim_t
+        self.n_fft = n_fft
+        self.hop = hop
+        self.n_bins = self.n_fft // 2 + 1
+        self.chunk_size = hop * (self.dim_t - 1)
+        self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(device)
+        self.target_name = target_name
+        self.blender = "blender" in model_name
+
+        self.dim_c = 4
+        out_c = self.dim_c * 4 if target_name == "*" else self.dim_c
+        self.freq_pad = torch.zeros([1, out_c, self.n_bins - self.dim_f, self.dim_t]).to(device)
+
+        self.n = L // 2
+
+    def stft(self, x):
+        x = x.reshape([-1, self.chunk_size])
+        x = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop,
+            window=self.window,
+            center=True,
+            return_complex=True,
+        )
+        x = torch.view_as_real(x)
+        x = x.permute([0, 3, 1, 2])
+        x = x.reshape([-1, 2, 2, self.n_bins, self.dim_t]).reshape([-1, self.dim_c, self.n_bins, self.dim_t])
+        return x[:, :, : self.dim_f]
+
+    def istft(self, x, freq_pad=None):
+        freq_pad = self.freq_pad.repeat([x.shape[0], 1, 1, 1]) if freq_pad is None else freq_pad
+        x = torch.cat([x, freq_pad], -2)
+        c = 4 * 2 if self.target_name == "*" else 2
+        x = x.reshape([-1, c, 2, self.n_bins, self.dim_t]).reshape([-1, 2, self.n_bins, self.dim_t])
+        x = x.permute([0, 2, 3, 1])
+        x = x.contiguous()
+        x = torch.view_as_complex(x)
+        x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True)
+        return x.reshape([-1, c, self.chunk_size])
+
+
+def get_models(device, dim_f, dim_t, n_fft):
+    return ConvTDFNetTrim(
+        device=device,
+        model_name="Conv-TDF",
+        target_name="vocals",
+        L=11,
+        dim_f=dim_f,
+        dim_t=dim_t,
+        n_fft=n_fft,
+    )
+
+
+class Predictor:
+    def __init__(self, args):
+        import onnxruntime as ort
+
+        available_providers = ort.get_available_providers()
+        logger.info("ONNX Runtime available providers: %s", available_providers)
+        if (
+            "CUDAExecutionProvider" in args.providers
+            and "CUDAExecutionProvider" not in available_providers
+        ):
+            raise RuntimeError(
+                "CUDAExecutionProvider is required for the FoxJoy ONNX model, "
+                "but the installed ONNX Runtime does not provide it. Install "
+                "the matching CUDA ONNX Runtime dependencies with this "
+                "project's runtime Python."
+            )
+        if (
+            "DmlExecutionProvider" in args.providers
+            and "DmlExecutionProvider" not in available_providers
+        ):
+            raise RuntimeError(
+                "DmlExecutionProvider is required for the FoxJoy ONNX model, "
+                "but the installed ONNX Runtime does not provide it. Install "
+                "requirments_cpu_py312.txt with this project's runtime Python."
+            )
+        self.args = args
+        self.model_ = get_models(device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft)
+        self.model = ort.InferenceSession(
+            os.path.join(args.onnx, self.model_.target_name + ".onnx"),
+            providers=args.providers,
+        )
+        active_providers = self.model.get_providers()
+        logger.info("ONNX Runtime active providers: %s", active_providers)
+        if (
+            "CUDAExecutionProvider" in args.providers
+            and (
+                not active_providers
+                or active_providers[0] != "CUDAExecutionProvider"
+            )
+        ):
+            raise RuntimeError(
+                "The FoxJoy ONNX model did not activate CUDAExecutionProvider; "
+                "check the CUDA 11/cuDNN 8 DLL installation."
+            )
+        if (
+            "DmlExecutionProvider" in args.providers
+            and (
+                not active_providers
+                or active_providers[0] != "DmlExecutionProvider"
+            )
+        ):
+            raise RuntimeError(
+                "The FoxJoy ONNX model did not activate DmlExecutionProvider; "
+                "check the ONNX Runtime DirectML installation."
+            )
+        logger.info("ONNX load done")
+
+    def demix(self, mix):
+        samples = mix.shape[-1]
+        margin = self.args.margin
+        chunk_size = self.args.chunks * 44100
+        assert not margin == 0, "margin cannot be zero!"
+        if margin > chunk_size:
+            margin = chunk_size
+
+        segmented_mix = {}
+
+        if self.args.chunks == 0 or samples < chunk_size:
+            chunk_size = samples
+
+        counter = -1
+        for skip in range(0, samples, chunk_size):
+            counter += 1
+
+            s_margin = 0 if counter == 0 else margin
+            end = min(skip + chunk_size + margin, samples)
+
+            start = skip - s_margin
+
+            segmented_mix[skip] = mix[:, start:end].copy()
+            if end == samples:
+                break
+
+        sources = self.demix_base(segmented_mix, margin_size=margin)
+        """
+        mix:(2,big_sample)
+        segmented_mix:offset->(2,small_sample)
+        sources:(1,2,big_sample)
+        """
+        return sources
+
+    def demix_base(self, mixes, margin_size):
+        chunked_sources = []
+        progress_bar = tqdm(total=len(mixes))
+        progress_bar.set_description("Processing")
+        for mix in mixes:
+            cmix = mixes[mix]
+            sources = []
+            n_sample = cmix.shape[1]
+            model = self.model_
+            trim = model.n_fft // 2
+            gen_size = model.chunk_size - 2 * trim
+            pad = gen_size - n_sample % gen_size
+            mix_p = np.concatenate((np.zeros((2, trim)), cmix, np.zeros((2, pad)), np.zeros((2, trim))), 1)
+            mix_waves = []
+            i = 0
+            while i < n_sample + pad:
+                waves = np.array(mix_p[:, i : i + model.chunk_size])
+                mix_waves.append(waves)
+                i += gen_size
+            mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(cpu)
+            with torch.no_grad():
+                _ort = self.model
+                spek = model.stft(mix_waves)
+                if self.args.denoise:
+                    spec_pred = (
+                        -_ort.run(None, {"input": -spek.cpu().numpy()})[0] * 0.5
+                        + _ort.run(None, {"input": spek.cpu().numpy()})[0] * 0.5
+                    )
+                    tar_waves = model.istft(torch.tensor(spec_pred))
+                else:
+                    tar_waves = model.istft(torch.tensor(_ort.run(None, {"input": spek.cpu().numpy()})[0]))
+                tar_signal = tar_waves[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).numpy()[:, :-pad]
+
+                start = 0 if mix == 0 else margin_size
+                end = None if mix == list(mixes.keys())[::-1][0] else -margin_size
+                if margin_size == 0:
+                    end = None
+                sources.append(tar_signal[:, start:end])
+
+                progress_bar.update(1)
+
+            chunked_sources.append(sources)
+        _sources = np.concatenate(chunked_sources, axis=-1)
+        # del self.model
+        progress_bar.close()
+        return _sources
+
+    def prediction(self, m, vocal_root, others_root, format):
+        os.makedirs(vocal_root, exist_ok=True)
+        os.makedirs(others_root, exist_ok=True)
+        basename = os.path.basename(m)
+        mix, rate = librosa.load(m, mono=False, sr=44100)
+        if mix.ndim == 1:
+            mix = np.asfortranarray([mix, mix])
+        mix = mix.T
+        sources = self.demix(mix.T)
+        opt = sources[0].T
+        if format in ["wav", "flac"]:
+            sf.write("%s/%s_main_vocal.%s" % (vocal_root, basename, format), mix - opt, rate)
+            sf.write("%s/%s_others.%s" % (others_root, basename, format), opt, rate)
+        else:
+            path_vocal = "%s/%s_main_vocal.wav" % (vocal_root, basename)
+            path_other = "%s/%s_others.wav" % (others_root, basename)
+            sf.write(path_vocal, mix - opt, rate)
+            sf.write(path_other, opt, rate)
+            opt_path_vocal = path_vocal[:-4] + ".%s" % format
+            opt_path_other = path_other[:-4] + ".%s" % format
+            if os.path.exists(path_vocal):
+                os.system('ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path_vocal, opt_path_vocal))
+                if os.path.exists(opt_path_vocal):
+                    try:
+                        os.remove(path_vocal)
+                    except:
+                        pass
+            if os.path.exists(path_other):
+                os.system('ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path_other, opt_path_other))
+                if os.path.exists(opt_path_other):
+                    try:
+                        os.remove(path_other)
+                    except:
+                        pass
+
+
+class MDXNetDereverb:
+    def __init__(self, chunks, providers):
+        self.onnx = os.path.join(
+            os.getenv("weight_uvr5_root", "assets/uvr5_weights"),
+            "onnx_dereverb_By_FoxJoy",
+        )
+        self.shifts = 10  # 'Predict with randomised equivariant stabilisation'
+        self.mixing = "min_mag"  # ['default','min_mag','max_mag']
+        self.chunks = chunks
+        self.providers = providers
+        self.margin = 44100
+        self.dim_t = 9
+        self.dim_f = 3072
+        self.n_fft = 6144
+        self.denoise = True
+        self.pred = Predictor(self)
+        self.device = cpu
+
+    def _path_audio_(self, input, others_root, vocal_root, format, is_hp3=False):
+        self.pred.prediction(input, vocal_root, others_root, format)

tools/uvr5/rotary_embedding_torch/__init__.py

@@ -0,0 +1,6 @@
+from .rotary_embedding_torch import (
+    apply_rotary_emb,
+    RotaryEmbedding,
+    apply_learned_rotations,
+    broadcat
+)

tools/uvr5/rotary_embedding_torch/rotary_embedding_torch.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+from math import pi, log
+import warnings
+
+warnings.filterwarnings(
+    "ignore",
+    message="`torch.cuda.amp.autocast.*is deprecated.*",
+    category=FutureWarning,
+)
+
+import torch
+from torch.nn import Module, ModuleList
+from torch.cuda.amp import autocast
+from torch import nn, einsum, broadcast_tensors, Tensor
+from einops import rearrange, repeat
+from typing import Literal
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def broadcat(tensors, dim=-1):
+    broadcasted_tensors = broadcast_tensors(*tensors)
+    return torch.cat(broadcasted_tensors, dim=dim)
+
+def rotate_half(x):
+    x = rearrange(x, '... (d r) -> ... d r', r=2)
+    (x1, x2) = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, '... d r -> ... (d r)')
+
+@autocast(enabled=False)
+def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2):
+    dtype = t.dtype
+    if t.ndim == 3:
+        seq_len = t.shape[seq_dim]
+        freqs = freqs[-seq_len:]
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+    assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+    (t_left, t, t_right) = (t[..., :start_index], t[..., start_index:end_index], t[..., end_index:])
+    t = t * freqs.cos() * scale + rotate_half(t) * freqs.sin() * scale
+    if t.device.type == 'privateuseone':
+        # DirectML rejects concatenation when one of the slices has a zero
+        # length.  Rotary embeddings normally cover the complete head, so both
+        # edge slices are empty; omitting them is mathematically identical.
+        parts = tuple(part for part in (t_left, t, t_right) if part.shape[-1] > 0)
+        out = parts[0] if len(parts) == 1 else torch.cat(parts, dim=-1)
+    else:
+        out = torch.cat((t_left, t, t_right), dim=-1)
+    return out.type(dtype)
+
+def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None):
+    if exists(freq_ranges):
+        rotations = einsum('..., f -> ... f', rotations, freq_ranges)
+        rotations = rearrange(rotations, '... r f -> ... (r f)')
+    rotations = repeat(rotations, '... n -> ... (n r)', r=2)
+    return apply_rotary_emb(rotations, t, start_index=start_index)
+
+class RotaryEmbedding(Module):
+
+    def __init__(self, dim, custom_freqs=None, freqs_for='lang', theta=10000, max_freq=10, num_freqs=1, learned_freq=False, use_xpos=False, xpos_scale_base=512, interpolate_factor=1.0, theta_rescale_factor=1.0, seq_before_head_dim=False, cache_if_possible=True):
+        super().__init__()
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+        self.freqs_for = freqs_for
+        if exists(custom_freqs):
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1.0 / theta ** (torch.arange(0, dim, 2)[:dim // 2].float() / dim)
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        self.cache_if_possible = cache_if_possible
+        self.tmp_store('cached_freqs', None)
+        self.tmp_store('cached_scales', None)
+        self.freqs = nn.Parameter(freqs, requires_grad=learned_freq)
+        self.learned_freq = learned_freq
+        self.tmp_store('dummy', torch.tensor(0))
+        self.seq_before_head_dim = seq_before_head_dim
+        self.default_seq_dim = -3 if seq_before_head_dim else -2
+        assert interpolate_factor >= 1.0
+        self.interpolate_factor = interpolate_factor
+        self.use_xpos = use_xpos
+        if not use_xpos:
+            self.tmp_store('scale', None)
+            return
+        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+        self.scale_base = xpos_scale_base
+        self.tmp_store('scale', scale)
+        self.apply_rotary_emb = staticmethod(apply_rotary_emb)
+
+    @property
+    def device(self):
+        return self.dummy.device
+
+    def tmp_store(self, key, value):
+        self.register_buffer(key, value, persistent=False)
+
+    def get_seq_pos(self, seq_len, device, dtype, offset=0):
+        return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor
+
+    def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+        assert not self.use_xpos or exists(scale), 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings'
+        (device, dtype, seq_len) = (t.device, t.dtype, t.shape[seq_dim])
+        seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)
+        freqs = self.forward(seq, seq_len=seq_len, offset=offset)
+        if seq_dim == -3:
+            freqs = rearrange(freqs, 'n d -> n 1 d')
+        return apply_rotary_emb(freqs, t, scale=default(scale, 1.0), seq_dim=seq_dim)
+
+    def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0):
+        (dtype, device, seq_dim) = (q.dtype, q.device, default(seq_dim, self.default_seq_dim))
+        (q_len, k_len) = (q.shape[seq_dim], k.shape[seq_dim])
+        assert q_len <= k_len
+        q_scale = k_scale = 1.0
+        if self.use_xpos:
+            seq = self.get_seq_pos(k_len, dtype=dtype, device=device)
+            q_scale = self.get_scale(seq[-q_len:]).type(dtype)
+            k_scale = self.get_scale(seq).type(dtype)
+        rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset)
+        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale ** (-1))
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+        return (rotated_q, rotated_k)
+
+    def rotate_queries_and_keys(self, q, k, seq_dim=None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+        assert self.use_xpos
+        (device, dtype, seq_len) = (q.device, q.dtype, q.shape[seq_dim])
+        seq = self.get_seq_pos(seq_len, dtype=dtype, device=device)
+        freqs = self.forward(seq, seq_len=seq_len)
+        scale = self.get_scale(seq, seq_len=seq_len).to(dtype)
+        if seq_dim == -3:
+            freqs = rearrange(freqs, 'n d -> n 1 d')
+            scale = rearrange(scale, 'n d -> n 1 d')
+        rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim)
+        rotated_k = apply_rotary_emb(freqs, k, scale=scale ** (-1), seq_dim=seq_dim)
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+        return (rotated_q, rotated_k)
+
+    def get_scale(self, t, seq_len=None, offset=0):
+        assert self.use_xpos
+        should_cache = self.cache_if_possible and exists(seq_len)
+        if should_cache and exists(self.cached_scales) and (seq_len + offset <= self.cached_scales.shape[0]):
+            return self.cached_scales[offset:offset + seq_len]
+        scale = 1.0
+        if self.use_xpos:
+            power = (t - len(t) // 2) / self.scale_base
+            scale = self.scale ** rearrange(power, 'n -> n 1')
+            scale = torch.cat((scale, scale), dim=-1)
+        if should_cache:
+            self.tmp_store('cached_scales', scale)
+        return scale
+
+    def get_axial_freqs(self, *dims):
+        Colon = slice(None)
+        all_freqs = []
+        for (ind, dim) in enumerate(dims):
+            if self.freqs_for == 'pixel':
+                pos = torch.linspace(-1, 1, steps=dim, device=self.device)
+            else:
+                pos = torch.arange(dim, device=self.device)
+            freqs = self.forward(pos, seq_len=dim)
+            all_axis = [None] * len(dims)
+            all_axis[ind] = Colon
+            new_axis_slice = (Ellipsis, *all_axis, Colon)
+            all_freqs.append(freqs[new_axis_slice])
+        all_freqs = broadcast_tensors(*all_freqs)
+        return torch.cat(all_freqs, dim=-1)
+
+    @autocast(enabled=False)
+    def forward(self, t, seq_len=None, offset=0):
+        should_cache = self.cache_if_possible and (not self.learned_freq) and exists(seq_len) and (self.freqs_for != 'pixel')
+        if should_cache and exists(self.cached_freqs) and (offset + seq_len <= self.cached_freqs.shape[0]):
+            return self.cached_freqs[offset:offset + seq_len].detach()
+        freqs = self.freqs
+        freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs)
+        freqs = repeat(freqs, '... n -> ... (n r)', r=2)
+        if should_cache:
+            self.tmp_store('cached_freqs', freqs.detach())
+        return freqs

tools/uvr5/vr.py

@@ -0,0 +1,350 @@
+import os
+
+parent_directory = os.path.dirname(os.path.abspath(__file__))
+import logging
+
+logger = logging.getLogger(__name__)
+
+import librosa
+import numpy as np
+import soundfile as sf
+import torch
+from tools.uvr5.lib.lib_v5 import nets_61968KB as Nets
+from tools.uvr5.lib.lib_v5 import spec_utils
+from tools.uvr5.lib.lib_v5.model_param_init import ModelParameters
+from tools.uvr5.lib.lib_v5.nets_new import CascadedNet
+from tools.uvr5.lib.utils import inference
+
+
+class AudioPre:
+    def __init__(self, agg, model_path, device, is_half, tta=False):
+        self.model_path = model_path
+        self.device = device
+        self.data = {
+            # Processing Options
+            "postprocess": False,
+            "tta": tta,
+            # Constants
+            "window_size": 512,
+            "agg": agg,
+            "high_end_process": "mirroring",
+        }
+        mp = ModelParameters("%s/lib/lib_v5/modelparams/4band_v2.json" % parent_directory)
+        model = Nets.CascadedASPPNet(mp.param["bins"] * 2)
+        cpk = torch.load(model_path, map_location="cpu")
+        model.load_state_dict(cpk)
+        model.eval()
+        if is_half:
+            model = model.half().to(device)
+        else:
+            model = model.to(device)
+
+        self.mp = mp
+        self.model = model
+
+    def _path_audio_(self, music_file, ins_root=None, vocal_root=None, format="flac", is_hp3=False):
+        if ins_root is None and vocal_root is None:
+            return "No save root."
+        name = os.path.basename(music_file)
+        if ins_root is not None:
+            os.makedirs(ins_root, exist_ok=True)
+        if vocal_root is not None:
+            os.makedirs(vocal_root, exist_ok=True)
+        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
+        bands_n = len(self.mp.param["band"])
+        # print(bands_n)
+        for d in range(bands_n, 0, -1):
+            bp = self.mp.param["band"][d]
+            if d == bands_n:  # high-end band
+                (
+                    X_wave[d],
+                    _,
+                ) = librosa.core.load(  # 理论上librosa读取可能对某些音频有bug，应该上ffmpeg读取，但是太麻烦了弃坑
+                    music_file,
+                    sr=bp["sr"],
+                    mono=False,
+                    dtype=np.float32,
+                    res_type=bp["res_type"],
+                )
+                if X_wave[d].ndim == 1:
+                    X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])
+            else:  # lower bands
+                X_wave[d] = librosa.core.resample(
+                    X_wave[d + 1],
+                    orig_sr=self.mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+            # Stft of wave source
+            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(
+                X_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                self.mp.param["mid_side"],
+                self.mp.param["mid_side_b2"],
+                self.mp.param["reverse"],
+            )
+            # pdb.set_trace()
+            if d == bands_n and self.data["high_end_process"] != "none":
+                input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
+                    self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]
+                )
+                input_high_end = X_spec_s[d][:, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :]
+
+        X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)
+        aggresive_set = float(self.data["agg"] / 100)
+        aggressiveness = {
+            "value": aggresive_set,
+            "split_bin": self.mp.param["band"][1]["crop_stop"],
+        }
+        with torch.no_grad():
+            pred, X_mag, X_phase = inference(X_spec_m, self.device, self.model, aggressiveness, self.data)
+        # Postprocess
+        if self.data["postprocess"]:
+            pred_inv = np.clip(X_mag - pred, 0, np.inf)
+            pred = spec_utils.mask_silence(pred, pred_inv)
+        y_spec_m = pred * X_phase
+        v_spec_m = X_spec_m - y_spec_m
+
+        if is_hp3 == True:
+            ins_root, vocal_root = vocal_root, ins_root
+
+        if ins_root is not None:
+            if self.data["high_end_process"].startswith("mirroring"):
+                input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], y_spec_m, input_high_end, self.mp)
+                wav_instrument = spec_utils.cmb_spectrogram_to_wave(
+                    y_spec_m, self.mp, input_high_end_h, input_high_end_
+                )
+            else:
+                wav_instrument = spec_utils.cmb_spectrogram_to_wave(y_spec_m, self.mp)
+            logger.info("%s instruments done" % name)
+            if is_hp3 == True:
+                head = "vocal_"
+            else:
+                head = "instrument_"
+            if format in ["wav", "flac"]:
+                sf.write(
+                    os.path.join(
+                        ins_root,
+                        head + "{}_{}.{}".format(name, self.data["agg"], format),
+                    ),
+                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )  #
+            else:
+                path = os.path.join(ins_root, head + "{}_{}.wav".format(name, self.data["agg"]))
+                sf.write(
+                    path,
+                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )
+                if os.path.exists(path):
+                    opt_format_path = path[:-4] + ".%s" % format
+                    cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path)
+                    print(cmd)
+                    os.system(cmd)
+                    if os.path.exists(opt_format_path):
+                        try:
+                            os.remove(path)
+                        except:
+                            pass
+        if vocal_root is not None:
+            if is_hp3 == True:
+                head = "instrument_"
+            else:
+                head = "vocal_"
+            if self.data["high_end_process"].startswith("mirroring"):
+                input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], v_spec_m, input_high_end, self.mp)
+                wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp, input_high_end_h, input_high_end_)
+            else:
+                wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp)
+            logger.info("%s vocals done" % name)
+            if format in ["wav", "flac"]:
+                sf.write(
+                    os.path.join(
+                        vocal_root,
+                        head + "{}_{}.{}".format(name, self.data["agg"], format),
+                    ),
+                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )
+            else:
+                path = os.path.join(vocal_root, head + "{}_{}.wav".format(name, self.data["agg"]))
+                sf.write(
+                    path,
+                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )
+                if os.path.exists(path):
+                    opt_format_path = path[:-4] + ".%s" % format
+                    cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path)
+                    print(cmd)
+                    os.system(cmd)
+                    if os.path.exists(opt_format_path):
+                        try:
+                            os.remove(path)
+                        except:
+                            pass
+
+
+class AudioPreDeEcho:
+    def __init__(self, agg, model_path, device, is_half, tta=False):
+        self.model_path = model_path
+        self.device = device
+        self.data = {
+            # Processing Options
+            "postprocess": False,
+            "tta": tta,
+            # Constants
+            "window_size": 512,
+            "agg": agg,
+            "high_end_process": "mirroring",
+        }
+        mp = ModelParameters("%s/lib/lib_v5/modelparams/4band_v3.json" % parent_directory)
+        nout = 64 if "DeReverb" in model_path else 48
+        model = CascadedNet(mp.param["bins"] * 2, nout)
+        cpk = torch.load(model_path, map_location="cpu")
+        model.load_state_dict(cpk)
+        model.eval()
+        if is_half:
+            model = model.half().to(device)
+        else:
+            model = model.to(device)
+
+        self.mp = mp
+        self.model = model
+
+    def _path_audio_(
+        self, music_file, vocal_root=None, ins_root=None, format="flac", is_hp3=False
+    ):  # 3个VR模型vocal和ins是反的
+        if ins_root is None and vocal_root is None:
+            return "No save root."
+        name = os.path.basename(music_file)
+        if ins_root is not None:
+            os.makedirs(ins_root, exist_ok=True)
+        if vocal_root is not None:
+            os.makedirs(vocal_root, exist_ok=True)
+        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
+        bands_n = len(self.mp.param["band"])
+        # print(bands_n)
+        for d in range(bands_n, 0, -1):
+            bp = self.mp.param["band"][d]
+            if d == bands_n:  # high-end band
+                (
+                    X_wave[d],
+                    _,
+                ) = librosa.core.load(  # 理论上librosa读取可能对某些音频有bug，应该上ffmpeg读取，但是太麻烦了弃坑
+                    music_file,
+                    sr=bp["sr"],
+                    mono=False,
+                    dtype=np.float32,
+                    res_type=bp["res_type"],
+                )
+                if X_wave[d].ndim == 1:
+                    X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])
+            else:  # lower bands
+                X_wave[d] = librosa.core.resample(
+                    X_wave[d + 1],
+                    orig_sr=self.mp.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=bp["res_type"],
+                )
+            # Stft of wave source
+            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(
+                X_wave[d],
+                bp["hl"],
+                bp["n_fft"],
+                self.mp.param["mid_side"],
+                self.mp.param["mid_side_b2"],
+                self.mp.param["reverse"],
+            )
+            # pdb.set_trace()
+            if d == bands_n and self.data["high_end_process"] != "none":
+                input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
+                    self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]
+                )
+                input_high_end = X_spec_s[d][:, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :]
+
+        X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)
+        aggresive_set = float(self.data["agg"] / 100)
+        aggressiveness = {
+            "value": aggresive_set,
+            "split_bin": self.mp.param["band"][1]["crop_stop"],
+        }
+        with torch.no_grad():
+            pred, X_mag, X_phase = inference(X_spec_m, self.device, self.model, aggressiveness, self.data)
+        # Postprocess
+        if self.data["postprocess"]:
+            pred_inv = np.clip(X_mag - pred, 0, np.inf)
+            pred = spec_utils.mask_silence(pred, pred_inv)
+        y_spec_m = pred * X_phase
+        v_spec_m = X_spec_m - y_spec_m
+
+        if ins_root is not None:
+            if self.data["high_end_process"].startswith("mirroring"):
+                input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], y_spec_m, input_high_end, self.mp)
+                wav_instrument = spec_utils.cmb_spectrogram_to_wave(
+                    y_spec_m, self.mp, input_high_end_h, input_high_end_
+                )
+            else:
+                wav_instrument = spec_utils.cmb_spectrogram_to_wave(y_spec_m, self.mp)
+            logger.info("%s instruments done" % name)
+            if format in ["wav", "flac"]:
+                sf.write(
+                    os.path.join(
+                        ins_root,
+                        "vocal_{}_{}.{}".format(name, self.data["agg"], format),
+                    ),
+                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )  #
+            else:
+                path = os.path.join(ins_root, "vocal_{}_{}.wav".format(name, self.data["agg"]))
+                sf.write(
+                    path,
+                    (np.array(wav_instrument) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )
+                if os.path.exists(path):
+                    opt_format_path = path[:-4] + ".%s" % format
+                    cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path)
+                    print(cmd)
+                    os.system(cmd)
+                    if os.path.exists(opt_format_path):
+                        try:
+                            os.remove(path)
+                        except:
+                            pass
+        if vocal_root is not None:
+            if self.data["high_end_process"].startswith("mirroring"):
+                input_high_end_ = spec_utils.mirroring(self.data["high_end_process"], v_spec_m, input_high_end, self.mp)
+                wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp, input_high_end_h, input_high_end_)
+            else:
+                wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp)
+            logger.info("%s vocals done" % name)
+            if format in ["wav", "flac"]:
+                sf.write(
+                    os.path.join(
+                        vocal_root,
+                        "instrument_{}_{}.{}".format(name, self.data["agg"], format),
+                    ),
+                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )
+            else:
+                path = os.path.join(vocal_root, "instrument_{}_{}.wav".format(name, self.data["agg"]))
+                sf.write(
+                    path,
+                    (np.array(wav_vocals) * 32768).astype("int16"),
+                    self.mp.param["sr"],
+                )
+                if os.path.exists(path):
+                    opt_format_path = path[:-4] + ".%s" % format
+                    cmd = 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path, opt_format_path)
+                    print(cmd)
+                    os.system(cmd)
+                    if os.path.exists(opt_format_path):
+                        try:
+                            os.remove(path)
+                        except:
+                            pass

tools/uvr5/webui.py

@@ -0,0 +1,133 @@
+import logging
+import os
+import traceback
+
+import ffmpeg
+import torch
+
+from configs.config import Config, IS_GPU
+from tools.uvr5.bsroformer import Roformer_Loader
+from tools.uvr5.mdxnet import MDXNetDereverb
+from tools.uvr5.vr import AudioPre, AudioPreDeEcho
+from i18n.i18n import I18nAuto
+
+
+logger = logging.getLogger(__name__)
+i18n = I18nAuto()
+config = Config()
+weight_uvr5_root = os.getenv("weight_uvr5_root", "assets/uvr5_weights")
+
+
+def clean_path(path):
+    path = path or ""
+    if path.endswith(("\\", "/")):
+        path = path[:-1]
+    return path.replace("/", os.sep).replace("\\", os.sep).strip(" '\n\"\u202a")
+
+
+def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format0):
+    infos = []
+    try:
+        inp_root = clean_path(inp_root)
+        save_root_vocal = clean_path(save_root_vocal)
+        save_root_ins = clean_path(save_root_ins)
+        is_hp3 = "HP3" in model_name
+        if model_name == "onnx_dereverb_By_FoxJoy":
+            if config.dml:
+                providers = ["DmlExecutionProvider", "CPUExecutionProvider"]
+            elif IS_GPU:
+                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            else:
+                providers = ["CPUExecutionProvider"]
+            pre_fun = MDXNetDereverb(15, providers)
+        elif "roformer" in model_name.lower():
+            pre_fun = Roformer_Loader(
+                model_path=os.path.join(weight_uvr5_root, model_name + ".ckpt"),
+                config_path=os.path.join(weight_uvr5_root, model_name + ".yaml"),
+                device=config.device,
+                is_half=config.is_half,
+            )
+            if not os.path.exists(
+                os.path.join(weight_uvr5_root, model_name + ".yaml")
+            ):
+                infos.append(i18n("未找到Roformer模型配置文件，正在使用内置默认配置"))
+                yield "\n".join(infos)
+        else:
+            func = AudioPre if "DeEcho" not in model_name else AudioPreDeEcho
+            pre_fun = func(
+                agg=int(agg),
+                model_path=os.path.join(weight_uvr5_root, model_name + ".pth"),
+                device=config.device,
+                is_half=config.is_half,
+            )
+        if inp_root:
+            paths = [os.path.join(inp_root, name) for name in os.listdir(inp_root)]
+        else:
+            paths = [path.name for path in (paths or [])]
+        for path in paths:
+            inp_path = os.path.join(inp_root, path)
+            if not os.path.isfile(inp_path):
+                continue
+            need_reformat = True
+            done = False
+            try:
+                info = ffmpeg.probe(inp_path, cmd="ffprobe")
+                if (
+                    info["streams"][0]["channels"] == 2
+                    and info["streams"][0]["sample_rate"] == "44100"
+                ):
+                    need_reformat = False
+                    pre_fun._path_audio_(
+                        inp_path,
+                        save_root_ins,
+                        save_root_vocal,
+                        format0,
+                        is_hp3,
+                    )
+                    done = True
+            except:
+                traceback.print_exc()
+            if need_reformat:
+                tmp_path = "%s/%s.reformatted.wav" % (
+                    os.environ["TEMP"],
+                    os.path.basename(inp_path),
+                )
+                os.system(
+                    'ffmpeg -i "%s" -vn -acodec pcm_s16le -ac 2 -ar 44100 "%s" -y'
+                    % (inp_path, tmp_path)
+                )
+                inp_path = tmp_path
+            try:
+                if not done:
+                    pre_fun._path_audio_(
+                        inp_path,
+                        save_root_ins,
+                        save_root_vocal,
+                        format0,
+                        is_hp3,
+                    )
+                infos.append(i18n("%s → 成功") % os.path.basename(inp_path))
+                yield "\n".join(infos)
+            except Exception:
+                infos.append(
+                    "%s → %s\n%s"
+                    % (os.path.basename(inp_path), i18n("失败"), traceback.format_exc())
+                )
+                yield "\n".join(infos)
+    except Exception:
+        infos.append("%s\n%s" % (i18n("失败"), traceback.format_exc()))
+        yield "\n".join(infos)
+    finally:
+        try:
+            if model_name == "onnx_dereverb_By_FoxJoy":
+                del pre_fun.pred.model
+                del pre_fun.pred.model_
+            else:
+                del pre_fun.model
+                del pre_fun
+        except:
+            traceback.print_exc()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            logger.info("Executed torch.cuda.empty_cache()")
+    yield "\n".join(infos)