fix

Dockerfile

@@ -13,7 +13,6 @@ ENV PYTHONUNBUFFERED=1
 
 COPY requirements.txt requirements.txt
 RUN pip3 install --no-cache-dir -r requirements.txt
-RUN pip3 install --no-cache-dir --no-dependencies spleeter
 
 RUN mkdir -p $HOME/app/cache/transformers
 RUN chmod 777 $HOME/app/cache/transformers

app.py

@@ -16,7 +16,9 @@ import numpy as np
 import soundfile
 from pydub import AudioSegment
 import uuid
-from torchspleeter import split_to_parts
+from torchspleeter.utils import sound_split
+from torchspleeter.splitter import Splitter
+
 import logging
 
 logging.getLogger('numba').setLevel(logging.WARNING)
@@ -84,6 +86,7 @@ load_svc_model("vits_pretrain/sovits5.0-48k-debug.pth", model)
 model.eval()
 model.to(device)
 whisper_model = whisper.inference.load_model(os.path.join("whisper_pretrain", "medium.pt"))
+splitter_model = Splitter.from_pretrained(os.path.join("torchspleeter/models/2stems", "spleeter.pth")).to(device).eval()
 
 
 # warm up
@@ -187,7 +190,7 @@ def svc_main(sid, input_audio):
     uuid_value = uuid.uuid4()
     uuid_string = str(uuid_value)
     input_audio_tmp_file = f'{uuid_string}.wav'
-    tmpfile_path = '/tmp'
+    tmpfile_path = f'/tmp/{uuid_string}'
 
     #
     # prediction = separator.separate(audio)
@@ -196,9 +199,9 @@ def svc_main(sid, input_audio):
     if not os.path.exists(tmpfile_path):
         os.makedirs(tmpfile_path)
 
-    split_to_parts(input_audio_tmp_file, tmpfile_path)
+    sound_split(splitter_model, input_audio_tmp_file, tmpfile_path)
 
-    curr_tmp_path = os.path.join(tmpfile_path, os.path.splitext(input_audio_tmp_file)[0])
+    curr_tmp_path = tmpfile_path
     vocals_filepath = os.path.join(curr_tmp_path, 'vocals.wav')
     accompaniment_filepath = os.path.join(curr_tmp_path, 'accompaniment.wav')
 
@@ -344,10 +347,9 @@ def main():
                         vc_search.click(auto_search, [vc_input2], [vc_output1, vc_input3])
                         vc_submit.click(svc_main, [sid, vc_input3], [vc_output1, vc_output2])
 
-            app.queue(max_size=3, api_open=False).launch()
+            app.queue(max_size=3, api_open=False).launch(server_port=6006)
     except KeyboardInterrupt:
         app.close()
-        separator._get_session().close()
         sys.exit(0)
 
 

torchspleeter/.dockerignore

@@ -0,0 +1,5 @@
+data/
+models/
+output/
+docs/
+.vscode/

torchspleeter/.gitignore

@@ -0,0 +1,135 @@
+models/
+output/
+experiments/
+data/
+.vscode/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

torchspleeter/__init__.py

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

torchspleeter/checkpoints/2stems/testcheckpoint0.ckpt

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

torchspleeter/checkpoints/2stems/testcheckpoint1.ckpt

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

torchspleeter/command_interface.py

@@ -1,98 +0,0 @@
-"""
-
-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/data.py

@@ -0,0 +1,37 @@
+import random
+from pathlib import Path
+from typing import List
+
+import musdb
+import torch
+from torch.utils.data import Dataset
+
+
+class MusdbDataset(Dataset):
+    def __init__(
+        self, root="data/musdb18-wav", is_train: bool = True, targets: List[str] = None
+    ) -> None:
+        super().__init__()
+        root = Path(root)
+        assert root.exists(), f"Path does not exist: {root}"
+        self.mus = musdb.DB(
+            root=root,
+            subsets=["train" if is_train else "test"],
+            is_wav=True,
+        )
+        self.targets = [s for s in targets] if targets else ["vocals", "accompaniment"]
+
+    def __len__(self) -> int:
+        return len(self.mus)
+
+    def __getitem__(self, index):
+        track = self.mus.tracks[index]
+        track.chunk_duration = 5.0
+        track.chunk_start = random.uniform(0, track.duration - track.chunk_duration)
+        x_wav = torch.torch.tensor(track.audio.T, dtype=torch.float32)
+        y_target_wavs = {
+            name: torch.tensor(track.targets[name].audio.T, dtype=torch.float32)
+            for name in self.targets
+        }
+        # original audio (x) and stems (y == targets)
+        return x_wav, y_target_wavs

torchspleeter/estimator.py

@@ -1,163 +0,0 @@
-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')
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-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, device=device),
-            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/requirements.txt

@@ -0,0 +1,8 @@
+accelerate==0.14.0
+librosa
+musdb==0.4.0
+soundfile
+tensorflow==2.10.0
+torch==1.13.0
+torchaudio
+typer

torchspleeter/run.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import torch
+import typer
+from utils import sound_split
+
+app = typer.Typer()
+
+
+@app.command()
+def model_summary() -> None:
+    from unet import UNet
+
+    net = UNet()
+    print(net)
+
+
+@app.command()
+def test() -> None:
+    from unet import UNet
+
+    batch_size = 5
+    n_channels = 2
+    x = torch.randn(batch_size, n_channels, 512, 128)
+    print(x.shape)
+    net = UNet(in_channels=n_channels)
+    y = net.forward(x)
+    print(y.shape)
+
+
+@app.command()
+def split(
+    model_path: str = "models/2stems/model",
+    input: str = "data/audio_example.mp3",
+    output_dir: str = "output",
+    offset: float = 0,
+    duration: float = 30,
+    write_src: bool = False,
+) -> None:
+
+    from splitter import Splitter
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    splitter_model = Splitter.from_pretrained(model_path).to(device).eval()
+
+    sound_split(splitter_model, input, output_dir, write_src)
+
+
+if __name__ == "__main__":
+    app()

torchspleeter/splitter.py

@@ -0,0 +1,140 @@
+import math
+from typing import Dict, List, Tuple
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .unet import UNet
+
+
+def batchify(tensor: Tensor, T: int) -> Tensor:
+    """
+    partition tensor into segments of length T, zero pad any ragged samples
+    Args:
+        tensor(Tensor): BxCxFxL
+    Returns:
+        tensor of size (B*[L/T] x C x F x T)
+    """
+    # Zero pad the original tensor to an even multiple of T
+    orig_size = tensor.size(-1)
+    new_size = math.ceil(orig_size / T) * T
+    tensor = F.pad(tensor, [0, new_size - orig_size])
+    # Partition the tensor into multiple samples of length T and stack them into a batch
+    return torch.cat(torch.split(tensor, T, dim=-1), dim=0)
+
+
+class Splitter(nn.Module):
+    def __init__(self, stem_names: List[str] = None):
+        super(Splitter, self).__init__()
+
+        # 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)
+
+        self.stems = nn.ModuleDict({'vocals': UNet(in_channels=2),
+                                    'accompaniment': UNet(in_channels=2)})
+
+    def compute_stft(self, wav: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        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: Tensor) -> Tensor:
+        """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 forward(self, wav: Tensor) -> Dict[str, Tensor]:
+        """
+        Separates stereo wav into different tracks (1 predicted track per stem)
+        Args:
+            wav (tensor): 2 x L
+        Returns:
+            masked stfts by track name
+        """
+        # stft - 2 X F x L x 2
+        # stft_mag - 2 X F x L
+        stft, stft_mag = self.compute_stft(wav.squeeze())
+
+        L = stft.size(2)
+
+        # 1 x 2 x F x T
+        stft_mag = stft_mag.unsqueeze(-1).permute([3, 0, 1, 2])
+        stft_mag = batchify(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
+
+        # compute stems' mask
+        masks = {name: net(stft_mag) for name, net in self.stems.items()}
+
+        # compute denominator
+        mask_sum = sum([m ** 2 for m in masks.values()])
+        mask_sum += 1e-10
+
+        def apply_mask(mask):
+            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
+            return stft_masked
+
+        return {name: apply_mask(m) for name, m in masks.items()}
+
+    def separate(self, wav: Tensor) -> Dict[str, Tensor]:
+        """
+        Separates stereo wav into different tracks (1 predicted track per stem)
+        Args:
+            wav (tensor): 2 x L
+        Returns:
+            wavs by track name
+        """
+
+        stft_masks = self.forward(wav)
+
+        return {
+            name: self.inverse_stft(stft_masked)
+            for name, stft_masked in stft_masks.items()
+        }
+
+    @classmethod
+    def from_pretrained(cls, model_path: str):
+        checkpoint = torch.load(model_path)
+        model = cls()
+        model.load_state_dict(checkpoint)
+        return model
+
+

torchspleeter/test/test_estimator.py

@@ -1,45 +0,0 @@
-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/tf2pytorch.py

@@ -0,0 +1,122 @@
+from typing import Dict
+import numpy as np
+
+import os
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+import tensorflow as tf
+
+
+def parse_int_or_default(s: str, default: int = 0) -> int:
+    try:
+        return int(s)
+    except:
+        return default
+
+
+def tf2pytorch(checkpoint_path: str) -> Dict:
+    init_vars = tf.train.list_variables(checkpoint_path)
+
+    tf_vars = {}
+    for name, _ in init_vars:
+        try:
+            # print('Loading TF Weight {} with shape {}'.format(name, shape))
+            data = tf.train.load_variable(checkpoint_path, name)
+            tf_vars[name] = data
+        except Exception as e:
+            print(f"Load error: {name}")
+            raise
+
+    layer_idxs = set(
+        [
+            parse_int_or_default(name.split("/")[0].split("_")[-1], default=0)
+            for name in tf_vars.keys()
+            if "conv2d_transpose" in name
+        ]
+    )
+
+    n_layers_per_unet = 6
+    n_layers_in_chkpt = max(layer_idxs) + 1
+    assert (
+        n_layers_in_chkpt % 6 == 0
+    ), f"expected multiple of {n_layers_per_unet}... ie: {n_layers_per_unet} layers per unet & 1 unet per stem"
+    n_stems = n_layers_in_chkpt // n_layers_per_unet
+
+    stem_names = {
+        2: ["vocals", "accompaniment"],
+        4: ["vocals", "drums", "bass", "other"],
+        5: ["vocals", "piano", "drums", "bass", "other"],
+    }.get(n_stems, [])
+
+    assert stem_names, f"Unsupported stem count: {n_stems}"
+
+    state_dict = {}
+    tf_idx_conv = 0
+    tf_idx_tconv = 0
+    tf_idx_bn = 0
+
+    for stem_name in stem_names:
+        # Encoder Blocks (Down sampling)
+        for layer_idx in range(n_layers_per_unet):
+            prefix = f"stems.{stem_name}.encoder_layers.{layer_idx}"
+            conv_suffix = "" if tf_idx_conv == 0 else f"_{tf_idx_conv}"
+            bn_suffix = "" if tf_idx_bn == 0 else f"_{tf_idx_bn}"
+
+            state_dict[f"{prefix}.conv.weight"] = np.transpose(
+                tf_vars[f"conv2d{conv_suffix}/kernel"], (3, 2, 0, 1)
+            )
+            state_dict[f"{prefix}.conv.bias"] = tf_vars[f"conv2d{conv_suffix}/bias"]
+            tf_idx_conv += 1
+
+            state_dict[f"{prefix}.bn.weight"] = tf_vars[
+                f"batch_normalization{bn_suffix}/gamma"
+            ]
+            state_dict[f"{prefix}.bn.bias"] = tf_vars[
+                f"batch_normalization{bn_suffix}/beta"
+            ]
+            state_dict[f"{prefix}.bn.running_mean"] = tf_vars[
+                f"batch_normalization{bn_suffix}/moving_mean"
+            ]
+            state_dict[f"{prefix}.bn.running_var"] = tf_vars[
+                f"batch_normalization{bn_suffix}/moving_variance"
+            ]
+            tf_idx_bn += 1
+
+        # Decoder Blocks (Up sampling)
+        for layer_idx in range(n_layers_per_unet):
+            prefix = f"stems.{stem_name}.decoder_layers.{layer_idx}"
+            tconv_suffix = "" if tf_idx_tconv == 0 else f"_{tf_idx_tconv}"
+            bn_suffix = f"_{tf_idx_bn}"
+
+            state_dict[f"{prefix}.tconv.weight"] = np.transpose(
+                tf_vars[f"conv2d_transpose{tconv_suffix}/kernel"], (3, 2, 0, 1)
+            )
+            state_dict[f"{prefix}.tconv.bias"] = tf_vars[
+                f"conv2d_transpose{tconv_suffix}/bias"
+            ]
+            tf_idx_tconv += 1
+
+            state_dict[f"{prefix}.bn.weight"] = tf_vars[
+                f"batch_normalization{bn_suffix}/gamma"
+            ]
+            state_dict[f"{prefix}.bn.bias"] = tf_vars[
+                f"batch_normalization{bn_suffix}/beta"
+            ]
+            state_dict[f"{prefix}.bn.running_mean"] = tf_vars[
+                f"batch_normalization{bn_suffix}/moving_mean"
+            ]
+            state_dict[f"{prefix}.bn.running_var"] = tf_vars[
+                f"batch_normalization{bn_suffix}/moving_variance"
+            ]
+            tf_idx_bn += 1
+
+        # Final conv2d
+        state_dict[f"stems.{stem_name}.up_final.weight"] = np.transpose(
+            tf_vars[f"conv2d_{tf_idx_conv}/kernel"], (3, 2, 0, 1)
+        )
+        state_dict[f"stems.{stem_name}.up_final.bias"] = tf_vars[
+            f"conv2d_{tf_idx_conv}/bias"
+        ]
+        tf_idx_conv += 1
+
+    return state_dict

torchspleeter/train.py

@@ -0,0 +1,143 @@
+import os
+from datetime import datetime
+from pathlib import Path
+
+import torch
+import typer
+from accelerate import Accelerator
+from accelerate.utils import LoggerType
+from torch import Tensor
+from torch.optim import AdamW
+
+# from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from data import MusdbDataset
+from splitter import Splitter
+
+DISABLE_TQDM = os.environ.get("DISABLE_TQDM", False)
+app = typer.Typer(pretty_exceptions_show_locals=False)
+
+
+def spectrogram_loss(masked_target: Tensor, original: Tensor) -> Tensor:
+    """
+    masked_target (Tensor): a masked STFT generated by applying a net's
+        estimated mask for source S to the ground truth STFT for source S
+    original (Tensor): an original input mixture
+    """
+    square_difference = torch.square(masked_target - original)
+    loss_value = torch.mean(square_difference)
+    return loss_value
+
+
+@app.command()
+def train(
+    dataset: str = "data/musdb18-wav",
+    output_dir: str = None,
+    fp16: bool = False,
+    cpu: bool = True,
+    max_steps: int = 100,
+    num_train_epochs: int = 1,
+    per_device_train_batch_size: int = 1,
+    effective_batch_size: int = 4,
+    max_grad_norm: float = 0.0,
+) -> None:
+    if not output_dir:
+        now_str = datetime.now().strftime("%Y%m%d-%H%M%S")
+        output_dir = f"experiments/{now_str}"
+    output_dir = Path(output_dir)
+    logging_dir = output_dir / "tracker_logs"
+    accelerator = Accelerator(
+        fp16=fp16,
+        cpu=cpu,
+        logging_dir=logging_dir,
+        log_with=[LoggerType.TENSORBOARD],
+    )
+    accelerator.init_trackers(logging_dir / "run")
+
+    train_dataset = MusdbDataset(root=dataset, is_train=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        shuffle=True,
+        batch_size=per_device_train_batch_size,
+    )
+
+    model = Splitter(stem_names=[s for s in train_dataset.targets])
+    optimizer = AdamW(
+        model.parameters(),
+        lr=1e-3,
+        eps=1e-8,
+    )
+    model, optimizer, train_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader
+    )
+
+    num_train_steps = (
+        max_steps if max_steps > 0 else len(train_dataloader) * num_train_epochs
+    )
+    accelerator.print(f"Num train steps: {num_train_steps}")
+
+    step_batch_size = per_device_train_batch_size * accelerator.num_processes
+    gradient_accumulation_steps = max(
+        1,
+        effective_batch_size // step_batch_size,
+    )
+
+    accelerator.print(
+        f"Gradient Accumulation Steps: {gradient_accumulation_steps}\nEffective Batch Size: {gradient_accumulation_steps * step_batch_size}"
+    )
+    global_step = 0
+    while global_step < num_train_steps:
+        accelerator.wait_for_everyone()
+        # accelerator.print(f"global step: {global_step}")
+        # accelerator.print("running train...")
+        model.train()
+        batch_iterator = tqdm(
+            train_dataloader,
+            desc="Batch",
+            disable=((not accelerator.is_local_main_process) or DISABLE_TQDM),
+        )
+        for batch_idx, batch in enumerate(batch_iterator):
+            assert per_device_train_batch_size == 1, "For now limit to 1."
+            x_wav, y_target_wavs = batch
+            predictions = model(x_wav)
+            stem_losses = []
+            for name, masked_stft in predictions.items():
+                target_stft, _ = model.compute_stft(y_target_wavs[name].squeeze())
+                loss = spectrogram_loss(
+                    masked_target=masked_stft,
+                    original=target_stft,
+                )
+                stem_losses.append(loss)
+                accelerator.log({f"train-loss-{name}": 1.0 * loss}, step=global_step)
+
+            total_loss = (
+                torch.sum(torch.stack(stem_losses)) / gradient_accumulation_steps
+            )
+            accelerator.print(f"global step: {global_step}\tloss: {total_loss:.4f}")
+            accelerator.log({f"train-loss": 1.0 * total_loss}, step=global_step)
+            accelerator.backward(total_loss)
+            if (batch_idx + 1) % gradient_accumulation_steps == 0:
+                if max_grad_norm > 0:
+                    accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
+                optimizer.step()
+                optimizer.zero_grad()
+                global_step += 1
+
+    accelerator.wait_for_everyone()
+    accelerator.end_training()
+    accelerator.print(f"Saving model to {output_dir}...")
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_model.save_pretrained(
+        output_dir,
+        save_function=accelerator.save,
+        state_dict=accelerator.get_state_dict(model),
+    )
+
+    accelerator.wait_for_everyone()
+    accelerator.print("DONE!")
+
+
+if __name__ == "__main__":
+    app()

torchspleeter/unet.py

@@ -1,97 +1,97 @@
+from typing import Tuple
 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]
-
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=5, stride=(2, 2))
+        self.bn = nn.BatchNorm2d(
+            num_features=out_channels,
+            track_running_stats=True,
+            eps=1e-3,
+            momentum=0.01,
+        )
+        self.relu = nn.LeakyReLU(negative_slope=0.2)
 
-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 forward(self, input: Tensor) -> Tuple[Tensor, Tensor]:
+        down = self.conv(F.pad(input, (1, 2, 1, 2), "constant", 0))
+        return down, self.relu(self.bn(down))
 
 
-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))
+class DecoderBlock(nn.Module):
+    def __init__(
+        self, in_channels: int, out_channels: int, dropout_prob: float = 0.0
+    ) -> None:
+        super().__init__()
+        self.tconv = nn.ConvTranspose2d(
+            in_channels, out_channels, kernel_size=5, stride=2
+        )
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(
+            out_channels, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+        self.dropout = nn.Dropout(dropout_prob) if dropout_prob > 0 else nn.Identity()
 
-    return nn.Sequential(*layers)
+    def forward(self, input: Tensor) -> Tensor:
+        up = self.tconv(input)
+        # reverse padding
+        l, r, t, b = 1, 2, 1, 2
+        up = up[:, :, l:-r, t:-b]
+        return self.dropout(self.bn(self.relu(up)))
 
 
 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 __init__(
+        self,
+        n_layers: int = 6,
+        in_channels: int = 1,
+    ) -> None:
+        super().__init__()
+
+        # DownSample layers
+        down_set = [in_channels] + [2 ** (i + 4) for i in range(n_layers)]
+        self.encoder_layers = nn.ModuleList(
+            [
+                EncoderBlock(in_channels=in_ch, out_channels=out_ch)
+                for in_ch, out_ch in zip(down_set[:-1], down_set[1:])
+            ]
         )
 
-    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
-
+        # UpSample layers
+        up_set = [1] + [2 ** (i + 4) for i in range(n_layers)]
+        up_set.reverse()
+        self.decoder_layers = nn.ModuleList(
+            [
+                DecoderBlock(
+                    # doubled for concatenated inputs (skip connections)
+                    in_channels=in_ch if i == 0 else in_ch * 2,
+                    out_channels=out_ch,
+                    #   50 % dropout... first 3 layers only
+                    dropout_prob=0.5 if i < 3 else 0,
+                )
+                for i, (in_ch, out_ch) in enumerate(zip(up_set[:-1], up_set[1:]))
+            ]
+        )
 
-if __name__ == '__main__':
-    net = UNet(14)
-    print(net(torch.rand(1, 14, 20, 48)).shape)
+        # reconstruct the final mask same as the original channels
+        self.up_final = nn.Conv2d(1, in_channels, kernel_size=4, dilation=2, padding=3)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, input: Tensor) -> Tensor:
+        encoder_outputs_pre_act = []
+        x = input
+        for down in self.encoder_layers:
+            conv, x = down(x)
+            encoder_outputs_pre_act.append(conv)
+
+        for i, up in enumerate(self.decoder_layers):
+            if i == 0:
+                x = up(encoder_outputs_pre_act.pop())
+            else:
+                # merge skip connection
+                x = up(torch.concat([encoder_outputs_pre_act.pop(), x], axis=1))
+        mask = self.sigmoid(self.up_final(x))
+        return mask * input

torchspleeter/utils.py

@@ -0,0 +1,42 @@
+import librosa
+import soundfile
+from pathlib import Path
+
+import torch
+
+from .splitter import Splitter
+
+
+def sound_split(
+        model: Splitter,
+        input: str = "data/audio_example.mp3",
+        output_dir: str = "output",
+        write_src: bool = False,
+) -> None:
+    sr = 44100
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # splitter = Splitter.from_pretrained(model_path).to(device).eval()
+
+    # load wav audio
+    fpath_src = Path(input)
+    wav, _ = librosa.load(
+        fpath_src,
+        mono=False,
+        res_type="kaiser_fast",
+        sr=sr,
+    )
+    wav = torch.Tensor(wav).to(device)
+
+    # normalize audio
+    # wav_torch = wav / (wav.max() + 1e-8)
+
+    with torch.no_grad():
+        stems = model.separate(wav)
+
+    if write_src:
+        stems["input"] = wav
+    for name, stem in stems.items():
+        fpath_dst = Path(output_dir) / f"{name}.wav"
+        print(f"Writing {fpath_dst}")
+        fpath_dst.parent.mkdir(exist_ok=True)
+        soundfile.write(fpath_dst, stem.cpu().detach().numpy().T, sr, "PCM_16")

warm_up.py

@@ -1,5 +1,5 @@
 from spleeter.separator import Separator
 
-separator = Separator('spleeter:2stems')
+separator = Separator('spleeter.pth:2stems')
 # warm up
 separator.separate_to_file('warm.wav', '/tmp/warm')