fix

app.py

@@ -14,10 +14,9 @@ import gradio as gr
 import librosa
 import numpy as np
 import soundfile
-from spleeter.separator import Separator
 from pydub import AudioSegment
 import uuid
-
+from torchspleeter import split_to_parts
 import logging
 
 logging.getLogger('numba').setLevel(logging.WARNING)
@@ -84,11 +83,11 @@ model = SynthesizerInfer(
 load_svc_model("vits_pretrain/sovits5.0-48k-debug.pth", model)
 model.eval()
 model.to(device)
-separator = Separator('spleeter:2stems')
 whisper_model = whisper.inference.load_model(os.path.join("whisper_pretrain", "medium.pt"))
 
+
 # warm up
-separator.separate_to_file('warm.wav', '/tmp/warm')
+# separator.separate_to_file('warm.wav', '/tmp/warm')
 
 
 def svc_change(argswave, argsspk):
@@ -196,7 +195,8 @@ def svc_main(sid, input_audio):
     soundfile.write(input_audio_tmp_file, audio, sampling_rate, format="wav")
     if not os.path.exists(tmpfile_path):
         os.makedirs(tmpfile_path)
-    separator.separate_to_file(input_audio_tmp_file, tmpfile_path)
+
+    split_to_parts(input_audio_tmp_file, tmpfile_path, models='torchspleeter/checkpoints/2stems/testcheckpoint1.ckpt')
 
     curr_tmp_path = os.path.join(tmpfile_path, os.path.splitext(input_audio_tmp_file)[0])
     vocals_filepath = os.path.join(curr_tmp_path, 'vocals.wav')
@@ -346,8 +346,8 @@ def main():
 
             app.queue(max_size=3, api_open=False).launch()
     except KeyboardInterrupt:
-        separator._get_session().close()
         app.close()
+        separator._get_session().close()
         sys.exit(0)
 
 

torchspleeter/__init__.py

@@ -0,0 +1,4 @@
+from torchspleeter.command_interface import *
+
+
+version="0.1.5"

torchspleeter/checkpoints/2stems/testcheckpoint0.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:385ea3372c6a3ceee01f6ded5504bb7ee1e9f0101950ae58869dc18382deb75c
+size 59050239

torchspleeter/checkpoints/2stems/testcheckpoint1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e4d6ede1ecad091468550773e77934aac3f1e039c0697fc9039aba9b935e344
+size 59033471

torchspleeter/command_interface.py

@@ -0,0 +1,98 @@
+"""
+
+This provides an interface to interact with the spleeter system on
+
+
+"""
+
+import os
+from torchspleeter.estimator import Estimator
+import argparse
+import uuid
+import numpy as np
+import librosa
+import soundfile
+import torch
+import pydub
+import os
+import shutil
+
+
+def split_to_parts(inputaudiofile, outputdir, instruments=2, models=[]):
+    filedata = pydub.AudioSegment.from_file(inputaudiofile)
+    sr = filedata.frame_rate
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # es = Estimator(2, './checkpoints/2stems/model').to(device)
+    # es = Estimator(2, ['./checkpoints/2stems/testcheckpoint0.ckpt','./checkpoints/2stems/testcheckpoint1.ckpt']).to(device)
+    es = Estimator()
+    es.eval()
+
+    # load wav audio
+    testaudiofile = inputaudiofile
+    channels = filedata.channels
+    mono_selection = False
+    if not os.path.exists(outputdir):
+        os.makedirs(outputdir, exist_ok=True)
+    if channels == 1:
+        mono_selection = True
+        multichannel = pydub.AudioSegment.from_mono_audiosegments(filedata, filedata)
+        os.makedirs(os.path.join(outputdir, 'tmp'), exist_ok=True)
+        testaudiofile = os.path.join(outputdir, "tmp" + str(uuid.uuid4()) + "." + testaudiofile.split('.')[-1])
+        # testaudiofile=testaudiofile.split('.')[0]+"-stereo."+testaudiofile.split('.')[-1]
+        multichannel.export(out_f=testaudiofile, format=testaudiofile.split('.')[-1])
+    print(mono_selection)
+    print(channels)
+    wav, _ = librosa.load(testaudiofile, mono=False, res_type='kaiser_fast', sr=sr)
+    wav = torch.Tensor(wav).to(device)
+    if mono_selection:
+        shutil.rmtree(os.path.join(outputdir, "tmp"))
+        # os.remove(testaudiofile)
+    wavs = es.separate(wav)
+    outputname = str(uuid.uuid4())
+    returnarray = []
+    for i in range(len(wavs)):
+        finaloutput = os.path.join(outputdir, outputname)
+        fname = '-out_{}.wav'.format(i)
+        fname = finaloutput + fname
+        print('Writing ', fname)
+        soundfile.write(fname, wavs[i].cpu().detach().numpy().T, sr, "PCM_16")
+        returnarray.append(fname)
+        # write_wav(fname, np.asfortranarray(wavs[i].squeeze().numpy()), sr)
+    return returnarray
+
+
+def get_file_list(dirname):
+    outputfilelist = []
+    for subdir, dirs, files in os.walk(dirname):
+        for file in files:
+            outputfilelist.append(os.path.join(subdir, file))
+
+    return outputfilelist
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='torchspleeter allows you to separate instrumentals from audio (vocals, instruments, background noise, etc) in a simple, cross platform manner')
+    parser.add_argument('-i', '--inputfile', help='Input Audio File to split into instrumentals', required=True)
+    parser.add_argument('-o', '--output', help='Output directory to deposit split audio', required=True)
+    parser.add_argument('-n', '--number', help="Number of instruments in the model (default 2)", required=False,
+                        default=2, type=int)
+    parser.add_argument('-m', '--modeldir',
+                        help="directory containing number of pre-converted torch compatible model components",
+                        required=False)
+    args = vars(parser.parse_args())
+    print(args)
+    if args['modeldir'] is not None:
+        modelfiles = get_file_list(args['modeldir'])
+        if len(modelfiles) != args['number']:
+            raise ValueError("You must have the same number of models as you do number of instruments!")
+    else:
+        args['modeldir'] = []
+    outputfiles = split_to_parts(args['inputfile'], args['output'], args['number'], args['modeldir'])
+    print("Your output files are:")
+    for item in outputfiles:
+        print(item)
+
+
+if __name__ == "__main__":
+    main()

torchspleeter/estimator.py

@@ -0,0 +1,160 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+import tqdm
+# from torchaudio.functional import istft
+
+from torchspleeter.unet import UNet
+#from .util import tf2pytorch
+
+import os
+dirname = os.path.dirname(__file__)
+defaultmodel0 = os.path.join(dirname, 'checkpoints/2stems/testcheckpoint0.ckpt')
+defaultmodel1 = os.path.join(dirname, 'checkpoints/2stems/testcheckpoint1.ckpt')
+
+def load_ckpt(model, ckpt):
+    state_dict = model.state_dict()
+    for k, v in ckpt.items():
+        if k in state_dict:
+            target_shape = state_dict[k].shape
+            assert target_shape == v.shape
+            state_dict.update({k: torch.from_numpy(v)})
+        else:
+            print('Ignore ', k)
+
+    model.load_state_dict(state_dict)
+    return model
+
+
+def pad_and_partition(tensor, T):
+    """
+    pads zero and partition tensor into segments of length T
+
+    Args:
+        tensor(Tensor): BxCxFxL
+
+    Returns:
+        tensor of size (B*[L/T] x C x F x T)
+    """
+    old_size = tensor.size(3)
+    new_size = math.ceil(old_size/T) * T
+    tensor = F.pad(tensor, [0, new_size - old_size])
+    [b, c, t, f] = tensor.shape
+    split = new_size // T
+    return torch.cat(torch.split(tensor, T, dim=3), dim=0)
+
+
+class Estimator(nn.Module):
+    def __init__(self, num_instrumments=2, checkpoint_path=None):
+        super(Estimator, self).__init__()
+        if checkpoint_path is None:
+            checkpoint_path=[defaultmodel0,defaultmodel1]
+        else:
+            if len(checkpoint_path)<1:
+                checkpoint_path=[defaultmodel0,defaultmodel1]
+        # stft config
+        self.F = 1024
+        self.T = 512
+        self.win_length = 4096
+        self.hop_length = 1024
+        self.win = nn.Parameter(
+            torch.hann_window(self.win_length),
+            requires_grad=False
+        )
+
+        ckpts=[]
+        if len(checkpoint_path) != num_instrumments:
+            raise ValueError("You must submit as many models as there are instruments!")
+        for ckpt_path in checkpoint_path:
+            ckpts.append(torch.load(ckpt_path))
+
+        #self.ckpts = ckpt #torch.load(checkpoint_path)#, num_instrumments)
+
+        #ckpts = #tf2pytorch(checkpoint_path, num_instrumments)
+
+        # filter
+        self.instruments = nn.ModuleList()
+        for i in range(num_instrumments):
+            print('Loading model for instrumment {}'.format(i))
+            net = UNet(2)
+            ckpt = ckpts[i]
+            net = load_ckpt(net, ckpt)
+            net.eval()  # change mode to eval
+            self.instruments.append(net)
+
+    def compute_stft(self, wav):
+        """
+        Computes stft feature from wav
+
+        Args:
+            wav (Tensor): B x L
+        """
+
+        stft = torch.stft(wav, n_fft=self.win_length, hop_length=self.hop_length, window=self.win,
+                          center=True, return_complex=False, pad_mode='constant')
+
+        # only keep freqs smaller than self.F
+        stft = stft[:, :self.F, :, :]
+        real = stft[:, :, :, 0]
+        im = stft[:, :, :, 1]
+        mag = torch.sqrt(real ** 2 + im ** 2)
+
+        return stft, mag
+
+    def inverse_stft(self, stft):
+        """Inverses stft to wave form"""
+
+        pad = self.win_length // 2 + 1 - stft.size(1)
+        stft = F.pad(stft, (0, 0, 0, 0, 0, pad))
+        wav = torch.istft(stft, self.win_length, hop_length=self.hop_length, center=True,
+                    window=self.win)
+        return wav.detach()
+
+    def separate(self, wav):
+        """
+        Separates stereo wav into different tracks corresponding to different instruments
+
+        Args:
+            wav (tensor): 2 x L
+        """
+
+        # stft - 2 X F x L x 2
+        # stft_mag - 2 X F x L
+        stft, stft_mag = self.compute_stft(wav)
+
+        L = stft.size(2)
+
+        # 1 x 2 x F x T
+        stft_mag = stft_mag.unsqueeze(-1).permute([3, 0, 1, 2])
+        stft_mag = pad_and_partition(stft_mag, self.T)  # B x 2 x F x T
+        stft_mag = stft_mag.transpose(2, 3)  # B x 2 x T x F
+
+        B = stft_mag.shape[0]
+
+        # compute instruments' mask
+        masks = []
+        for net in self.instruments:
+            mask = net(stft_mag)
+            masks.append(mask)
+
+        # compute denominator
+        mask_sum = sum([m ** 2 for m in masks])
+        mask_sum += 1e-10
+
+        wavs = []
+        for mask in tqdm.tqdm(masks):
+            mask = (mask ** 2 + 1e-10/2)/(mask_sum)
+            mask = mask.transpose(2, 3)  # B x 2 X F x T
+
+            mask = torch.cat(
+                torch.split(mask, 1, dim=0), dim=3)
+
+            mask = mask.squeeze(0)[:,:,:L].unsqueeze(-1) # 2 x F x L x 1
+            stft_masked = stft *  mask
+            wav_masked = self.inverse_stft(stft_masked)
+
+            wavs.append(wav_masked)
+
+        return wavs

torchspleeter/test/test_estimator.py

@@ -0,0 +1,45 @@
+import numpy as np
+import librosa
+import soundfile
+import torch
+import pydub
+import os
+from torchspleeter.estimator import Estimator
+dirname = os.path.dirname(__file__)
+testfilename = os.path.join(dirname, 'checkpoints/2stems/audio_example.mp3')
+
+if __name__ == '__main__':
+    sr = 44100
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    #es = Estimator(2, './checkpoints/2stems/model').to(device)
+    #es = Estimator(2, ['./checkpoints/2stems/testcheckpoint0.ckpt','./checkpoints/2stems/testcheckpoint1.ckpt']).to(device)
+    es=Estimator()
+    es.eval()
+
+    # load wav audio
+    testaudiofile=testfilename
+    filestats=pydub.AudioSegment.from_file(testaudiofile)
+    channels=filestats.channels
+    mono_selection=False
+    if channels==1:
+        mono_selection=True
+        multichannel=pydub.AudioSegment.from_mono_audiosegments(filestats,filestats)
+        testaudiofile=testaudiofile.split('.')[0]+"-stereo."+testaudiofile.split('.')[-1]
+        multichannel.export(out_f=testaudiofile,format=testaudiofile.split('.')[-1])
+    print(mono_selection)
+    print(channels)
+    wav, _ = librosa.load(testaudiofile, mono=False, res_type='kaiser_fast',sr=sr)
+    wav = torch.Tensor(wav).to(device)
+    if mono_selection:
+        os.remove(testaudiofile)
+
+
+    # normalize audio
+    # wav_torch = wav / (wav.max() + 1e-8)
+
+    wavs = es.separate(wav)
+    for i in range(len(wavs)):
+        fname = 'output/out_{}.wav'.format(i)
+        print('Writing ',fname)
+        soundfile.write(fname, wavs[i].cpu().detach().numpy().T, sr, "PCM_16")
+        # write_wav(fname, np.asfortranarray(wavs[i].squeeze().numpy()), sr)

torchspleeter/unet.py

@@ -0,0 +1,97 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+class CustomPad(nn.Module):
+    def __init__(self, padding_setting=(1, 2, 1, 2)):
+        super(CustomPad, self).__init__()
+        self.padding_setting = padding_setting
+
+    def forward(self, x):
+        return F.pad(x, self.padding_setting, "constant", 0)
+
+
+class CustomTransposedPad(nn.Module):
+    def __init__(self, padding_setting=(1, 2, 1, 2)):
+        super(CustomTransposedPad, self).__init__()
+        self.padding_setting = padding_setting
+
+    def forward(self, x):
+        l,r,t,b = self.padding_setting
+        return x[:,:,l:-r,t:-b]
+
+
+def down_block(in_filters, out_filters):
+    return nn.Sequential(CustomPad(),
+                         nn.Conv2d(in_filters, out_filters, kernel_size=5, stride=2,padding=0)), \
+           nn.Sequential(
+               nn.BatchNorm2d(out_filters, track_running_stats=True, eps=1e-3, momentum=0.01),
+               nn.LeakyReLU(0.2))
+
+
+def up_block(in_filters, out_filters, dropout=False):
+    layers = [
+        nn.ConvTranspose2d(in_filters, out_filters, kernel_size=5,stride=2),
+        CustomTransposedPad(),
+        nn.ReLU(),
+        nn.BatchNorm2d(out_filters, track_running_stats=True, eps=1e-3, momentum=0.01)
+    ]
+    if dropout:
+        layers.append(nn.Dropout(0.5))
+
+    return nn.Sequential(*layers)
+
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=2):
+        super(UNet, self).__init__()
+        self.down1_conv, self.down1_act = down_block(in_channels, 16)
+        self.down2_conv, self.down2_act = down_block(16, 32)
+        self.down3_conv, self.down3_act = down_block(32, 64)
+        self.down4_conv, self.down4_act = down_block(64, 128)
+        self.down5_conv, self.down5_act = down_block(128, 256)
+        self.down6_conv, self.down6_act = down_block(256, 512)
+
+        self.up1 = up_block(512, 256, dropout=True)
+        self.up2 = up_block(512, 128, dropout=True)
+        self.up3 = up_block(256, 64, dropout=True)
+        self.up4 = up_block(128, 32)
+        self.up5 = up_block(64, 16)
+        self.up6 = up_block(32, 1)
+        self.up7 = nn.Sequential(
+            nn.Conv2d(1, in_channels, kernel_size=4, dilation=2, padding=3),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        d1_conv = self.down1_conv(x)
+        d1 = self.down1_act(d1_conv)
+
+        d2_conv = self.down2_conv(d1)
+        d2 = self.down2_act(d2_conv)
+
+        d3_conv = self.down3_conv(d2)
+        d3 = self.down3_act(d3_conv)
+
+        d4_conv = self.down4_conv(d3)
+        d4 = self.down4_act(d4_conv)
+
+        d5_conv = self.down5_conv(d4)
+        d5 = self.down5_act(d5_conv)
+
+        d6_conv = self.down6_conv(d5)
+        d6 = self.down6_act(d6_conv)
+
+        u1 = self.up1(d6_conv)
+        u2 = self.up2(torch.cat([d5_conv, u1], axis=1))
+        u3 = self.up3(torch.cat([d4_conv, u2], axis=1))
+        u4 = self.up4(torch.cat([d3_conv, u3], axis=1))
+        u5 = self.up5(torch.cat([d2_conv, u4], axis=1))
+        u6 = self.up6(torch.cat([d1_conv, u5], axis=1))
+        u7 = self.up7(u6)
+        return u7 * x
+
+
+if __name__ == '__main__':
+    net = UNet(14)
+    print(net(torch.rand(1, 14, 20, 48)).shape)