Spaces:
Runtime error
Runtime error
File size: 3,991 Bytes
df578bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import argparse
import os
import data.utils
import model.utils as model_utils
from test import predict_song
from model.waveunet import Waveunet
def main(args):
# MODEL
num_features = [args.features*i for i in range(1, args.levels+1)] if args.feature_growth == "add" else \
[args.features*2**i for i in range(0, args.levels)]
target_outputs = int(args.output_size * args.sr)
model = Waveunet(args.channels, num_features, args.channels, args.instruments, kernel_size=args.kernel_size,
target_output_size=target_outputs, depth=args.depth, strides=args.strides,
conv_type=args.conv_type, res=args.res, separate=args.separate)
if args.cuda:
model = model_utils.DataParallel(model)
print("move model to gpu")
model.cuda()
print("Loading model from checkpoint " + str(args.load_model))
state = model_utils.load_model(model, None, args.load_model, args.cuda)
print('Step', state['step'])
preds = predict_song(args, args.input, model)
output_folder = os.path.dirname(args.input) if args.output is None else args.output
for inst in preds.keys():
data.utils.write_wav(os.path.join(output_folder, os.path.basename(args.input) + "_" + inst + ".wav"), preds[inst], args.sr)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--instruments', type=str, nargs='+', default=["bass", "drums", "other", "vocals"],
help="List of instruments to separate (default: \"bass drums other vocals\")")
parser.add_argument('--cuda', action='store_true',
help='Use CUDA (default: False)')
parser.add_argument('--features', type=int, default=32,
help='Number of feature channels per layer')
parser.add_argument('--load_model', type=str, default='checkpoints/waveunet/model',
help='Reload a previously trained model')
parser.add_argument('--batch_size', type=int, default=4,
help="Batch size")
parser.add_argument('--levels', type=int, default=6,
help="Number of DS/US blocks")
parser.add_argument('--depth', type=int, default=1,
help="Number of convs per block")
parser.add_argument('--sr', type=int, default=44100,
help="Sampling rate")
parser.add_argument('--channels', type=int, default=2,
help="Number of input audio channels")
parser.add_argument('--kernel_size', type=int, default=5,
help="Filter width of kernels. Has to be an odd number")
parser.add_argument('--output_size', type=float, default=2.0,
help="Output duration")
parser.add_argument('--strides', type=int, default=4,
help="Strides in Waveunet")
parser.add_argument('--conv_type', type=str, default="gn",
help="Type of convolution (normal, BN-normalised, GN-normalised): normal/bn/gn")
parser.add_argument('--res', type=str, default="fixed",
help="Resampling strategy: fixed sinc-based lowpass filtering or learned conv layer: fixed/learned")
parser.add_argument('--separate', type=int, default=1,
help="Train separate model for each source (1) or only one (0)")
parser.add_argument('--feature_growth', type=str, default="double",
help="How the features in each layer should grow, either (add) the initial number of features each time, or multiply by 2 (double)")
parser.add_argument('--input', type=str, default=os.path.join("audio_examples", "Cristina Vane - So Easy", "mix.mp3"),
help="Path to input mixture to be separated")
parser.add_argument('--output', type=str, default=None, help="Output path (same folder as input path if not set)")
args = parser.parse_args()
main(args)
|