273 lines
15 KiB
Python
273 lines
15 KiB
Python
class ModelData():
|
|
def __init__(self, model_name: str,
|
|
selected_process_method=ENSEMBLE_MODE,
|
|
is_secondary_model=False,
|
|
primary_model_primary_stem=None,
|
|
is_primary_model_primary_stem_only=False,
|
|
is_primary_model_secondary_stem_only=False,
|
|
is_pre_proc_model=False,
|
|
is_dry_check=False):
|
|
|
|
self.is_gpu_conversion = 0 if root.is_gpu_conversion_var.get() else -1
|
|
self.is_normalization = root.is_normalization_var.get()
|
|
self.is_primary_stem_only = root.is_primary_stem_only_var.get()
|
|
self.is_secondary_stem_only = root.is_secondary_stem_only_var.get()
|
|
self.is_denoise = root.is_denoise_var.get()
|
|
self.mdx_batch_size = 1 if root.mdx_batch_size_var.get() == DEF_OPT else int(root.mdx_batch_size_var.get())
|
|
self.is_mdx_ckpt = False
|
|
self.wav_type_set = root.wav_type_set
|
|
self.mp3_bit_set = root.mp3_bit_set_var.get()
|
|
self.save_format = root.save_format_var.get()
|
|
self.is_invert_spec = root.is_invert_spec_var.get()
|
|
self.is_mixer_mode = root.is_mixer_mode_var.get()
|
|
self.demucs_stems = root.demucs_stems_var.get()
|
|
self.demucs_source_list = []
|
|
self.demucs_stem_count = 0
|
|
self.mixer_path = MDX_MIXER_PATH
|
|
self.model_name = model_name
|
|
self.process_method = selected_process_method
|
|
self.model_status = False if self.model_name == CHOOSE_MODEL or self.model_name == NO_MODEL else True
|
|
self.primary_stem = None
|
|
self.secondary_stem = None
|
|
self.is_ensemble_mode = False
|
|
self.ensemble_primary_stem = None
|
|
self.ensemble_secondary_stem = None
|
|
self.primary_model_primary_stem = primary_model_primary_stem
|
|
self.is_secondary_model = is_secondary_model
|
|
self.secondary_model = None
|
|
self.secondary_model_scale = None
|
|
self.demucs_4_stem_added_count = 0
|
|
self.is_demucs_4_stem_secondaries = False
|
|
self.is_4_stem_ensemble = False
|
|
self.pre_proc_model = None
|
|
self.pre_proc_model_activated = False
|
|
self.is_pre_proc_model = is_pre_proc_model
|
|
self.is_dry_check = is_dry_check
|
|
self.model_samplerate = 44100
|
|
self.model_capacity = 32, 128
|
|
self.is_vr_51_model = False
|
|
self.is_demucs_pre_proc_model_inst_mix = False
|
|
self.manual_download_Button = None
|
|
self.secondary_model_4_stem = []
|
|
self.secondary_model_4_stem_scale = []
|
|
self.secondary_model_4_stem_names = []
|
|
self.secondary_model_4_stem_model_names_list = []
|
|
self.all_models = []
|
|
self.secondary_model_other = None
|
|
self.secondary_model_scale_other = None
|
|
self.secondary_model_bass = None
|
|
self.secondary_model_scale_bass = None
|
|
self.secondary_model_drums = None
|
|
self.secondary_model_scale_drums = None
|
|
|
|
if selected_process_method == ENSEMBLE_MODE:
|
|
partitioned_name = model_name.partition(ENSEMBLE_PARTITION)
|
|
self.process_method = partitioned_name[0]
|
|
self.model_name = partitioned_name[2]
|
|
self.model_and_process_tag = model_name
|
|
self.ensemble_primary_stem, self.ensemble_secondary_stem = root.return_ensemble_stems()
|
|
self.is_ensemble_mode = True if not is_secondary_model and not is_pre_proc_model else False
|
|
self.is_4_stem_ensemble = True if root.ensemble_main_stem_var.get() == FOUR_STEM_ENSEMBLE and self.is_ensemble_mode else False
|
|
self.pre_proc_model_activated = root.is_demucs_pre_proc_model_activate_var.get() if not self.ensemble_primary_stem == VOCAL_STEM else False
|
|
|
|
if self.process_method == VR_ARCH_TYPE:
|
|
self.is_secondary_model_activated = root.vr_is_secondary_model_activate_var.get() if not self.is_secondary_model else False
|
|
self.aggression_setting = float(int(root.aggression_setting_var.get())/100)
|
|
self.is_tta = root.is_tta_var.get()
|
|
self.is_post_process = root.is_post_process_var.get()
|
|
self.window_size = int(root.window_size_var.get())
|
|
self.batch_size = 1 if root.batch_size_var.get() == DEF_OPT else int(root.batch_size_var.get())
|
|
self.crop_size = int(root.crop_size_var.get())
|
|
self.is_high_end_process = 'mirroring' if root.is_high_end_process_var.get() else 'None'
|
|
self.post_process_threshold = float(root.post_process_threshold_var.get())
|
|
self.model_capacity = 32, 128
|
|
self.model_path = os.path.join(VR_MODELS_DIR, f"{self.model_name}.pth")
|
|
self.get_model_hash()
|
|
if self.model_hash:
|
|
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_PARAMS
|
|
if self.model_data:
|
|
vr_model_param = os.path.join(VR_PARAM_DIR, "{}.json".format(self.model_data["vr_model_param"]))
|
|
self.primary_stem = self.model_data["primary_stem"]
|
|
self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
|
|
self.vr_model_param = ModelParameters(vr_model_param)
|
|
self.model_samplerate = self.vr_model_param.param['sr']
|
|
if "nout" in self.model_data.keys() and "nout_lstm" in self.model_data.keys():
|
|
self.model_capacity = self.model_data["nout"], self.model_data["nout_lstm"]
|
|
self.is_vr_51_model = True
|
|
else:
|
|
self.model_status = False
|
|
|
|
if self.process_method == MDX_ARCH_TYPE:
|
|
self.is_secondary_model_activated = root.mdx_is_secondary_model_activate_var.get() if not is_secondary_model else False
|
|
self.margin = int(root.margin_var.get())
|
|
self.chunks = root.determine_auto_chunks(root.chunks_var.get(), self.is_gpu_conversion) if root.is_chunk_mdxnet_var.get() else 0
|
|
self.get_mdx_model_path()
|
|
self.get_model_hash()
|
|
if self.model_hash:
|
|
self.model_data = self.get_model_data(MDX_HASH_DIR, root.mdx_hash_MAPPER)
|
|
if self.model_data:
|
|
self.compensate = self.model_data["compensate"] if root.compensate_var.get() == AUTO_SELECT else float(root.compensate_var.get())
|
|
self.mdx_dim_f_set = self.model_data["mdx_dim_f_set"]
|
|
self.mdx_dim_t_set = self.model_data["mdx_dim_t_set"]
|
|
self.mdx_n_fft_scale_set = self.model_data["mdx_n_fft_scale_set"]
|
|
self.primary_stem = self.model_data["primary_stem"]
|
|
self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
|
|
else:
|
|
self.model_status = False
|
|
|
|
if self.process_method == DEMUCS_ARCH_TYPE:
|
|
self.is_secondary_model_activated = root.demucs_is_secondary_model_activate_var.get() if not is_secondary_model else False
|
|
if not self.is_ensemble_mode:
|
|
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 False
|
|
self.overlap = float(root.overlap_var.get())
|
|
self.margin_demucs = int(root.margin_demucs_var.get())
|
|
self.chunks_demucs = root.determine_auto_chunks(root.chunks_demucs_var.get(), self.is_gpu_conversion)
|
|
self.shifts = int(root.shifts_var.get())
|
|
self.is_split_mode = root.is_split_mode_var.get()
|
|
self.segment = root.segment_var.get()
|
|
self.is_chunk_demucs = root.is_chunk_demucs_var.get()
|
|
self.is_demucs_combine_stems = root.is_demucs_combine_stems_var.get()
|
|
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()
|
|
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()
|
|
self.get_demucs_model_path()
|
|
self.get_demucs_model_data()
|
|
|
|
self.model_basename = os.path.splitext(os.path.basename(self.model_path))[0] if self.model_status else None
|
|
self.pre_proc_model_activated = self.pre_proc_model_activated if not self.is_secondary_model else False
|
|
|
|
self.is_primary_model_primary_stem_only = is_primary_model_primary_stem_only
|
|
self.is_primary_model_secondary_stem_only = is_primary_model_secondary_stem_only
|
|
|
|
if self.is_secondary_model_activated and self.model_status:
|
|
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:
|
|
for key in DEMUCS_4_SOURCE_LIST:
|
|
self.secondary_model_data(key)
|
|
self.secondary_model_4_stem.append(self.secondary_model)
|
|
self.secondary_model_4_stem_scale.append(self.secondary_model_scale)
|
|
self.secondary_model_4_stem_names.append(key)
|
|
self.demucs_4_stem_added_count = sum(i is not None for i in self.secondary_model_4_stem)
|
|
self.is_secondary_model_activated = False if all(i is None for i in self.secondary_model_4_stem) else True
|
|
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_count
|
|
if self.is_secondary_model_activated:
|
|
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]
|
|
self.is_demucs_4_stem_secondaries = True
|
|
else:
|
|
primary_stem = self.ensemble_primary_stem if self.is_ensemble_mode and self.process_method == DEMUCS_ARCH_TYPE else self.primary_stem
|
|
self.secondary_model_data(primary_stem)
|
|
|
|
if self.process_method == DEMUCS_ARCH_TYPE and not is_secondary_model:
|
|
if self.demucs_stem_count >= 3 and self.pre_proc_model_activated:
|
|
self.pre_proc_model_activated = True
|
|
self.pre_proc_model = root.process_determine_demucs_pre_proc_model(self.primary_stem)
|
|
self.is_demucs_pre_proc_model_inst_mix = root.is_demucs_pre_proc_model_inst_mix_var.get() if self.pre_proc_model else False
|
|
|
|
def secondary_model_data(self, primary_stem):
|
|
secondary_model_data = root.process_determine_secondary_model(self.process_method, primary_stem, self.is_primary_stem_only, self.is_secondary_stem_only)
|
|
self.secondary_model = secondary_model_data[0]
|
|
self.secondary_model_scale = secondary_model_data[1]
|
|
self.is_secondary_model_activated = False if not self.secondary_model else True
|
|
if self.secondary_model:
|
|
self.is_secondary_model_activated = False if self.secondary_model.model_basename == self.model_basename else True
|
|
|
|
def get_mdx_model_path(self):
|
|
|
|
if self.model_name.endswith(CKPT):
|
|
# self.chunks = 0
|
|
# self.is_mdx_batch_mode = True
|
|
self.is_mdx_ckpt = True
|
|
|
|
ext = '' if self.is_mdx_ckpt else ONNX
|
|
|
|
for file_name, chosen_mdx_model in root.mdx_name_select_MAPPER.items():
|
|
if self.model_name in chosen_mdx_model:
|
|
self.model_path = os.path.join(MDX_MODELS_DIR, f"{file_name}{ext}")
|
|
break
|
|
else:
|
|
self.model_path = os.path.join(MDX_MODELS_DIR, f"{self.model_name}{ext}")
|
|
|
|
self.mixer_path = os.path.join(MDX_MODELS_DIR, f"mixer_val.ckpt")
|
|
|
|
def get_demucs_model_path(self):
|
|
|
|
demucs_newer = [True for x in DEMUCS_NEWER_TAGS if x in self.model_name]
|
|
demucs_model_dir = DEMUCS_NEWER_REPO_DIR if demucs_newer else DEMUCS_MODELS_DIR
|
|
|
|
for file_name, chosen_model in root.demucs_name_select_MAPPER.items():
|
|
if self.model_name in chosen_model:
|
|
self.model_path = os.path.join(demucs_model_dir, file_name)
|
|
break
|
|
else:
|
|
self.model_path = os.path.join(DEMUCS_NEWER_REPO_DIR, f'{self.model_name}.yaml')
|
|
|
|
def get_demucs_model_data(self):
|
|
|
|
self.demucs_version = DEMUCS_V4
|
|
|
|
for key, value in DEMUCS_VERSION_MAPPER.items():
|
|
if value in self.model_name:
|
|
self.demucs_version = key
|
|
|
|
self.demucs_source_list = DEMUCS_2_SOURCE if DEMUCS_UVR_MODEL in self.model_name else DEMUCS_4_SOURCE
|
|
self.demucs_source_map = DEMUCS_2_SOURCE_MAPPER if DEMUCS_UVR_MODEL in self.model_name else DEMUCS_4_SOURCE_MAPPER
|
|
self.demucs_stem_count = 2 if DEMUCS_UVR_MODEL in self.model_name else 4
|
|
|
|
if not self.is_ensemble_mode:
|
|
self.primary_stem = PRIMARY_STEM if self.demucs_stems == ALL_STEMS else self.demucs_stems
|
|
self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
|
|
|
|
def get_model_data(self, model_hash_dir, hash_mapper):
|
|
|
|
model_settings_json = os.path.join(model_hash_dir, "{}.json".format(self.model_hash))
|
|
|
|
if os.path.isfile(model_settings_json):
|
|
return json.load(open(model_settings_json))
|
|
else:
|
|
for hash, settings in hash_mapper.items():
|
|
if self.model_hash in hash:
|
|
return settings
|
|
else:
|
|
return self.get_model_data_from_popup()
|
|
|
|
def get_model_data_from_popup(self):
|
|
|
|
if not self.is_dry_check:
|
|
confirm = tk.messagebox.askyesno(title=UNRECOGNIZED_MODEL[0],
|
|
message=f"\"{self.model_name}\"{UNRECOGNIZED_MODEL[1]}",
|
|
parent=root)
|
|
|
|
if confirm:
|
|
if self.process_method == VR_ARCH_TYPE:
|
|
root.pop_up_vr_param(self.model_hash)
|
|
return root.vr_model_params
|
|
if self.process_method == MDX_ARCH_TYPE:
|
|
root.pop_up_mdx_model(self.model_hash, self.model_path)
|
|
return root.mdx_model_params
|
|
else:
|
|
return None
|
|
else:
|
|
return None
|
|
|
|
def get_model_hash(self):
|
|
self.model_hash = None
|
|
|
|
if not os.path.isfile(self.model_path):
|
|
self.model_status = False
|
|
self.model_hash is None
|
|
else:
|
|
if model_hash_table:
|
|
for (key, value) in model_hash_table.items():
|
|
if self.model_path == key:
|
|
self.model_hash = value
|
|
break
|
|
|
|
if not self.model_hash:
|
|
try:
|
|
with open(self.model_path, 'rb') as f:
|
|
f.seek(- 10000 * 1024, 2)
|
|
self.model_hash = hashlib.md5(f.read()).hexdigest()
|
|
except:
|
|
self.model_hash = hashlib.md5(open(self.model_path,'rb').read()).hexdigest()
|
|
|
|
table_entry = {self.model_path: self.model_hash}
|
|
model_hash_table.update(table_entry)
|