| import os | |
| import sys | |
| import time | |
| import yaml | |
| import torch | |
| import codecs | |
| import hashlib | |
| import logging | |
| import platform | |
| import warnings | |
| import requests | |
| import onnxruntime | |
| from importlib import metadata, import_module | |
| now_dir = os.getcwd() | |
| sys.path.append(now_dir) | |
| from main.configs.config import Config | |
| translations = Config().translations | |
| class Separator: | |
| def __init__(self, logger=logging.getLogger(__name__), log_level=logging.INFO, log_formatter=None, model_file_dir="assets/models/uvr5", output_dir=None, output_format="wav", output_bitrate=None, normalization_threshold=0.9, output_single_stem=None, invert_using_spec=False, sample_rate=44100, mdx_params={"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False}, demucs_params={"segment_size": "Default", "shifts": 2, "overlap": 0.25, "segments_enabled": True}): | |
| self.logger = logger | |
| self.log_level = log_level | |
| self.log_formatter = log_formatter | |
| self.log_handler = logging.StreamHandler() | |
| if self.log_formatter is None: self.log_formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(module)s - %(message)s") | |
| self.log_handler.setFormatter(self.log_formatter) | |
| if not self.logger.hasHandlers(): self.logger.addHandler(self.log_handler) | |
| if log_level > logging.DEBUG: warnings.filterwarnings("ignore") | |
| self.logger.info(translations["separator_info"].format(output_dir=output_dir, output_format=output_format)) | |
| self.model_file_dir = model_file_dir | |
| if output_dir is None: | |
| output_dir = now_dir | |
| self.logger.info(translations["output_dir_is_none"]) | |
| self.output_dir = output_dir | |
| os.makedirs(self.model_file_dir, exist_ok=True) | |
| os.makedirs(self.output_dir, exist_ok=True) | |
| self.output_format = output_format | |
| self.output_bitrate = output_bitrate | |
| if self.output_format is None: self.output_format = "wav" | |
| self.normalization_threshold = normalization_threshold | |
| if normalization_threshold <= 0 or normalization_threshold > 1: raise ValueError(translations[">0or=1"]) | |
| self.output_single_stem = output_single_stem | |
| if output_single_stem is not None: self.logger.debug(translations["output_single"].format(output_single_stem=output_single_stem)) | |
| self.invert_using_spec = invert_using_spec | |
| if self.invert_using_spec: self.logger.debug(translations["step2"]) | |
| self.sample_rate = int(sample_rate) | |
| self.arch_specific_params = {"MDX": mdx_params, "Demucs": demucs_params} | |
| self.torch_device = None | |
| self.torch_device_cpu = None | |
| self.torch_device_mps = None | |
| self.onnx_execution_provider = None | |
| self.model_instance = None | |
| self.model_is_uvr_vip = False | |
| self.model_friendly_name = None | |
| self.setup_accelerated_inferencing_device() | |
| def setup_accelerated_inferencing_device(self): | |
| system_info = self.get_system_info() | |
| self.log_onnxruntime_packages() | |
| self.setup_torch_device(system_info) | |
| def get_system_info(self): | |
| os_name = platform.system() | |
| os_version = platform.version() | |
| self.logger.info(f"{translations['os']}: {os_name} {os_version}") | |
| system_info = platform.uname() | |
| self.logger.info(translations["platform_info"].format(system_info=system_info, node=system_info.node, release=system_info.release, machine=system_info.machine, processor=system_info.processor)) | |
| python_version = platform.python_version() | |
| self.logger.info(f"{translations['name_ver'].format(name='python')}: {python_version}") | |
| pytorch_version = torch.__version__ | |
| self.logger.info(f"{translations['name_ver'].format(name='pytorch')}: {pytorch_version}") | |
| return system_info | |
| def log_onnxruntime_packages(self): | |
| onnxruntime_gpu_package = self.get_package_distribution("onnxruntime-gpu") | |
| onnxruntime_cpu_package = self.get_package_distribution("onnxruntime") | |
| if onnxruntime_gpu_package is not None: self.logger.info(f"{translations['install_onnx'].format(pu='GPU')}: {onnxruntime_gpu_package.version}") | |
| if onnxruntime_cpu_package is not None: self.logger.info(f"{translations['install_onnx'].format(pu='CPU')}: {onnxruntime_cpu_package.version}") | |
| def setup_torch_device(self, system_info): | |
| hardware_acceleration_enabled = False | |
| ort_providers = onnxruntime.get_available_providers() | |
| self.torch_device_cpu = torch.device("cpu") | |
| if torch.cuda.is_available(): | |
| self.configure_cuda(ort_providers) | |
| hardware_acceleration_enabled = True | |
| elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available() and system_info.processor == "arm": | |
| self.configure_mps(ort_providers) | |
| hardware_acceleration_enabled = True | |
| if not hardware_acceleration_enabled: | |
| self.logger.info(translations["running_in_cpu"]) | |
| self.torch_device = self.torch_device_cpu | |
| self.onnx_execution_provider = ["CPUExecutionProvider"] | |
| def configure_cuda(self, ort_providers): | |
| self.logger.info(translations["running_in_cuda"]) | |
| self.torch_device = torch.device("cuda") | |
| if "CUDAExecutionProvider" in ort_providers: | |
| self.logger.info(translations["onnx_have"].format(have='CUDAExecutionProvider')) | |
| self.onnx_execution_provider = ["CUDAExecutionProvider"] | |
| else: self.logger.warning(translations["onnx_not_have"].format(have='CUDAExecutionProvider')) | |
| def configure_mps(self, ort_providers): | |
| self.logger.info(translations["set_torch_mps"]) | |
| self.torch_device_mps = torch.device("mps") | |
| self.torch_device = self.torch_device_mps | |
| if "CoreMLExecutionProvider" in ort_providers: | |
| self.logger.info(translations["onnx_have"].format(have='CoreMLExecutionProvider')) | |
| self.onnx_execution_provider = ["CoreMLExecutionProvider"] | |
| else: self.logger.warning(translations["onnx_not_have"].format(have='CoreMLExecutionProvider')) | |
| def get_package_distribution(self, package_name): | |
| try: | |
| return metadata.distribution(package_name) | |
| except metadata.PackageNotFoundError: | |
| self.logger.debug(translations["python_not_install"].format(package_name=package_name)) | |
| return None | |
| def get_model_hash(self, model_path): | |
| self.logger.debug(translations["hash"].format(model_path=model_path)) | |
| try: | |
| with open(model_path, "rb") as f: | |
| f.seek(-10000 * 1024, 2) | |
| return hashlib.md5(f.read()).hexdigest() | |
| except IOError as e: | |
| self.logger.error(translations["ioerror"].format(e=e)) | |
| return hashlib.md5(open(model_path, "rb").read()).hexdigest() | |
| def download_file_if_not_exists(self, url, output_path): | |
| if os.path.isfile(output_path): | |
| self.logger.debug(translations["cancel_download"].format(output_path=output_path)) | |
| return | |
| self.logger.debug(translations["download_model"].format(url=url, output_path=output_path)) | |
| response = requests.get(url, stream=True, timeout=300) | |
| if response.status_code == 200: | |
| from tqdm import tqdm | |
| progress_bar = tqdm(total=int(response.headers.get("content-length", 0)), ncols=100, unit="byte") | |
| with open(output_path, "wb") as f: | |
| for chunk in response.iter_content(chunk_size=8192): | |
| progress_bar.update(len(chunk)) | |
| f.write(chunk) | |
| progress_bar.close() | |
| else: raise RuntimeError(translations["download_error"].format(url=url, status_code=response.status_code)) | |
| def print_uvr_vip_message(self): | |
| if self.model_is_uvr_vip: | |
| self.logger.warning(translations["vip_model"].format(model_friendly_name=self.model_friendly_name)) | |
| self.logger.warning(translations["vip_print"]) | |
| def list_supported_model_files(self): | |
| response = requests.get(codecs.decode("uggcf://uhttvatsnpr.pb/NauC/Ivrganzrfr-EIP-Cebwrpg/enj/znva/wfba/hie_zbqryf.wfba", "rot13")) | |
| response.raise_for_status() | |
| model_downloads_list = response.json() | |
| self.logger.debug(translations["load_download_json"]) | |
| return {"MDX": {**model_downloads_list["mdx_download_list"], **model_downloads_list["mdx_download_vip_list"]}, "Demucs": {key: value for key, value in model_downloads_list["demucs_download_list"].items() if key.startswith("Demucs v4")}} | |
| def download_model_files(self, model_filename): | |
| model_path = os.path.join(self.model_file_dir, model_filename) | |
| supported_model_files_grouped = self.list_supported_model_files() | |
| yaml_config_filename = None | |
| self.logger.debug(translations["search_model"].format(model_filename=model_filename)) | |
| for model_type, model_list in supported_model_files_grouped.items(): | |
| for model_friendly_name, model_download_list in model_list.items(): | |
| self.model_is_uvr_vip = "VIP" in model_friendly_name | |
| model_repo_url_prefix = codecs.decode("uggcf://uhttvatsnpr.pb/NauC/Ivrganzrfr-EIP-Cebwrpg/erfbyir/znva/hie5_zbqryf", "rot13") | |
| if isinstance(model_download_list, str) and model_download_list == model_filename: | |
| self.logger.debug(translations["single_model"].format(model_friendly_name=model_friendly_name)) | |
| self.model_friendly_name = model_friendly_name | |
| try: | |
| self.download_file_if_not_exists(f"{model_repo_url_prefix}/MDX/{model_filename}", model_path) | |
| except RuntimeError: | |
| self.logger.warning(translations["not_found_model"]) | |
| self.download_file_if_not_exists(f"{model_repo_url_prefix}/Demucs/{model_filename}", model_path) | |
| self.print_uvr_vip_message() | |
| self.logger.debug(translations["single_model_path"].format(model_path=model_path)) | |
| return model_filename, model_type, model_friendly_name, model_path, yaml_config_filename | |
| elif isinstance(model_download_list, dict): | |
| this_model_matches_input_filename = False | |
| for file_name, file_url in model_download_list.items(): | |
| if file_name == model_filename or file_url == model_filename: | |
| self.logger.debug(translations["find_model"].format(model_filename=model_filename, model_friendly_name=model_friendly_name)) | |
| this_model_matches_input_filename = True | |
| if this_model_matches_input_filename: | |
| self.logger.debug(translations["find_models"].format(model_friendly_name=model_friendly_name)) | |
| self.model_friendly_name = model_friendly_name | |
| self.print_uvr_vip_message() | |
| for config_key, config_value in model_download_list.items(): | |
| self.logger.debug(f"{translations['find_path']}: {config_key} -> {config_value}") | |
| if config_value.startswith("http"): self.download_file_if_not_exists(config_value, os.path.join(self.model_file_dir, config_key)) | |
| elif config_key.endswith(".ckpt"): | |
| try: | |
| self.download_file_if_not_exists(f"{model_repo_url_prefix}/Demucs/{config_key}", os.path.join(self.model_file_dir, config_key)) | |
| except RuntimeError: | |
| self.logger.warning(translations["not_found_model_warehouse"]) | |
| if model_filename.endswith(".yaml"): | |
| self.logger.warning(translations["yaml_warning"].format(model_filename=model_filename)) | |
| self.logger.warning(translations["yaml_warning_2"].format(config_key=config_key)) | |
| self.logger.warning(translations["yaml_warning_3"]) | |
| model_filename = config_key | |
| model_path = os.path.join(self.model_file_dir, f"{model_filename}") | |
| yaml_config_filename = config_value | |
| yaml_config_filepath = os.path.join(self.model_file_dir, yaml_config_filename) | |
| try: | |
| self.download_file_if_not_exists(f"{model_repo_url_prefix}/mdx_c_configs/{yaml_config_filename}", yaml_config_filepath) | |
| except RuntimeError: | |
| self.logger.debug(translations["yaml_debug"]) | |
| else: self.download_file_if_not_exists(f"{model_repo_url_prefix}/Demucs/{config_value}", os.path.join(self.model_file_dir, config_value)) | |
| self.logger.debug(translations["download_model_friendly"].format(model_friendly_name=model_friendly_name, model_path=model_path)) | |
| return model_filename, model_type, model_friendly_name, model_path, yaml_config_filename | |
| raise ValueError(translations["not_found_model_2"].format(model_filename=model_filename)) | |
| def load_model_data_from_yaml(self, yaml_config_filename): | |
| model_data_yaml_filepath = os.path.join(self.model_file_dir, yaml_config_filename) if not os.path.exists(yaml_config_filename) else yaml_config_filename | |
| self.logger.debug(translations["load_yaml"].format(model_data_yaml_filepath=model_data_yaml_filepath)) | |
| model_data = yaml.load(open(model_data_yaml_filepath, encoding="utf-8"), Loader=yaml.FullLoader) | |
| self.logger.debug(translations["load_yaml_2"].format(model_data=model_data)) | |
| if "roformer" in model_data_yaml_filepath: model_data["is_roformer"] = True | |
| return model_data | |
| def load_model_data_using_hash(self, model_path): | |
| self.logger.debug(translations["hash_md5"]) | |
| model_hash = self.get_model_hash(model_path) | |
| self.logger.debug(translations["model_hash"].format(model_path=model_path, model_hash=model_hash)) | |
| mdx_model_data_path = codecs.decode("uggcf://uhttvatsnpr.pb/NauC/Ivrganzrfr-EIP-Cebwrpg/enj/znva/wfba/zbqry_qngn.wfba", "rot13") | |
| self.logger.debug(translations["mdx_data"].format(mdx_model_data_path=mdx_model_data_path)) | |
| response = requests.get(mdx_model_data_path) | |
| response.raise_for_status() | |
| mdx_model_data_object = response.json() | |
| self.logger.debug(translations["load_mdx"]) | |
| if model_hash in mdx_model_data_object: model_data = mdx_model_data_object[model_hash] | |
| else: raise ValueError(translations["model_not_support"].format(model_hash=model_hash)) | |
| self.logger.debug(translations["uvr_json"].format(model_hash=model_hash, model_data=model_data)) | |
| return model_data | |
| def load_model(self, model_filename): | |
| self.logger.info(translations["loading_model"].format(model_filename=model_filename)) | |
| load_model_start_time = time.perf_counter() | |
| model_filename, model_type, model_friendly_name, model_path, yaml_config_filename = self.download_model_files(model_filename) | |
| self.logger.debug(translations["download_model_friendly_2"].format(model_friendly_name=model_friendly_name, model_path=model_path)) | |
| if model_path.lower().endswith(".yaml"): yaml_config_filename = model_path | |
| common_params = {"logger": self.logger, "log_level": self.log_level, "torch_device": self.torch_device, "torch_device_cpu": self.torch_device_cpu, "torch_device_mps": self.torch_device_mps, "onnx_execution_provider": self.onnx_execution_provider, "model_name": model_filename.split(".")[0], "model_path": model_path, "model_data": self.load_model_data_from_yaml(yaml_config_filename) if yaml_config_filename is not None else self.load_model_data_using_hash(model_path), "output_format": self.output_format, "output_bitrate": self.output_bitrate, "output_dir": self.output_dir, "normalization_threshold": self.normalization_threshold, "output_single_stem": self.output_single_stem, "invert_using_spec": self.invert_using_spec, "sample_rate": self.sample_rate} | |
| separator_classes = {"MDX": "mdx_separator.MDXSeparator", "Demucs": "demucs_separator.DemucsSeparator"} | |
| if model_type not in self.arch_specific_params or model_type not in separator_classes: raise ValueError(translations["model_type_not_support"].format(model_type=model_type)) | |
| if model_type == "Demucs" and sys.version_info < (3, 10): raise Exception(translations["demucs_not_support_python<3.10"]) | |
| self.logger.debug(f"{translations['import_module']} {model_type}: {separator_classes[model_type]}") | |
| module_name, class_name = separator_classes[model_type].split(".") | |
| separator_class = getattr(import_module(f"main.library.architectures.{module_name}"), class_name) | |
| self.logger.debug(f"{translations['initialization']} {model_type}: {separator_class}") | |
| self.model_instance = separator_class(common_config=common_params, arch_config=self.arch_specific_params[model_type]) | |
| self.logger.debug(translations["loading_model_success"]) | |
| self.logger.info(f"{translations['loading_model_duration']}: {time.strftime('%H:%M:%S', time.gmtime(int(time.perf_counter() - load_model_start_time)))}") | |
| def separate(self, audio_file_path): | |
| self.logger.info(f"{translations['starting_separator']}: {audio_file_path}") | |
| separate_start_time = time.perf_counter() | |
| self.logger.debug(translations["normalization"].format(normalization_threshold=self.normalization_threshold)) | |
| output_files = self.model_instance.separate(audio_file_path) | |
| self.model_instance.clear_gpu_cache() | |
| self.model_instance.clear_file_specific_paths() | |
| self.print_uvr_vip_message() | |
| self.logger.debug(translations["separator_success_3"]) | |
| self.logger.info(f"{translations['separator_duration']}: {time.strftime('%H:%M:%S', time.gmtime(int(time.perf_counter() - separate_start_time)))}") | |
| return output_files | |
| def download_model_and_data(self, model_filename): | |
| self.logger.info(translations["loading_separator_model"].format(model_filename=model_filename)) | |
| model_filename, model_type, model_friendly_name, model_path, yaml_config_filename = self.download_model_files(model_filename) | |
| if model_path.lower().endswith(".yaml"): yaml_config_filename = model_path | |
| self.logger.info(translations["downloading_model"].format(model_type=model_type, model_friendly_name=model_friendly_name, model_path=model_path, model_data_dict_size=len(self.load_model_data_from_yaml(yaml_config_filename) if yaml_config_filename is not None else self.load_model_data_using_hash(model_path)))) |