mosibi/RVC_HFv2

1# Copyright (c) Facebook, Inc. and its affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import math8 9import julius10from torch import nn11 12from .utils import capture_init, center_trim13 14 15class BLSTM(nn.Module):16    def __init__(self, dim, layers=1):17        super().__init__()18        self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)19        self.linear = nn.Linear(2 * dim, dim)20 21    def forward(self, x):22        x = x.permute(2, 0, 1)23        x = self.lstm(x)[0]24        x = self.linear(x)25        x = x.permute(1, 2, 0)26        return x27 28 29def rescale_conv(conv, reference):30    std = conv.weight.std().detach()31    scale = (std / reference)**0.532    conv.weight.data /= scale33    if conv.bias is not None:34        conv.bias.data /= scale35 36 37def rescale_module(module, reference):38    for sub in module.modules():39        if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d)):40            rescale_conv(sub, reference)41 42 43class Demucs(nn.Module):44    @capture_init45    def __init__(self,46                 sources,47                 audio_channels=2,48                 channels=64,49                 depth=6,50                 rewrite=True,51                 glu=True,52                 rescale=0.1,53                 resample=True,54                 kernel_size=8,55                 stride=4,56                 growth=2.,57                 lstm_layers=2,58                 context=3,59                 normalize=False,60                 samplerate=44100,61                 segment_length=4 * 10 * 44100):62        """63        Args:64            sources (list[str]): list of source names65            audio_channels (int): stereo or mono66            channels (int): first convolution channels67            depth (int): number of encoder/decoder layers68            rewrite (bool): add 1x1 convolution to each encoder layer69                and a convolution to each decoder layer.70                For the decoder layer, `context` gives the kernel size.71            glu (bool): use glu instead of ReLU72            resample_input (bool): upsample x2 the input and downsample /2 the output.73            rescale (int): rescale initial weights of convolutions74                to get their standard deviation closer to `rescale`75            kernel_size (int): kernel size for convolutions76            stride (int): stride for convolutions77            growth (float): multiply (resp divide) number of channels by that78                for each layer of the encoder (resp decoder)79            lstm_layers (int): number of lstm layers, 0 = no lstm80            context (int): kernel size of the convolution in the81                decoder before the transposed convolution. If > 1,82                will provide some context from neighboring time83                steps.84            samplerate (int): stored as meta information for easing85                future evaluations of the model.86            segment_length (int): stored as meta information for easing87                future evaluations of the model. Length of the segments on which88                the model was trained.89        """90 91        super().__init__()92        self.audio_channels = audio_channels93        self.sources = sources94        self.kernel_size = kernel_size95        self.context = context96        self.stride = stride97        self.depth = depth98        self.resample = resample99        self.channels = channels100        self.normalize = normalize101        self.samplerate = samplerate102        self.segment_length = segment_length103 104        self.encoder = nn.ModuleList()105        self.decoder = nn.ModuleList()106 107        if glu:108            activation = nn.GLU(dim=1)109            ch_scale = 2110        else:111            activation = nn.ReLU()112            ch_scale = 1113        in_channels = audio_channels114        for index in range(depth):115            encode = []116            encode += [nn.Conv1d(in_channels, channels, kernel_size, stride), nn.ReLU()]117            if rewrite:118                encode += [nn.Conv1d(channels, ch_scale * channels, 1), activation]119            self.encoder.append(nn.Sequential(*encode))120 121            decode = []122            if index > 0:123                out_channels = in_channels124            else:125                out_channels = len(self.sources) * audio_channels126            if rewrite:127                decode += [nn.Conv1d(channels, ch_scale * channels, context), activation]128            decode += [nn.ConvTranspose1d(channels, out_channels, kernel_size, stride)]129            if index > 0:130                decode.append(nn.ReLU())131            self.decoder.insert(0, nn.Sequential(*decode))132            in_channels = channels133            channels = int(growth * channels)134 135        channels = in_channels136 137        if lstm_layers:138            self.lstm = BLSTM(channels, lstm_layers)139        else:140            self.lstm = None141 142        if rescale:143            rescale_module(self, reference=rescale)144 145    def valid_length(self, length):146        """147        Return the nearest valid length to use with the model so that148        there is no time steps left over in a convolutions, e.g. for all149        layers, size of the input - kernel_size % stride = 0.150 151        If the mixture has a valid length, the estimated sources152        will have exactly the same length when context = 1. If context > 1,153        the two signals can be center trimmed to match.154 155        For training, extracts should have a valid length.For evaluation156        on full tracks we recommend passing `pad = True` to :method:`forward`.157        """158        if self.resample:159            length *= 2160        for _ in range(self.depth):161            length = math.ceil((length - self.kernel_size) / self.stride) + 1162            length = max(1, length)163            length += self.context - 1164        for _ in range(self.depth):165            length = (length - 1) * self.stride + self.kernel_size166 167        if self.resample:168            length = math.ceil(length / 2)169        return int(length)170 171    def forward(self, mix):172        x = mix173 174        if self.normalize:175            mono = mix.mean(dim=1, keepdim=True)176            mean = mono.mean(dim=-1, keepdim=True)177            std = mono.std(dim=-1, keepdim=True)178        else:179            mean = 0180            std = 1181 182        x = (x - mean) / (1e-5 + std)183 184        if self.resample:185            x = julius.resample_frac(x, 1, 2)186 187        saved = []188        for encode in self.encoder:189            x = encode(x)190            saved.append(x)191        if self.lstm:192            x = self.lstm(x)193        for decode in self.decoder:194            skip = center_trim(saved.pop(-1), x)195            x = x + skip196            x = decode(x)197 198        if self.resample:199            x = julius.resample_frac(x, 2, 1)200        x = x * std + mean201        x = x.view(x.size(0), len(self.sources), self.audio_channels, x.size(-1))202        return x203