Spaces:
Sleeping
Sleeping
File size: 5,551 Bytes
56b2bcf |
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 |
import gc
import requests
import subprocess
import logging
import sys
from bs4 import BeautifulSoup
import torch, pdb, os, warnings, librosa
import soundfile as sf
from tqdm import tqdm
import numpy as np
import torch
now_dir = os.getcwd()
sys.path.append(now_dir)
import mdx
branch = "https://github.com/NaJeongMo/Colab-for-MDX_B"
model_params = "https://raw.githubusercontent.com/TRvlvr/application_data/main/mdx_model_data/model_data.json"
_Models = "https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/"
# _models = "https://pastebin.com/raw/jBzYB8vz"
_models = "https://raw.githubusercontent.com/TRvlvr/application_data/main/filelists/download_checks.json"
#stem_naming = "https://pastebin.com/raw/mpH4hRcF"
file_folder = "Colab-for-MDX_B"
model_ids = requests.get(_models).json()
model_ids = model_ids["mdx_download_list"].values()
#print(model_ids)
model_params = requests.get(model_params).json()
#stem_naming = requests.get(stem_naming).json()
stem_naming = {
"Vocals": "Instrumental",
"Other": "Instruments",
"Instrumental": "Vocals",
"Drums": "Drumless",
"Bass": "Bassless"
}
os.makedirs("tmp_models", exist_ok=True)
warnings.filterwarnings("ignore")
cpu = torch.device("cpu")
if torch.cuda.is_available():
device = torch.device("cuda:0")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
def get_model_list():
return model_ids
def id_to_ptm(mkey):
if mkey in model_ids:
mpath = f"{now_dir}/tmp_models/{mkey}"
if not os.path.exists(f'{now_dir}/tmp_models/{mkey}'):
print('Downloading model...',end=' ')
subprocess.run(
["wget", _Models+mkey, "-O", mpath]
)
print(f'saved to {mpath}')
# get_ipython().system(f'gdown {model_id} -O /content/tmp_models/{mkey}')
return mpath
else:
return mpath
else:
mpath = f'models/{mkey}'
return mpath
def prepare_mdx(onnx,custom_param=False, dim_f=None, dim_t=None, n_fft=None, stem_name=None, compensation=None):
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
if custom_param:
assert not (dim_f is None or dim_t is None or n_fft is None or compensation is None), 'Custom parameter selected, but incomplete parameters are provided.'
mdx_model = mdx.MDX_Model(
device,
dim_f = dim_f,
dim_t = dim_t,
n_fft = n_fft,
stem_name=stem_name,
compensation=compensation
)
else:
model_hash = mdx.MDX.get_hash(onnx)
if model_hash in model_params:
mp = model_params.get(model_hash)
mdx_model = mdx.MDX_Model(
device,
dim_f = mp["mdx_dim_f_set"],
dim_t = 2**mp["mdx_dim_t_set"],
n_fft = mp["mdx_n_fft_scale_set"],
stem_name=mp["primary_stem"],
compensation=compensation if not custom_param and compensation is not None else mp["compensate"]
)
return mdx_model
def run_mdx(onnx, mdx_model,filename, output_format='wav',diff=False,suffix=None,diff_suffix=None, denoise=False, m_threads=2):
mdx_sess = mdx.MDX(onnx,mdx_model)
print(f"Processing: {filename}")
if filename.lower().endswith('.wav'):
wave, sr = librosa.load(filename, mono=False, sr=44100)
else:
temp_wav = 'temp_audio.wav'
subprocess.run(['ffmpeg', '-i', filename, '-ar', '44100', '-ac', '2', temp_wav]) # Convert to WAV format
wave, sr = librosa.load(temp_wav, mono=False, sr=44100)
os.remove(temp_wav)
#wave, sr = librosa.load(filename,mono=False, sr=44100)
# normalizing input wave gives better output
peak = max(np.max(wave), abs(np.min(wave)))
wave /= peak
if denoise:
wave_processed = -(mdx_sess.process_wave(-wave, m_threads)) + (mdx_sess.process_wave(wave, m_threads))
wave_processed *= 0.5
else:
wave_processed = mdx_sess.process_wave(wave, m_threads)
# return to previous peak
wave_processed *= peak
stem_name = mdx_model.stem_name if suffix is None else suffix # use suffix if provided
save_path = os.path.basename(os.path.splitext(filename)[0])
#vocals_save_path = os.path.join(vocals_folder, f"{save_path}_{stem_name}.{output_format}")
#instrumental_save_path = os.path.join(instrumental_folder, f"{save_path}_{stem_name}.{output_format}")
save_path = f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.{output_format}"
save_path = os.path.join(
'audios',
save_path
)
sf.write(
save_path,
wave_processed.T,
sr
)
print(f'done, saved to: {save_path}')
if diff:
diff_stem_name = stem_naming.get(stem_name) if diff_suffix is None else diff_suffix # use suffix if provided
stem_name = f"{stem_name}_diff" if diff_stem_name is None else diff_stem_name
save_path = f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.{output_format}"
save_path = os.path.join(
'audio-others',
save_path
)
sf.write(
save_path,
(-wave_processed.T*mdx_model.compensation)+wave.T,
sr
)
print(f'invert done, saved to: {save_path}')
del mdx_sess, wave_processed, wave
gc.collect()
if __name__ == "__main__":
print() |