Florii/Ultimate-Vocal-Remover
0
1import audioread2import librosa3 4import os5import sys6import json7import time8from tqdm import tqdm9import pickle10import hashlib11import logging12import traceback13import shutil14import soundfile as sf15 16import torch17 18from gui_data.constants import *19from gui_data.old_data_check import file_check, remove_unneeded_yamls, remove_temps20from lib_v5.vr_network.model_param_init import ModelParameters21from lib_v5 import spec_utils22from pathlib import Path23from separate import SeperateAttributes, SeperateDemucs, SeperateMDX, SeperateVR, save_format24from typing import List25 26 27logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.INFO)28logging.info('UVR BEGIN')29 30PREVIOUS_PATCH_WIN = 'UVR_Patch_1_12_23_14_54'31 32is_dnd_compatible = True33banner_placement = -234 35def save_data(data):36 """37 Saves given data as a .pkl (pickle) file38 39 Paramters:40 data(dict):41 Dictionary containing all the necessary data to save42 """43 # Open data file, create it if it does not exist44 with open('data.pkl', 'wb') as data_file:45 pickle.dump(data, data_file)46 47def load_data() -> dict:48 """49 Loads saved pkl file and returns the stored data50 51 Returns(dict):52 Dictionary containing all the saved data53 """54 try:55 with open('data.pkl', 'rb') as data_file: # Open data file56 data = pickle.load(data_file)57 58 return data59 except (ValueError, FileNotFoundError):60 # Data File is corrupted or not found so recreate it61 62 save_data(data=DEFAULT_DATA)63 64 return load_data()65 66def load_model_hash_data(dictionary):67 '''Get the model hash dictionary'''68 69 with open(dictionary) as d:70 data = d.read()71 72 return json.loads(data)73 74# Change the current working directory to the directory75# this file sits in76if getattr(sys, 'frozen', False):77 # If the application is run as a bundle, the PyInstaller bootloader78 # extends the sys module by a flag frozen=True and sets the app79 # path into variable _MEIPASS'.80 BASE_PATH = sys._MEIPASS81else:82 BASE_PATH = os.path.dirname(os.path.abspath(__file__))83 84os.chdir(BASE_PATH) # Change the current working directory to the base path85 86debugger = []87 88#--Constants--89#Models90MODELS_DIR = os.path.join(BASE_PATH, 'models')91VR_MODELS_DIR = os.path.join(MODELS_DIR, 'VR_Models')92MDX_MODELS_DIR = os.path.join(MODELS_DIR, 'MDX_Net_Models')93DEMUCS_MODELS_DIR = os.path.join(MODELS_DIR, 'Demucs_Models')94DEMUCS_NEWER_REPO_DIR = os.path.join(DEMUCS_MODELS_DIR, 'v3_v4_repo')95MDX_MIXER_PATH = os.path.join(BASE_PATH, 'lib_v5', 'mixer.ckpt')96 97#Cache & Parameters98VR_HASH_DIR = os.path.join(VR_MODELS_DIR, 'model_data')99VR_HASH_JSON = os.path.join(VR_MODELS_DIR, 'model_data', 'model_data.json')100MDX_HASH_DIR = os.path.join(MDX_MODELS_DIR, 'model_data')101MDX_HASH_JSON = os.path.join(MDX_MODELS_DIR, 'model_data', 'model_data.json')102DEMUCS_MODEL_NAME_SELECT = os.path.join(DEMUCS_MODELS_DIR, 'model_data', 'model_name_mapper.json')103MDX_MODEL_NAME_SELECT = os.path.join(MDX_MODELS_DIR, 'model_data', 'model_name_mapper.json')104ENSEMBLE_CACHE_DIR = os.path.join(BASE_PATH, 'gui_data', 'saved_ensembles')105SETTINGS_CACHE_DIR = os.path.join(BASE_PATH, 'gui_data', 'saved_settings')106VR_PARAM_DIR = os.path.join(BASE_PATH, 'lib_v5', 'vr_network', 'modelparams')107SAMPLE_CLIP_PATH = os.path.join(BASE_PATH, 'temp_sample_clips')108ENSEMBLE_TEMP_PATH = os.path.join(BASE_PATH, 'ensemble_temps')109 110#Style111ICON_IMG_PATH = os.path.join(BASE_PATH, 'gui_data', 'img', 'GUI-Icon.ico')112FONT_PATH = os.path.join(BASE_PATH, 'gui_data', 'fonts', 'centurygothic', 'GOTHIC.TTF')#ensemble_temps113 114#Other115COMPLETE_CHIME = os.path.join(BASE_PATH, 'gui_data', 'complete_chime.wav')116FAIL_CHIME = os.path.join(BASE_PATH, 'gui_data', 'fail_chime.wav')117CHANGE_LOG = os.path.join(BASE_PATH, 'gui_data', 'change_log.txt')118SPLASH_DOC = os.path.join(BASE_PATH, 'tmp', 'splash.txt')119 120file_check(os.path.join(MODELS_DIR, 'Main_Models'), VR_MODELS_DIR)121file_check(os.path.join(DEMUCS_MODELS_DIR, 'v3_repo'), DEMUCS_NEWER_REPO_DIR)122remove_unneeded_yamls(DEMUCS_MODELS_DIR)123 124remove_temps(ENSEMBLE_TEMP_PATH)125remove_temps(SAMPLE_CLIP_PATH)126remove_temps(os.path.join(BASE_PATH, 'img'))127 128if not os.path.isdir(ENSEMBLE_TEMP_PATH):129 os.mkdir(ENSEMBLE_TEMP_PATH)130 131if not os.path.isdir(SAMPLE_CLIP_PATH):132 os.mkdir(SAMPLE_CLIP_PATH)133 134model_hash_table = {}135data = load_data()136 137class ModelData():138 def __init__(self, model_name: str, 139 selected_process_method=ENSEMBLE_MODE, 140 is_secondary_model=False, 141 primary_model_primary_stem=None, 142 is_primary_model_primary_stem_only=False, 143 is_primary_model_secondary_stem_only=False, 144 is_pre_proc_model=False,145 is_dry_check=False):146 147 self.is_gpu_conversion = 0 if root.is_gpu_conversion_var.get() else -1148 self.is_normalization = root.is_normalization_var.get()149 self.is_primary_stem_only = root.is_primary_stem_only_var.get()150 self.is_secondary_stem_only = root.is_secondary_stem_only_var.get()151 self.is_denoise = root.is_denoise_var.get()152 self.mdx_batch_size = 1 if root.mdx_batch_size_var.get() == DEF_OPT else int(root.mdx_batch_size_var.get())153 self.is_mdx_ckpt = False154 self.wav_type_set = root.wav_type_set155 self.mp3_bit_set = root.mp3_bit_set_var.get()156 self.save_format = root.save_format_var.get()157 self.is_invert_spec = root.is_invert_spec_var.get()158 self.is_mixer_mode = root.is_mixer_mode_var.get()159 self.demucs_stems = root.demucs_stems_var.get()160 self.demucs_source_list = []161 self.demucs_stem_count = 0162 self.mixer_path = MDX_MIXER_PATH163 self.model_name = model_name164 self.process_method = selected_process_method165 self.model_status = False if self.model_name == CHOOSE_MODEL or self.model_name == NO_MODEL else True166 self.primary_stem = None167 self.secondary_stem = None168 self.is_ensemble_mode = False169 self.ensemble_primary_stem = None170 self.ensemble_secondary_stem = None171 self.primary_model_primary_stem = primary_model_primary_stem172 self.is_secondary_model = is_secondary_model173 self.secondary_model = None174 self.secondary_model_scale = None175 self.demucs_4_stem_added_count = 0176 self.is_demucs_4_stem_secondaries = False177 self.is_4_stem_ensemble = False178 self.pre_proc_model = None179 self.pre_proc_model_activated = False180 self.is_pre_proc_model = is_pre_proc_model181 self.is_dry_check = is_dry_check182 self.model_samplerate = 44100183 self.model_capacity = 32, 128184 self.is_vr_51_model = False185 self.is_demucs_pre_proc_model_inst_mix = False186 self.manual_download_Button = None187 self.secondary_model_4_stem = []188 self.secondary_model_4_stem_scale = []189 self.secondary_model_4_stem_names = []190 self.secondary_model_4_stem_model_names_list = []191 self.all_models = []192 self.secondary_model_other = None193 self.secondary_model_scale_other = None194 self.secondary_model_bass = None195 self.secondary_model_scale_bass = None196 self.secondary_model_drums = None197 self.secondary_model_scale_drums = None198 199 if selected_process_method == ENSEMBLE_MODE:200 partitioned_name = model_name.partition(ENSEMBLE_PARTITION)201 self.process_method = partitioned_name[0]202 self.model_name = partitioned_name[2]203 self.model_and_process_tag = model_name204 self.ensemble_primary_stem, self.ensemble_secondary_stem = root.return_ensemble_stems()205 self.is_ensemble_mode = True if not is_secondary_model and not is_pre_proc_model else False206 self.is_4_stem_ensemble = True if root.ensemble_main_stem_var.get() == FOUR_STEM_ENSEMBLE and self.is_ensemble_mode else False207 self.pre_proc_model_activated = root.is_demucs_pre_proc_model_activate_var.get() if not self.ensemble_primary_stem == VOCAL_STEM else False208 209 if self.process_method == VR_ARCH_TYPE:210 self.is_secondary_model_activated = root.vr_is_secondary_model_activate_var.get() if not self.is_secondary_model else False211 self.aggression_setting = float(int(root.aggression_setting_var.get())/100)212 self.is_tta = root.is_tta_var.get()213 self.is_post_process = root.is_post_process_var.get()214 self.window_size = int(root.window_size_var.get())215 self.batch_size = 1 if root.batch_size_var.get() == DEF_OPT else int(root.batch_size_var.get())216 self.crop_size = int(root.crop_size_var.get())217 self.is_high_end_process = 'mirroring' if root.is_high_end_process_var.get() else 'None'218 self.post_process_threshold = float(root.post_process_threshold_var.get())219 self.model_capacity = 32, 128220 self.model_path = os.path.join(VR_MODELS_DIR, f"{self.model_name}.pth")221 self.get_model_hash()222 if self.model_hash:223 self.model_data = self.get_model_data(VR_HASH_DIR, root.vr_hash_MAPPER) if not self.model_hash == WOOD_INST_MODEL_HASH else WOOD_INST_PARAMS224 if self.model_data:225 vr_model_param = os.path.join(VR_PARAM_DIR, "{}.json".format(self.model_data["vr_model_param"]))226 self.primary_stem = self.model_data["primary_stem"]227 self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]228 self.vr_model_param = ModelParameters(vr_model_param)229 self.model_samplerate = self.vr_model_param.param['sr']230 if "nout" in self.model_data.keys() and "nout_lstm" in self.model_data.keys():231 self.model_capacity = self.model_data["nout"], self.model_data["nout_lstm"]232 self.is_vr_51_model = True233 else:234 self.model_status = False235 236 if self.process_method == MDX_ARCH_TYPE:237 self.is_secondary_model_activated = root.mdx_is_secondary_model_activate_var.get() if not is_secondary_model else False238 self.margin = int(root.margin_var.get())239 self.chunks = root.determine_auto_chunks(root.chunks_var.get(), self.is_gpu_conversion) if root.is_chunk_mdxnet_var.get() else 0240 self.get_mdx_model_path()241 self.get_model_hash()242 if self.model_hash:243 self.model_data = self.get_model_data(MDX_HASH_DIR, root.mdx_hash_MAPPER)244 if self.model_data:245 self.compensate = self.model_data["compensate"] if root.compensate_var.get() == AUTO_SELECT else float(root.compensate_var.get())246 self.mdx_dim_f_set = self.model_data["mdx_dim_f_set"]247 self.mdx_dim_t_set = self.model_data["mdx_dim_t_set"]248 self.mdx_n_fft_scale_set = self.model_data["mdx_n_fft_scale_set"]249 self.primary_stem = self.model_data["primary_stem"]250 self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]251 else:252 self.model_status = False253 254 if self.process_method == DEMUCS_ARCH_TYPE:255 self.is_secondary_model_activated = root.demucs_is_secondary_model_activate_var.get() if not is_secondary_model else False256 if not self.is_ensemble_mode:257 self.pre_proc_model_activated = root.is_demucs_pre_proc_model_activate_var.get() if not root.demucs_stems_var.get() in [VOCAL_STEM, INST_STEM] else False258 self.overlap = float(root.overlap_var.get())259 self.margin_demucs = int(root.margin_demucs_var.get())260 self.chunks_demucs = root.determine_auto_chunks(root.chunks_demucs_var.get(), self.is_gpu_conversion)261 self.shifts = int(root.shifts_var.get())262 self.is_split_mode = root.is_split_mode_var.get()263 self.segment = root.segment_var.get()264 self.is_chunk_demucs = root.is_chunk_demucs_var.get()265 self.is_demucs_combine_stems = root.is_demucs_combine_stems_var.get()266 self.is_primary_stem_only = root.is_primary_stem_only_var.get() if self.is_ensemble_mode else root.is_primary_stem_only_Demucs_var.get() 267 self.is_secondary_stem_only = root.is_secondary_stem_only_var.get() if self.is_ensemble_mode else root.is_secondary_stem_only_Demucs_var.get()268 self.get_demucs_model_path()269 self.get_demucs_model_data()270 271 self.model_basename = os.path.splitext(os.path.basename(self.model_path))[0] if self.model_status else None272 self.pre_proc_model_activated = self.pre_proc_model_activated if not self.is_secondary_model else False273 274 self.is_primary_model_primary_stem_only = is_primary_model_primary_stem_only275 self.is_primary_model_secondary_stem_only = is_primary_model_secondary_stem_only276 277 if self.is_secondary_model_activated and self.model_status:278 if (not self.is_ensemble_mode and root.demucs_stems_var.get() == ALL_STEMS and self.process_method == DEMUCS_ARCH_TYPE) or self.is_4_stem_ensemble:279 for key in DEMUCS_4_SOURCE_LIST:280 self.secondary_model_data(key)281 self.secondary_model_4_stem.append(self.secondary_model)282 self.secondary_model_4_stem_scale.append(self.secondary_model_scale)283 self.secondary_model_4_stem_names.append(key)284 self.demucs_4_stem_added_count = sum(i is not None for i in self.secondary_model_4_stem)285 self.is_secondary_model_activated = False if all(i is None for i in self.secondary_model_4_stem) else True286 self.demucs_4_stem_added_count = self.demucs_4_stem_added_count - 1 if self.is_secondary_model_activated else self.demucs_4_stem_added_count287 if self.is_secondary_model_activated:288 self.secondary_model_4_stem_model_names_list = [None if i is None else i.model_basename for i in self.secondary_model_4_stem]289 self.is_demucs_4_stem_secondaries = True 290 else:291 primary_stem = self.ensemble_primary_stem if self.is_ensemble_mode and self.process_method == DEMUCS_ARCH_TYPE else self.primary_stem292 self.secondary_model_data(primary_stem)293 294 if self.process_method == DEMUCS_ARCH_TYPE and not is_secondary_model:295 if self.demucs_stem_count >= 3 and self.pre_proc_model_activated:296 self.pre_proc_model_activated = True297 self.pre_proc_model = root.process_determine_demucs_pre_proc_model(self.primary_stem)298 self.is_demucs_pre_proc_model_inst_mix = root.is_demucs_pre_proc_model_inst_mix_var.get() if self.pre_proc_model else False299 300 def secondary_model_data(self, primary_stem):301 secondary_model_data = root.process_determine_secondary_model(self.process_method, primary_stem, self.is_primary_stem_only, self.is_secondary_stem_only)302 self.secondary_model = secondary_model_data[0]303 self.secondary_model_scale = secondary_model_data[1]304 self.is_secondary_model_activated = False if not self.secondary_model else True305 if self.secondary_model:306 self.is_secondary_model_activated = False if self.secondary_model.model_basename == self.model_basename else True307 308 def get_mdx_model_path(self):309 310 if self.model_name.endswith(CKPT):311 # self.chunks = 0312 # self.is_mdx_batch_mode = True313 self.is_mdx_ckpt = True314 315 ext = '' if self.is_mdx_ckpt else ONNX316 317 for file_name, chosen_mdx_model in root.mdx_name_select_MAPPER.items():318 if self.model_name in chosen_mdx_model:319 self.model_path = os.path.join(MDX_MODELS_DIR, f"{file_name}{ext}")320 break321 else:322 self.model_path = os.path.join(MDX_MODELS_DIR, f"{self.model_name}{ext}")323 324 self.mixer_path = os.path.join(MDX_MODELS_DIR, f"mixer_val.ckpt")325 326 def get_demucs_model_path(self):327 328 demucs_newer = [True for x in DEMUCS_NEWER_TAGS if x in self.model_name]329 demucs_model_dir = DEMUCS_NEWER_REPO_DIR if demucs_newer else DEMUCS_MODELS_DIR330 331 for file_name, chosen_model in root.demucs_name_select_MAPPER.items():332 if self.model_name in chosen_model:333 self.model_path = os.path.join(demucs_model_dir, file_name)334 break335 else:336 self.model_path = os.path.join(DEMUCS_NEWER_REPO_DIR, f'{self.model_name}.yaml')337 338 def get_demucs_model_data(self):339 340 self.demucs_version = DEMUCS_V4341 342 for key, value in DEMUCS_VERSION_MAPPER.items():343 if value in self.model_name:344 self.demucs_version = key345 346 self.demucs_source_list = DEMUCS_2_SOURCE if DEMUCS_UVR_MODEL in self.model_name else DEMUCS_4_SOURCE347 self.demucs_source_map = DEMUCS_2_SOURCE_MAPPER if DEMUCS_UVR_MODEL in self.model_name else DEMUCS_4_SOURCE_MAPPER348 self.demucs_stem_count = 2 if DEMUCS_UVR_MODEL in self.model_name else 4349 350 if not self.is_ensemble_mode:351 self.primary_stem = PRIMARY_STEM if self.demucs_stems == ALL_STEMS else self.demucs_stems352 self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]353 354 def get_model_data(self, model_hash_dir, hash_mapper):355 model_settings_json = os.path.join(model_hash_dir, "{}.json".format(self.model_hash))356 357 if os.path.isfile(model_settings_json):358 return json.load(open(model_settings_json))359 else:360 for hash, settings in hash_mapper.items():361 if self.model_hash in hash:362 return settings363 else:364 return self.get_model_data_from_popup()365 366 def get_model_data_from_popup(self):367 return None368 369 def get_model_hash(self):370 self.model_hash = None371 372 if not os.path.isfile(self.model_path):373 self.model_status = False374 self.model_hash is None375 else:376 if model_hash_table:377 for (key, value) in model_hash_table.items():378 if self.model_path == key:379 self.model_hash = value380 break381 382 if not self.model_hash:383 try:384 with open(self.model_path, 'rb') as f:385 f.seek(- 10000 * 1024, 2)386 self.model_hash = hashlib.md5(f.read()).hexdigest()387 except:388 self.model_hash = hashlib.md5(open(self.model_path,'rb').read()).hexdigest()389 390 table_entry = {self.model_path: self.model_hash}391 model_hash_table.update(table_entry)392 393 394class Ensembler():395 def __init__(self, is_manual_ensemble=False):396 self.is_save_all_outputs_ensemble = root.is_save_all_outputs_ensemble_var.get()397 chosen_ensemble_name = '{}'.format(root.chosen_ensemble_var.get().replace(" ", "_")) if not root.chosen_ensemble_var.get() == CHOOSE_ENSEMBLE_OPTION else 'Ensembled'398 ensemble_algorithm = root.ensemble_type_var.get().partition("/")399 ensemble_main_stem_pair = root.ensemble_main_stem_var.get().partition("/")400 time_stamp = round(time.time())401 self.audio_tool = MANUAL_ENSEMBLE402 self.main_export_path = Path(root.export_path_var.get())403 self.chosen_ensemble = f"_{chosen_ensemble_name}" if root.is_append_ensemble_name_var.get() else ''404 ensemble_folder_name = self.main_export_path if self.is_save_all_outputs_ensemble else ENSEMBLE_TEMP_PATH405 self.ensemble_folder_name = os.path.join(ensemble_folder_name, '{}_Outputs_{}'.format(chosen_ensemble_name, time_stamp))406 self.is_testing_audio = f"{time_stamp}_" if root.is_testing_audio_var.get() else ''407 self.primary_algorithm = ensemble_algorithm[0]408 self.secondary_algorithm = ensemble_algorithm[2]409 self.ensemble_primary_stem = ensemble_main_stem_pair[0]410 self.ensemble_secondary_stem = ensemble_main_stem_pair[2]411 self.is_normalization = root.is_normalization_var.get()412 self.wav_type_set = root.wav_type_set413 self.mp3_bit_set = root.mp3_bit_set_var.get()414 self.save_format = root.save_format_var.get()415 if not is_manual_ensemble:416 os.mkdir(self.ensemble_folder_name)417 418 def ensemble_outputs(self, audio_file_base, export_path, stem, is_4_stem=False, is_inst_mix=False):419 """Processes the given outputs and ensembles them with the chosen algorithm"""420 421 if is_4_stem:422 algorithm = root.ensemble_type_var.get()423 stem_tag = stem424 else:425 if is_inst_mix:426 algorithm = self.secondary_algorithm427 stem_tag = f"{self.ensemble_secondary_stem} {INST_STEM}"428 else:429 algorithm = self.primary_algorithm if stem == PRIMARY_STEM else self.secondary_algorithm430 stem_tag = self.ensemble_primary_stem if stem == PRIMARY_STEM else self.ensemble_secondary_stem431 432 stem_outputs = self.get_files_to_ensemble(folder=export_path, prefix=audio_file_base, suffix=f"_({stem_tag}).wav")433 audio_file_output = f"{self.is_testing_audio}{audio_file_base}{self.chosen_ensemble}_({stem_tag})"434 stem_save_path = os.path.join('{}'.format(self.main_export_path),'{}.wav'.format(audio_file_output))435 436 if stem_outputs:437 spec_utils.ensemble_inputs(stem_outputs, algorithm, self.is_normalization, self.wav_type_set, stem_save_path)438 save_format(stem_save_path, self.save_format, self.mp3_bit_set)439 440 if self.is_save_all_outputs_ensemble:441 for i in stem_outputs:442 save_format(i, self.save_format, self.mp3_bit_set)443 else:444 for i in stem_outputs:445 try:446 os.remove(i)447 except Exception as e:448 print(e)449 450 def ensemble_manual(self, audio_inputs, audio_file_base, is_bulk=False):451 """Processes the given outputs and ensembles them with the chosen algorithm"""452 453 is_mv_sep = True454 455 if is_bulk:456 number_list = list(set([os.path.basename(i).split("_")[0] for i in audio_inputs]))457 for n in number_list:458 current_list = [i for i in audio_inputs if os.path.basename(i).startswith(n)]459 audio_file_base = os.path.basename(current_list[0]).split('.wav')[0]460 stem_testing = "instrum" if "Instrumental" in audio_file_base else "vocals"461 if is_mv_sep:462 audio_file_base = audio_file_base.split("_")463 audio_file_base = f"{audio_file_base[1]}_{audio_file_base[2]}_{stem_testing}"464 self.ensemble_manual_process(current_list, audio_file_base, is_bulk)465 else:466 self.ensemble_manual_process(audio_inputs, audio_file_base, is_bulk)467 468 def ensemble_manual_process(self, audio_inputs, audio_file_base, is_bulk):469 470 algorithm = root.choose_algorithm_var.get()471 algorithm_text = "" if is_bulk else f"_({root.choose_algorithm_var.get()})"472 stem_save_path = os.path.join('{}'.format(self.main_export_path),'{}{}{}.wav'.format(self.is_testing_audio, audio_file_base, algorithm_text))473 spec_utils.ensemble_inputs(audio_inputs, algorithm, self.is_normalization, self.wav_type_set, stem_save_path)474 save_format(stem_save_path, self.save_format, self.mp3_bit_set)475 476 def get_files_to_ensemble(self, folder="", prefix="", suffix=""):477 """Grab all the files to be ensembled"""478 479 return [os.path.join(folder, i) for i in os.listdir(folder) if i.startswith(prefix) and i.endswith(suffix)]480 481 482def secondary_stem(stem):483 """Determines secondary stem"""484 485 for key, value in STEM_PAIR_MAPPER.items():486 if stem in key:487 secondary_stem = value488 489 return secondary_stem490 491 492class UVRInterface:493 def __init__(self) -> None:494 pass495 496 def assemble_model_data(self, model=None, arch_type=ENSEMBLE_MODE, is_dry_check=False) -> List[ModelData]:497 if arch_type == ENSEMBLE_STEM_CHECK:498 model_data = self.model_data_table499 missing_models = [model.model_status for model in model_data if not model.model_status]500 501 if missing_models or not model_data:502 model_data: List[ModelData] = [ModelData(model_name, is_dry_check=is_dry_check) for model_name in self.ensemble_model_list]503 self.model_data_table = model_data504 505 if arch_type == ENSEMBLE_MODE:506 model_data: List[ModelData] = [ModelData(model_name) for model_name in self.ensemble_listbox_get_all_selected_models()]507 if arch_type == ENSEMBLE_CHECK:508 model_data: List[ModelData] = [ModelData(model)]509 if arch_type == VR_ARCH_TYPE or arch_type == VR_ARCH_PM:510 model_data: List[ModelData] = [ModelData(model, VR_ARCH_TYPE)]511 if arch_type == MDX_ARCH_TYPE:512 model_data: List[ModelData] = [ModelData(model, MDX_ARCH_TYPE)]513 if arch_type == DEMUCS_ARCH_TYPE:514 model_data: List[ModelData] = [ModelData(model, DEMUCS_ARCH_TYPE)]#515 516 return model_data517 518 def create_sample(self, audio_file, sample_path=SAMPLE_CLIP_PATH):519 try:520 with audioread.audio_open(audio_file) as f:521 track_length = int(f.duration)522 except Exception as e:523 print('Audioread failed to get duration. Trying Librosa...')524 y, sr = librosa.load(audio_file, mono=False, sr=44100)525 track_length = int(librosa.get_duration(y=y, sr=sr))526 527 clip_duration = int(root.model_sample_mode_duration_var.get())528 529 if track_length >= clip_duration:530 offset_cut = track_length//3531 off_cut = offset_cut + track_length532 if not off_cut >= clip_duration:533 offset_cut = 0534 name_apped = f'{clip_duration}_second_'535 else:536 offset_cut, clip_duration = 0, track_length537 name_apped = ''538 539 sample = librosa.load(audio_file, offset=offset_cut, duration=clip_duration, mono=False, sr=44100)[0].T540 audio_sample = os.path.join(sample_path, f'{os.path.splitext(os.path.basename(audio_file))[0]}_{name_apped}sample.wav')541 sf.write(audio_sample, sample, 44100)542 543 return audio_sample544 545 def verify_audio(self, audio_file, is_process=True, sample_path=None):546 is_good = False547 error_data = ''548 549 if os.path.isfile(audio_file):550 try:551 librosa.load(audio_file, duration=3, mono=False, sr=44100) if not type(sample_path) is str else self.create_sample(audio_file, sample_path)552 is_good = True553 except Exception as e:554 error_name = f'{type(e).__name__}'555 traceback_text = ''.join(traceback.format_tb(e.__traceback__))556 message = f'{error_name}: "{e}"\n{traceback_text}"'557 if is_process:558 audio_base_name = os.path.basename(audio_file)559 self.error_log_var.set(f'Error Loading the Following File:\n\n\"{audio_base_name}\"\n\nRaw Error Details:\n\n{message}')560 else:561 error_data = AUDIO_VERIFICATION_CHECK(audio_file, message)562 563 if is_process:564 return is_good565 else:566 return is_good, error_data567 568 def cached_sources_clear(self):569 self.vr_cache_source_mapper = {}570 self.mdx_cache_source_mapper = {}571 self.demucs_cache_source_mapper = {}572 573 def cached_model_source_holder(self, process_method, sources, model_name=None):574 if process_method == VR_ARCH_TYPE:575 self.vr_cache_source_mapper = {**self.vr_cache_source_mapper, **{model_name: sources}}576 if process_method == MDX_ARCH_TYPE:577 self.mdx_cache_source_mapper = {**self.mdx_cache_source_mapper, **{model_name: sources}}578 if process_method == DEMUCS_ARCH_TYPE:579 self.demucs_cache_source_mapper = {**self.demucs_cache_source_mapper, **{model_name: sources}}580 581 def cached_source_callback(self, process_method, model_name=None):582 model, sources = None, None583 584 if process_method == VR_ARCH_TYPE:585 mapper = self.vr_cache_source_mapper586 if process_method == MDX_ARCH_TYPE:587 mapper = self.mdx_cache_source_mapper588 if process_method == DEMUCS_ARCH_TYPE:589 mapper = self.demucs_cache_source_mapper590 591 for key, value in mapper.items():592 if model_name in key:593 model = key594 sources = value595 596 return model, sources597 598 def cached_source_model_list_check(self, model_list: List[ModelData]):599 model: ModelData600 primary_model_names = lambda process_method:[model.model_basename if model.process_method == process_method else None for model in model_list]601 secondary_model_names = lambda process_method:[model.secondary_model.model_basename if model.is_secondary_model_activated and model.process_method == process_method else None for model in model_list]602 603 self.vr_primary_model_names = primary_model_names(VR_ARCH_TYPE)604 self.mdx_primary_model_names = primary_model_names(MDX_ARCH_TYPE)605 self.demucs_primary_model_names = primary_model_names(DEMUCS_ARCH_TYPE)606 self.vr_secondary_model_names = secondary_model_names(VR_ARCH_TYPE)607 self.mdx_secondary_model_names = secondary_model_names(MDX_ARCH_TYPE)608 self.demucs_secondary_model_names = [model.secondary_model.model_basename if model.is_secondary_model_activated and model.process_method == DEMUCS_ARCH_TYPE and not model.secondary_model is None else None for model in model_list]609 self.demucs_pre_proc_model_name = [model.pre_proc_model.model_basename if model.pre_proc_model else None for model in model_list]#list(dict.fromkeys())610 611 for model in model_list:612 if model.process_method == DEMUCS_ARCH_TYPE and model.is_demucs_4_stem_secondaries:613 if not model.is_4_stem_ensemble:614 self.demucs_secondary_model_names = model.secondary_model_4_stem_model_names_list615 break616 else:617 for i in model.secondary_model_4_stem_model_names_list:618 self.demucs_secondary_model_names.append(i)619 620 self.all_models = self.vr_primary_model_names + self.mdx_primary_model_names + self.demucs_primary_model_names + self.vr_secondary_model_names + self.mdx_secondary_model_names + self.demucs_secondary_model_names + self.demucs_pre_proc_model_name621 622 def process(self, model_name, arch_type, audio_file, export_path, is_model_sample_mode=False, is_4_stem_ensemble=False, set_progress_func=None, console_write=print) -> SeperateAttributes:623 stime = time.perf_counter()624 time_elapsed = lambda:f'Time Elapsed: {time.strftime("%H:%M:%S", time.gmtime(int(time.perf_counter() - stime)))}'625 626 if arch_type==ENSEMBLE_MODE:627 model_list, ensemble = self.assemble_model_data(), Ensembler()628 export_path = ensemble.ensemble_folder_name629 is_ensemble = True630 else:631 model_list = self.assemble_model_data(model_name, arch_type)632 is_ensemble = False633 self.cached_source_model_list_check(model_list)634 model = model_list[0]635 636 if self.verify_audio(audio_file):637 audio_file = self.create_sample(audio_file) if is_model_sample_mode else audio_file638 else:639 print(f'"{os.path.basename(audio_file)}\" is missing or currupted.\n')640 exit()641 642 audio_file_base = f"{os.path.splitext(os.path.basename(audio_file))[0]}"643 audio_file_base = audio_file_base if is_ensemble else f"{round(time.time())}_{audio_file_base}"644 audio_file_base = audio_file_base if not is_ensemble else f"{audio_file_base}_{model.model_basename}"645 if not is_ensemble:646 audio_file_base = f"{audio_file_base}_{model.model_basename}"647 648 if not is_ensemble:649 export_path = os.path.join(Path(export_path), model.model_basename, os.path.splitext(os.path.basename(audio_file))[0])650 if not os.path.isdir(export_path):651 os.makedirs(export_path) 652 653 if set_progress_func is None:654 pbar = tqdm(total=1)655 self._progress = 0656 def set_progress_func(step, inference_iterations=0):657 progress_curr = step + inference_iterations658 pbar.update(progress_curr-self._progress)659 self._progress = progress_curr660 661 def postprocess():662 pbar.close()663 else:664 def postprocess():665 pass666 667 process_data = {668 'model_data': model,669 'export_path': export_path,670 'audio_file_base': audio_file_base,671 'audio_file': audio_file,672 'set_progress_bar': set_progress_func,673 'write_to_console': lambda progress_text, base_text='': console_write(base_text + progress_text),674 'process_iteration': lambda:None,675 'cached_source_callback': self.cached_source_callback,676 'cached_model_source_holder': self.cached_model_source_holder,677 'list_all_models': self.all_models,678 'is_ensemble_master': is_ensemble,679 'is_4_stem_ensemble': is_ensemble and is_4_stem_ensemble680 }681 if model.process_method == VR_ARCH_TYPE:682 seperator = SeperateVR(model, process_data)683 if model.process_method == MDX_ARCH_TYPE:684 seperator = SeperateMDX(model, process_data)685 if model.process_method == DEMUCS_ARCH_TYPE:686 seperator = SeperateDemucs(model, process_data)687 688 seperator.seperate()689 postprocess()690 691 if is_ensemble:692 audio_file_base = audio_file_base.replace(f"_{model.model_basename}", "")693 console_write(ENSEMBLING_OUTPUTS)694 695 if is_4_stem_ensemble:696 for output_stem in DEMUCS_4_SOURCE_LIST:697 ensemble.ensemble_outputs(audio_file_base, export_path, output_stem, is_4_stem=True)698 else:699 if not root.is_secondary_stem_only_var.get():700 ensemble.ensemble_outputs(audio_file_base, export_path, PRIMARY_STEM)701 if not root.is_primary_stem_only_var.get():702 ensemble.ensemble_outputs(audio_file_base, export_path, SECONDARY_STEM)703 ensemble.ensemble_outputs(audio_file_base, export_path, SECONDARY_STEM, is_inst_mix=True)704 705 console_write(DONE)706 707 if is_model_sample_mode:708 if os.path.isfile(audio_file):709 os.remove(audio_file)710 711 torch.cuda.empty_cache()712 713 if is_ensemble and len(os.listdir(export_path)) == 0:714 shutil.rmtree(export_path)715 console_write(f'Process Complete, using time: {time_elapsed()}\nOutput path: {export_path}')716 self.cached_sources_clear()717 return seperator718 719 720class RootWrapper:721 def __init__(self, var) -> None:722 self.var=var723 724 def set(self, val):725 self.var=val726 727 def get(self):728 return self.var729 730class FakeRoot:731 def __init__(self) -> None:732 self.wav_type_set = 'PCM_16'733 self.vr_hash_MAPPER = load_model_hash_data(VR_HASH_JSON)734 self.mdx_hash_MAPPER = load_model_hash_data(MDX_HASH_JSON)735 self.mdx_name_select_MAPPER = load_model_hash_data(MDX_MODEL_NAME_SELECT)736 self.demucs_name_select_MAPPER = load_model_hash_data(DEMUCS_MODEL_NAME_SELECT)737 738 def __getattribute__(self, __name: str):739 try:740 return super().__getattribute__(__name)741 except AttributeError:742 wrapped=RootWrapper(None)743 super().__setattr__(__name, wrapped)744 return wrapped745 746 def load_saved_settings(self, loaded_setting: dict, process_method=None):747 """Loads user saved application settings or resets to default"""748 749 for key, value in DEFAULT_DATA.items():750 if not key in loaded_setting.keys():751 loaded_setting = {**loaded_setting, **{key:value}}752 loaded_setting['batch_size'] = DEF_OPT753 754 is_ensemble = True if process_method == ENSEMBLE_MODE else False755 756 if not process_method or process_method == VR_ARCH_PM or is_ensemble:757 self.vr_model_var.set(loaded_setting['vr_model'])758 self.aggression_setting_var.set(loaded_setting['aggression_setting'])759 self.window_size_var.set(loaded_setting['window_size'])760 self.batch_size_var.set(loaded_setting['batch_size'])761 self.crop_size_var.set(loaded_setting['crop_size'])762 self.is_tta_var.set(loaded_setting['is_tta'])763 self.is_output_image_var.set(loaded_setting['is_output_image'])764 self.is_post_process_var.set(loaded_setting['is_post_process'])765 self.is_high_end_process_var.set(loaded_setting['is_high_end_process'])766 self.post_process_threshold_var.set(loaded_setting['post_process_threshold'])767 self.vr_voc_inst_secondary_model_var.set(loaded_setting['vr_voc_inst_secondary_model'])768 self.vr_other_secondary_model_var.set(loaded_setting['vr_other_secondary_model'])769 self.vr_bass_secondary_model_var.set(loaded_setting['vr_bass_secondary_model'])770 self.vr_drums_secondary_model_var.set(loaded_setting['vr_drums_secondary_model'])771 self.vr_is_secondary_model_activate_var.set(loaded_setting['vr_is_secondary_model_activate'])772 self.vr_voc_inst_secondary_model_scale_var.set(loaded_setting['vr_voc_inst_secondary_model_scale'])773 self.vr_other_secondary_model_scale_var.set(loaded_setting['vr_other_secondary_model_scale'])774 self.vr_bass_secondary_model_scale_var.set(loaded_setting['vr_bass_secondary_model_scale'])775 self.vr_drums_secondary_model_scale_var.set(loaded_setting['vr_drums_secondary_model_scale'])776 777 if not process_method or process_method == DEMUCS_ARCH_TYPE or is_ensemble:778 self.demucs_model_var.set(loaded_setting['demucs_model'])779 self.segment_var.set(loaded_setting['segment'])780 self.overlap_var.set(loaded_setting['overlap'])781 self.shifts_var.set(loaded_setting['shifts'])782 self.chunks_demucs_var.set(loaded_setting['chunks_demucs'])783 self.margin_demucs_var.set(loaded_setting['margin_demucs'])784 self.is_chunk_demucs_var.set(loaded_setting['is_chunk_demucs'])785 self.is_chunk_mdxnet_var.set(loaded_setting['is_chunk_mdxnet'])786 self.is_primary_stem_only_Demucs_var.set(loaded_setting['is_primary_stem_only_Demucs'])787 self.is_secondary_stem_only_Demucs_var.set(loaded_setting['is_secondary_stem_only_Demucs'])788 self.is_split_mode_var.set(loaded_setting['is_split_mode'])789 self.is_demucs_combine_stems_var.set(loaded_setting['is_demucs_combine_stems'])790 self.demucs_voc_inst_secondary_model_var.set(loaded_setting['demucs_voc_inst_secondary_model'])791 self.demucs_other_secondary_model_var.set(loaded_setting['demucs_other_secondary_model'])792 self.demucs_bass_secondary_model_var.set(loaded_setting['demucs_bass_secondary_model'])793 self.demucs_drums_secondary_model_var.set(loaded_setting['demucs_drums_secondary_model'])794 self.demucs_is_secondary_model_activate_var.set(loaded_setting['demucs_is_secondary_model_activate'])795 self.demucs_voc_inst_secondary_model_scale_var.set(loaded_setting['demucs_voc_inst_secondary_model_scale'])796 self.demucs_other_secondary_model_scale_var.set(loaded_setting['demucs_other_secondary_model_scale'])797 self.demucs_bass_secondary_model_scale_var.set(loaded_setting['demucs_bass_secondary_model_scale'])798 self.demucs_drums_secondary_model_scale_var.set(loaded_setting['demucs_drums_secondary_model_scale'])799 self.demucs_stems_var.set(loaded_setting['demucs_stems'])800 # self.update_stem_checkbox_labels(self.demucs_stems_var.get(), demucs=True)801 self.demucs_pre_proc_model_var.set(data['demucs_pre_proc_model'])802 self.is_demucs_pre_proc_model_activate_var.set(data['is_demucs_pre_proc_model_activate'])803 self.is_demucs_pre_proc_model_inst_mix_var.set(data['is_demucs_pre_proc_model_inst_mix'])804 805 if not process_method or process_method == MDX_ARCH_TYPE or is_ensemble:806 self.mdx_net_model_var.set(loaded_setting['mdx_net_model'])807 self.chunks_var.set(loaded_setting['chunks'])808 self.margin_var.set(loaded_setting['margin'])809 self.compensate_var.set(loaded_setting['compensate'])810 self.is_denoise_var.set(loaded_setting['is_denoise'])811 self.is_invert_spec_var.set(loaded_setting['is_invert_spec'])812 self.is_mixer_mode_var.set(loaded_setting['is_mixer_mode'])813 self.mdx_batch_size_var.set(loaded_setting['mdx_batch_size'])814 self.mdx_voc_inst_secondary_model_var.set(loaded_setting['mdx_voc_inst_secondary_model'])815 self.mdx_other_secondary_model_var.set(loaded_setting['mdx_other_secondary_model'])816 self.mdx_bass_secondary_model_var.set(loaded_setting['mdx_bass_secondary_model'])817 self.mdx_drums_secondary_model_var.set(loaded_setting['mdx_drums_secondary_model'])818 self.mdx_is_secondary_model_activate_var.set(loaded_setting['mdx_is_secondary_model_activate'])819 self.mdx_voc_inst_secondary_model_scale_var.set(loaded_setting['mdx_voc_inst_secondary_model_scale'])820 self.mdx_other_secondary_model_scale_var.set(loaded_setting['mdx_other_secondary_model_scale'])821 self.mdx_bass_secondary_model_scale_var.set(loaded_setting['mdx_bass_secondary_model_scale'])822 self.mdx_drums_secondary_model_scale_var.set(loaded_setting['mdx_drums_secondary_model_scale'])823 824 if not process_method or is_ensemble:825 self.is_save_all_outputs_ensemble_var.set(loaded_setting['is_save_all_outputs_ensemble'])826 self.is_append_ensemble_name_var.set(loaded_setting['is_append_ensemble_name'])827 self.chosen_audio_tool_var.set(loaded_setting['chosen_audio_tool'])828 self.choose_algorithm_var.set(loaded_setting['choose_algorithm'])829 self.time_stretch_rate_var.set(loaded_setting['time_stretch_rate'])830 self.pitch_rate_var.set(loaded_setting['pitch_rate'])831 self.is_primary_stem_only_var.set(loaded_setting['is_primary_stem_only'])832 self.is_secondary_stem_only_var.set(loaded_setting['is_secondary_stem_only'])833 self.is_testing_audio_var.set(loaded_setting['is_testing_audio'])834 self.is_add_model_name_var.set(loaded_setting['is_add_model_name'])835 self.is_accept_any_input_var.set(loaded_setting["is_accept_any_input"])836 self.is_task_complete_var.set(loaded_setting['is_task_complete'])837 self.is_create_model_folder_var.set(loaded_setting['is_create_model_folder'])838 self.mp3_bit_set_var.set(loaded_setting['mp3_bit_set'])839 self.save_format_var.set(loaded_setting['save_format'])840 self.wav_type_set_var.set(loaded_setting['wav_type_set'])841 self.user_code_var.set(loaded_setting['user_code'])842 843 self.is_gpu_conversion_var.set(loaded_setting['is_gpu_conversion'])844 self.is_normalization_var.set(loaded_setting['is_normalization'])845 self.help_hints_var.set(loaded_setting['help_hints_var'])846 847 self.model_sample_mode_var.set(loaded_setting['model_sample_mode'])848 self.model_sample_mode_duration_var.set(loaded_setting['model_sample_mode_duration'])849 850 851root = FakeRoot()852root.load_saved_settings(DEFAULT_DATA)