File size: 6,490 Bytes
a8a84ee |
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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
# coding: utf-8
__author__ = 'Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/'
import argparse
import time
import librosa
from tqdm import tqdm
import sys
import os
import glob
import torch
import numpy as np
import soundfile as sf
import torch.nn as nn
# Using the embedded version of Python can also correctly import the utils module.
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
from utils import demix_track, demix_track_demucs, get_model_from_config
import warnings
warnings.filterwarnings("ignore")
def run_folder(model, args, config, device, verbose=False):
start_time = time.time()
model.eval()
all_mixtures_path = glob.glob(args.input_folder + '/*.*')
all_mixtures_path.sort()
print('Total files found: {}'.format(len(all_mixtures_path)))
instruments = config.training.instruments
if config.training.target_instrument is not None:
instruments = [config.training.target_instrument]
if not os.path.isdir(args.store_dir):
os.mkdir(args.store_dir)
if not verbose:
all_mixtures_path = tqdm(all_mixtures_path, desc="Total progress")
if args.disable_detailed_pbar:
detailed_pbar = False
else:
detailed_pbar = True
for path in all_mixtures_path:
print("Starting processing track: ", path)
if not verbose:
all_mixtures_path.set_postfix({'track': os.path.basename(path)})
try:
# mix, sr = sf.read(path)
mix, sr = librosa.load(path, sr=44100, mono=False)
except Exception as e:
print('Can read track: {}'.format(path))
print('Error message: {}'.format(str(e)))
continue
# Convert mono to stereo if needed
if len(mix.shape) == 1:
mix = np.stack([mix, mix], axis=0)
mix_orig = mix.copy()
if 'normalize' in config.inference:
if config.inference['normalize'] is True:
mono = mix.mean(0)
mean = mono.mean()
std = mono.std()
mix = (mix - mean) / std
mixture = torch.tensor(mix, dtype=torch.float32)
if args.model_type == 'htdemucs':
res = demix_track_demucs(config, model, mixture, device, pbar=detailed_pbar)
else:
res = demix_track(config, model, mixture, device, pbar=detailed_pbar)
for instr in instruments:
estimates = res[instr].T
if 'normalize' in config.inference:
if config.inference['normalize'] is True:
estimates = estimates * std + mean
file_name, _ = os.path.splitext(os.path.basename(path))
output_file = os.path.join(args.store_dir, f"{file_name}_{instr}.wav")
sf.write(output_file, estimates, sr, subtype = 'FLOAT')
# Output "instrumental", which is an inverse of 'vocals' (or first stem in list if 'vocals' absent)
if args.extract_instrumental:
file_name, _ = os.path.splitext(os.path.basename(path))
instrum_file_name = os.path.join(args.store_dir, f"{file_name}_instrumental.wav")
if 'vocals' in instruments:
estimates = res['vocals'].T
else:
estimates = res[instruments[0]].T
if 'normalize' in config.inference:
if config.inference['normalize'] is True:
estimates = estimates * std + mean
sf.write(instrum_file_name, mix_orig.T - estimates, sr, subtype = 'FLOAT')
time.sleep(1)
print("Elapsed time: {:.2f} sec".format(time.time() - start_time))
def proc_folder(args):
parser = argparse.ArgumentParser()
parser.add_argument("--model_type", type=str, default='mdx23c',
help="One of bandit, bandit_v2, bs_roformer, htdemucs, mdx23c, mel_band_roformer, scnet, scnet_unofficial, segm_models, swin_upernet, torchseg")
parser.add_argument("--config_path", type=str, help="path to config file")
parser.add_argument("--start_check_point", type=str, default='', help="Initial checkpoint to valid weights")
parser.add_argument("--input_folder", type=str, help="folder with mixtures to process")
parser.add_argument("--store_dir", default="", type=str, help="path to store results as wav file")
parser.add_argument("--device_ids", nargs='+', type=int, default=0, help='list of gpu ids')
parser.add_argument("--extract_instrumental", action='store_true', help="invert vocals to get instrumental if provided")
parser.add_argument("--disable_detailed_pbar", action='store_true', help="disable detailed progress bar")
parser.add_argument("--force_cpu", action = 'store_true', help = "Force the use of CPU even if CUDA is available")
if args is None:
args = parser.parse_args()
else:
args = parser.parse_args(args)
device = "cpu"
if args.force_cpu:
device = "cpu"
elif torch.cuda.is_available():
print('CUDA is available, use --force_cpu to disable it.')
device = "cuda"
device = f'cuda:{args.device_ids}' if type(args.device_ids) == int else f'cuda:{args.device_ids[0]}'
elif torch.backends.mps.is_available():
device = "mps"
print("Using device: ", device)
model_load_start_time = time.time()
torch.backends.cudnn.benchmark = True
model, config = get_model_from_config(args.model_type, args.config_path)
if args.start_check_point != '':
print('Start from checkpoint: {}'.format(args.start_check_point))
if args.model_type == 'htdemucs':
state_dict = torch.load(args.start_check_point, map_location = device, weights_only=False)
# Fix for htdemucs pretrained models
if 'state' in state_dict:
state_dict = state_dict['state']
else:
state_dict = torch.load(args.start_check_point, map_location = device, weights_only=True)
model.load_state_dict(state_dict)
print("Instruments: {}".format(config.training.instruments))
# in case multiple CUDA GPUs are used and --device_ids arg is passed
if type(args.device_ids) != int:
model = nn.DataParallel(model, device_ids = args.device_ids)
model = model.to(device)
print("Model load time: {:.2f} sec".format(time.time() - model_load_start_time))
run_folder(model, args, config, device, verbose=True)
if __name__ == "__main__":
proc_folder(None)
|