import sys import os import yaml import time import traceback import gc gc.enable() import librosa from pathlib import Path BASE_DIR = Path(__file__).resolve().parent sys.path.append(str(BASE_DIR)) from extra_utils import hf_spaces_gpu, dw_file, extra_clear_torch_cache, nuclear_clear_model, emergency_ram_clear import torch import rich nn = torch.nn import json from tqdm import tqdm import numpy as np from typing import Literal, Optional, List, Tuple, Any, Dict from ml_collections import ConfigDict from omegaconf import OmegaConf import gradio as gr from audio import read, write, output_formats, subtractor, check, easy_resampler, ensemble_types, ensemble, multiread, get_audio_files_from_list, stereo_to_mono from args_parser import parse_separator_args, tobool from namer import Namer from i18n import _i18n import contextlib class PathNotExist(Exception): pass class PathsNotExist(Exception): pass class PathNotSpecified(Exception): pass class PathsNotSpecified(Exception): pass class FileIsNotAudio(Exception): pass class FilesIsNotAudio(Exception): pass class MixNotFound(Exception): pass class MixIsEmpty(Exception): pass class UnknownModelType(Exception): pass class DemixError(Exception): pass class ConfigNotLoaded(Exception): pass class ModelNotLoaded(Exception): pass class ModelStateDictError(Exception): pass HAS_OLD_AMP = False if hasattr(torch, "cuda"): if hasattr(torch.cuda, "amp"): if hasattr(torch.cuda.amp, "autocast"): HAS_OLD_AMP = True HAS_NEW_AMP = False if hasattr(torch, "amp"): if hasattr(torch.amp, "autocast"): HAS_NEW_AMP = True def get_autocast_context(device_type="cuda", enabled=True): if HAS_NEW_AMP: return torch.amp.autocast(device_type=device_type, enabled=enabled) elif HAS_OLD_AMP: return torch.cuda.amp.autocast(enabled=enabled) else: # Если AMP не поддерживается вообще return contextlib.nullcontext() # Или пустой контекст def _getWindowingArray(window_size: int, fade_size: int) -> torch.Tensor: """ Создать массив окна для плавного склеивания Args: window_size: Размер окна fade_size: Размер зоны затухания Returns: Массив окна """ fadein = torch.linspace(0, 1, fade_size) fadeout = torch.linspace(1, 0, fade_size) window = torch.ones(window_size) window[-fade_size:] = fadeout window[:fade_size] = fadein return window base_params = { "sec": { "type": "float", "component": "number", "minimum": 1, "maximum": 30, "step": 0.1, "default": 7, "info": "separation_segment_size_info" }, "size": { "type": "int", "component": "slider", "minimum": 128, "maximum": 1024, "step": 128, "default": 256, "info": "separation_segment_size_info" }, "wsize": { "type": "int", "component": "slider", "minimum": 320, "maximum": 1024, "step": 64, "default": 512, "info": "separation_window_size_info" }, "hop": { "type": "int", "component": "slider", "minimum": 512, "maximum": 2048, "step": 512, "default": 1024, "info": "separation_hop_info" }, "overlap": { "type": "int", "component": "slider", "minimum": 1, "maximum": 16, "step": 1, "default": 2, "info": "separation_overlap_info" }, "batch": { "type": "int", "component": "slider", "minimum": 1, "maximum": 16, "step": 1, "default": 1, "info": "separation_batch_size_info" }, "threshold" : { "type": "float", "component": "slider", "minimum": 0.1, "maximum": 0.3, "step": 0.1, "default": 0.2, }, "aggression": { "type": "int", "component": "slider", "minimum": 0, "maximum": 100, "step": 1, "default": 5, "info": "separation_aggresion_info" }, "enable": { "type": "bool", "component": "checkbox", "default": True }, "disable": { "type": "bool", "component": "checkbox", "default": False }, } add_params = { "mdxc": { "mdxc_segment_size": base_params["size"], "mdxc_batch_size": base_params["batch"], "mdxc_overlap": base_params["overlap"], "mdxc_denoise": base_params["disable"], "mdxc_override_segment": base_params["disable"] }, "demucs": { "demucs_segment": base_params["sec"], "demucs_batch_size": base_params["batch"], "demucs_overlap": base_params["overlap"], "demucs_denoise": base_params["disable"], "demucs_override_segment": base_params["disable"] }, "mdx": { "mdx_hop_length": base_params["hop"], "mdx_segment_size": base_params["size"], "mdx_batch_size": base_params["batch"], "mdx_overlap": base_params["overlap"], "mdx_denoise": base_params["disable"], "mdx_override_segment": base_params["disable"] }, "vr": { "vr_window_size": base_params["wsize"], "vr_batch_size": base_params["batch"], "vr_aggression": base_params["aggression"], "vr_post_process": base_params["disable"], "vr_post_process_threshold": base_params["threshold"], "vr_high_end_process": {**base_params["disable"], "info": "separation_hi-end_process_info"} }, "mvox": { "mvox_segment": base_params["sec"], "mvox_overlap": base_params["overlap"], "mvox_override_segment": base_params["disable"] } } add_params_list = [] add_params_group = [] add_params_args = {} for t_tab, t_components in add_params.items(): add_params_group.append(t_tab) for t_component, t_settings in t_components.items(): add_params_list.append(t_component) add_params_args[t_component] = {"default": t_settings["default"], "type": t_settings["type"]} def get_add_params(args): """Безопасно получает add_params из args""" if hasattr(args, 'add_params') and args.add_params is not None: return vars(args.add_params) return {} class MSSI: # Music Source Separation Inference def __init__(self, output_dir=".", output_format=output_formats[0], use_spec_invert=False, device="cuda" if torch.cuda.is_available() else "cpu", ): self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.model_types = ( "mel_band_roformer", "bs_roformer", "mdx23c", "scnet", "scnet_masked", "scnet_tran", "htdemucs", "bandit", "bandit_v2", "mdxnet", "vr", "medley_vox" ) self.custom_model_types = self.model_types[:9] self.output_format = output_format self.device = torch.device(device) self.use_spec_invert = use_spec_invert self.model = None self.model_module = None self.state_dict = {} self.model_loaded = False self.model_type = None self.ckpt_path = None self.conf_path = None self.config = None self.target_instrument = None self.instruments = [] self.input_mix = None self.input_file_name = None self.sample_rate = None self.selected_instruments = [] self.output_files_list = [] self.add_params = {} self.output_arrays: dict[str, np.ndarray] = {} def settings(self, output_dir=".", output_format=output_formats[0], use_spec_invert=False, ): self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.output_format = output_format self.use_spec_invert = use_spec_invert def set_add_params(self, **kwargs): self.add_params = kwargs def load_config(self, model_type: str, conf: str | Path): if not conf: raise PathNotSpecified(_i18n("path_not_specified")) self.conf_path = Path(conf) if not self.conf_path.exists(): self.conf_path = None raise PathNotExist(_i18n("path_not_exist")) if model_type not in self.model_types: raise UnknownModelType(_i18n("unknown_model_type", model_type=model_type)) self.model_type = model_type try: if self.model_type == "htdemucs": self.config = OmegaConf.load(self.conf_path) self.sample_rate = self.config.training.samplerate else: with self.conf_path.open("r", encoding="utf-8") as f: self.config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader)) self.sample_rate = self.config.audio.sample_rate self.target_instrument = self.config.training.target_instrument self.instruments = self.config.training.instruments print(_i18n("config_loaded")+": "+self.conf_path.name) except FileNotFoundError: self.config = None self.conf_path = None self.model_type = None self.target_instrument = None self.instruments = [] self.sample_rate = None raise FileNotFoundError(_i18n("config_not_found", path=conf)) from e except Exception as e: self.config = None self.conf_path = None self.model_type = None self.target_instrument = None self.instruments = [] self.sample_rate = None raise ValueError(_i18n("config_load_error", error=str(e))) from e def prefer_target_instrument(self): if self.target_instrument: return [self.target_instrument] else: return self.instruments def print_instruments(self): print(_i18n("stems")+": "+",".join(self.instruments)) print(_i18n("target_instrument")+": "+(self.target_instrument if self.target_instrument else _i18n("no"))) def load_model_instance(self): if self.config is None or self.model_type is None: raise ConfigNotLoaded(_i18n("config_is_not_loaded")) if self.model_type == "mdx23c": from models import mdx23c_tfc_tdf_v3 as module self.model_module = module.TFC_TDF_net self.model = self.model_module(self.config) del module elif self.model_type == "mdxnet": from models import mdx_net as module self.model_module = module.MDXNet self.model = self.model_module(**dict(self.config.model)) del module elif self.model_type == "vr": from models import vr_arch as module self.model_module = module.get_model self.model = self.model_module(self.config) del module elif self.model_type == "htdemucs": models_path = BASE_DIR / 'models' sys.path.append(str(models_path)) from demucs import get_model as module self.model_module = module self.model = self.model_module(self.config) del module elif self.model_type == "mel_band_roformer": if hasattr(self.config, "windowed"): from models.windowed_roformer import model as module self.model_module = module.MelBandRoformerWSA self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "conformer"): from models import bs_roformer as module self.model_module = module.MelBandConformer self.model = self.model_module(**dict(self.config.model)) del module else: from models import bs_roformer as module self.model_module = module.MelBandRoformer self.model = self.model_module(**dict(self.config.model)) del module elif self.model_type == "bs_roformer": if hasattr(self.config, "sw"): from models import bs_roformer as module self.model_module = module.BSRoformer_SW self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "fno"): from models import bs_roformer as module self.model_module = module.BSRoformer_FNO self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "hyperace"): from models import bs_roformer as module self.model_module = module.BSRoformerHyperACE self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "hyperace2"): from models import bs_roformer as module self.model_module = module.BSRoformerHyperACE_2 self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "conformer"): from models import bs_roformer as module self.model_module = module.BSConformer self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "conditional"): from models import bs_roformer as module self.model_module = module.BSRoformer_Conditional self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "unwa_inst_large_2"): from models import bs_roformer as module self.model_module = module.BSRoformer_2 self.model = self.model_module(**dict(self.config.model)) del module elif hasattr(self.config, "siamese"): from models import bs_roformer as module self.model_module = module.BSSiameseRoformer self.model = self.model_module(**dict(self.config.model)) del module else: from models import bs_roformer as module self.model_module = module.BSRoformer self.model = self.model_module(**dict(self.config.model)) del module elif self.model_type == "bandit": from models.bandit.core import model as module self.model_module = module.MultiMaskMultiSourceBandSplitRNNSimple self.model = self.model_module(**self.config.model) del module elif self.model_type == "bandit_v2": from models.bandit_v2 import bandit as module self.model_module = module.Bandit self.model = self.model_module(**self.config.kwargs) del module elif self.model_type == "scnet_unofficial": from models import scnet_unofficial as module self.model_module = module.SCNet self.model = self.model_module(**self.config.model) del module elif self.model_type == "scnet": from models import scnet as module self.model_module = module.SCNet self.model = self.model_module(**self.config.model) del module elif self.model_type == 'scnet_masked': from models.scnet import scnet_masked as module self.model_module = module.SCNet self.model = self.model_module(**self.config.model) del module elif self.model_type == 'scnet_tran': from models.scnet import scnet_tran as module self.model_module = module.SCNet_Tran self.model = self.model_module(**self.config.model) del module elif self.model_type == 'medley_vox': from models import medley_vox as module self.model_module = module.load_model_with_args self.model = self.model_module(self.config.model) del module else: raise UnknownModelType(_i18n("unknown_model_type", model_type=self.model_type)) def clear_model(self): """Стандартная очистка (сохранена для совместимости)""" # Принудительно переносим модель на CPU перед удалением if self.model is not None: if hasattr(self.model, "cpu"): self.model = self.model.cpu() del self.model self.model = None del self.model_module self.model_module = None self.model_loaded = False self.state_dict.clear() if hasattr(self.state_dict, 'clear') else None self.state_dict = {} self.config = None self.target_instrument = None self.instruments = [] self.output_arrays.clear() self.ckpt_path = None self.conf_path = None self.model_type = None gc.collect() gc.collect() extra_clear_torch_cache() nuclear_clear_model() emergency_ram_clear() self.clear_gpu_cache() def clear_gpu_cache(self): gc.collect() torch.clear_autocast_cache() if self.device.type == "mps": torch.mps.empty_cache() if self.device.type == "cuda": torch.cuda.synchronize() torch.cuda.ipc_collect() torch.cuda.empty_cache() def load_checkpoint(self, ckpt: str | Path): if not ckpt: raise PathNotSpecified(_i18n("path_not_specified")) self.ckpt_path = Path(ckpt) if not self.ckpt_path.exists(): self.ckpt_path = None raise PathNotExist(_i18n("path_not_exist")) if not self.model: self.ckpt_path = None raise ModelNotLoaded(_i18n("model_not_loaded")) if self.model_type == "mdxnet": try: self.model.init_onnx_session(self.ckpt_path, self.device, 0) self.model_loaded = True print(_i18n("checkpoint_loaded") + ": " + self.ckpt_path.name) except Exception as e: self.model_loaded = False self.ckpt_path = None self.clear_model() print(_i18n("load_checkpoint_error", error=e)) return else: try: try: self.state_dict = torch.load( self.ckpt_path, map_location=self.device, weights_only=True ) except torch.serialization.pickle.UnpicklingError: self.state_dict = torch.load( self.ckpt_path, map_location=self.device, weights_only=False ) except Exception as e: self.model_loaded = False self.ckpt_path = None self.clear_model() print(_i18n("load_checkpoint_error", error=e)) return if "state" in self.state_dict: self.state_dict = self.state_dict["state"] if "state_dict" in self.state_dict: self.state_dict = self.state_dict["state_dict"] if "model_state_dict" in self.state_dict: self.state_dict = self.state_dict["model_state_dict"] if self.model_type == "medley_vox": has_ema_keys = any(k.startswith("ema_model") for k in self.state_dict.keys()) if has_ema_keys: self.state_dict = {k: v for k, v in self.state_dict.items() if k.startswith("ema_model")} new_state_dict = {} for k, v in self.state_dict.items(): if k.startswith("ema_model.module."): new_key = k.replace("ema_model.module.", "") elif k.startswith("ema_model."): new_key = k.replace("ema_model.", "") elif k.startswith("online_model.module."): new_key = k.replace("online_model.module.", "") elif k.startswith("online_model."): new_key = k.replace("online_model.", "") elif k.startswith("module."): new_key = k.replace("module.", "") else: new_key = k if new_key not in ["initted", "step"]: new_state_dict[new_key] = v self.state_dict = new_state_dict del new_state_dict try: self.model.load_state_dict(self.state_dict) self.state_dict = {} self.model_loaded = True self.model.to(self.device) self.model.eval() print(_i18n("checkpoint_loaded") + ": " + self.ckpt_path.name) except RuntimeError as e: try: self.model.load_state_dict(self.state_dict, strict=False) self.state_dict = {} self.model_loaded = True self.model.to(self.device) self.model.eval() print(_i18n("load_state_dict_error", error=e)) print(_i18n("checkpoint_loaded") + ": " + self.ckpt_path.name) except RuntimeError as e_2: self.state_dict = {} self.model_loaded = False self.ckpt_path = None self.clear_model() print(_i18n("load_state_dict_error", error=e_2)) return def load_mix(self, path: str): self.input_file_name = None self.input_mix = None if self.config is None: raise ConfigNotLoaded(_i18n("config_is_not_loaded")) mono_bool = False if hasattr(self.config, "model"): if hasattr(self.config.model, "stereo"): mono_bool = False if self.config.model.stereo else True if not path: raise PathNotSpecified(_i18n("path_not_specified")) input_file = Path(path) if not input_file.exists(): raise PathNotExist(_i18n("path_not_exist")) if check(path): self.input_file_name = input_file.stem self.input_mix, _ = read(path=input_file, sr=self.sample_rate, mono=mono_bool) self.input_mix = self.input_mix.copy() print(_i18n("loaded_mix")+": "+input_file.name) print(_i18n("array_shape")+": "+str(self.input_mix.shape)) else: raise FileIsNotAudio(_i18n("file_is_not_audio", path=path)) def load_array(self, array: np.ndarray, orig_sr: int): self.input_file_name = "temp_array" if self.config is None: raise ConfigNotLoaded(_i18n("config_is_not_loaded")) mono_bool = False if hasattr(self.config, "model"): if hasattr(self.config.model, "stereo"): mono_bool = not self.config.model.stereo self.input_mix = easy_resampler(array.copy(), orig_sr, self.sample_rate) if orig_sr != self.sample_rate else array.copy() if mono_bool: self.input_mix = stereo_to_mono(self.input_mix) print(_i18n("loaded_mix")+": "+_i18n("from_array")) print(_i18n("array_shape")+": "+str(self.input_mix.shape)) def demix(self, add_text: str = ""): if self.input_mix is None: raise MixNotFound(_i18n("mix_not_found")) if self.input_mix.size == 0: raise MixIsEmpty(_i18n("mix_is_empty")) if not self.model_loaded: raise ModelNotLoaded(_i18n("model_not_loaded")) if self.model_type == "mdxnet": mix_tensor = torch.tensor(self.input_mix, dtype=torch.float32).to(self.device) batch_size = 1 dim_t = 256 hop_length: int = self.add_params.get("mdx_hop_length", 1024) batch_size: int = self.add_params.get("mdx_batch_size", 1) num_overlap: int = self.add_params.get("mdx_overlap", 2) denoise: bool = self.add_params.get("mdx_denoise", False) if self.add_params.get("mdx_override_segment", False): segment_size: int = self.add_params.get("mdx_segment_size", dim_t) else: segment_size: int = dim_t segment_size = round(segment_size / 128) * 128 stem_name = self.target_instrument chunk_size = hop_length * (segment_size - 1) fade_size = chunk_size // 10 step = chunk_size // num_overlap border = chunk_size - step self.model.post_init(segment_size, self.device) length_init = mix_tensor.shape[-1] if length_init > 2 * border and border > 0: wave = nn.functional.pad(mix_tensor, (border, border), mode="reflect") window = _getWindowingArray(chunk_size, fade_size).to(self.device) with torch.no_grad(): result = torch.zeros_like(wave, device=self.device) counter = torch.zeros_like(wave, device=self.device) i = 0 batch_data = [] batch_locations = [] denoise_str = " "+_i18n("denoise") if denoise else "" with tqdm(total=wave.shape[1], desc=_i18n("processing") + denoise_str + str(add_text), unit=_i18n("samples")) as progress_bar: while i < wave.shape[1]: part = wave[:, i : i + chunk_size] chunk_len = part.shape[-1] if chunk_len < chunk_size: pad_mode = "reflect" if chunk_len > chunk_size // 2 else "constant" part = nn.functional.pad( part, (0, chunk_size - chunk_len), mode=pad_mode, value=0 ) batch_data.append(part) batch_locations.append((i, chunk_len)) i += step if len(batch_data) >= batch_size or i >= wave.shape[1]: arr = torch.stack(batch_data, dim=0) for j, (start, seg_len) in enumerate(batch_locations): if denoise: processed_spec1 = self.model.forward(self.model.stft(arr[j : j + 1], chunk_size, hop_length, segment_size)) processed_spec2 = self.model.forward(self.model.stft(-(arr[j : j + 1]), chunk_size, hop_length, segment_size)) processed_wav = (self.model.istft(processed_spec1, chunk_size, hop_length, segment_size) + -self.model.istft(processed_spec2, chunk_size, hop_length, segment_size)) * 0.5 else: processed_spec = self.model.forward(self.model.stft(arr[j : j + 1], chunk_size, hop_length, segment_size)) processed_wav = self.model.istft(processed_spec, chunk_size, hop_length, segment_size) window_segment = window[..., :seg_len] result[:, start : start + seg_len] += ( processed_wav[0, :, :seg_len] * window_segment ) counter[:, start : start + seg_len] += window_segment batch_data.clear() batch_locations.clear() progress_bar.update(step) estimated_sources = result / counter if length_init > 2 * border and border > 0: estimated_sources = estimated_sources[..., border:-border] result_separation = estimated_sources.detach().cpu().numpy() result_separation = np.nan_to_num( result_separation, nan=0.0, posinf=0.0, neginf=0.0 ) self.output_arrays = {stem_name: result_separation} del mix_tensor, window, result, counter, batch_data, batch_locations del estimated_sources, result_separation if denoise: del processed_spec1, processed_spec2, processed_wav else: del processed_spec, processed_wav elif self.model_type == "vr": from models.vr_arch import spec_utils, NON_ACCOM_STEMS aggression: int = self.add_params.get("vr_aggression", 5) enable_post_process: bool = self.add_params.get("vr_post_process", False) high_end_process: bool = self.add_params.get("vr_high_end_process", False) post_process_threshold: float = self.add_params.get("vr_post_process_threshold", 0.2) batch_size: int = self.add_params.get("vr_batch_size", 1) window_size: int = self.add_params.get("vr_window_size", 512) sr = self.sample_rate model_sample_rate = self.model.model_params.param["sr"] primary_stem, secondary_stem = self.instruments[0], self.instruments[1] aggr = float(int(aggression) / 100) aggressiveness = { "value": aggr, "split_bin": self.model.model_params.param["band"][1]["crop_stop"], "aggr_correction": self.model.model_params.param.get("aggr_correction"), } input_high_end_h = None input_high_end = None X_wave, X_spec_s = {}, {} bands_n = len(self.model.model_params.param["band"]) for d in tqdm(range(bands_n, 0, -1), desc=_i18n("processing") + str(add_text), unit=_i18n("bands")): bp = self.model.model_params.param["band"][d] wav_resolution = bp["res_type"] if self.device.type == "mps": wav_resolution = "polyphase" if d == bands_n: X_wave[d], _ = librosa.resample( y=self.input_mix, orig_sr=self.sample_rate, target_sr=bp["sr"], res_type=wav_resolution, ) X_spec_s[d] = spec_utils.wave_to_spectrogram( X_wave[d], bp["hl"], bp["n_fft"], self.model.model_params, band=d, is_v51_model=self.config.model.is_vr5, ) if X_wave[d].ndim == 1: X_wave[d] = np.asarray([X_wave[d], X_wave[d]]) else: X_wave[d] = librosa.resample( X_wave[d + 1], orig_sr=self.model.model_params.param["band"][d + 1]["sr"], target_sr=bp["sr"], res_type=wav_resolution, ) X_spec_s[d] = spec_utils.wave_to_spectrogram( X_wave[d], bp["hl"], bp["n_fft"], self.model.model_params, band=d, is_v51_model=self.config.model.is_vr5, ) if d == bands_n and high_end_process: input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + ( self.model.model_params.param["pre_filter_stop"] - self.model.model_params.param["pre_filter_start"] ) input_high_end = X_spec_s[d][ :, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, : ] X_spec = spec_utils.combine_spectrograms( X_spec_s, self.model.model_params, is_v51_model=self.config.model.is_vr5 ) del X_wave, X_spec_s def spec_to_wav(spec, high_end_process, input_high_end, input_high_end_h): if ( high_end_process and isinstance(input_high_end, np.ndarray) and input_high_end_h is not None # Check if it's not None ): input_high_end_ = spec_utils.mirroring( "mirroring", spec, input_high_end, self.model.model_params ) wav = spec_utils.cmb_spectrogram_to_wave( spec, self.model.model_params, input_high_end_h, input_high_end_, is_v51_model=self.config.model.is_vr5, ) else: wav = spec_utils.cmb_spectrogram_to_wave( spec, self.model.model_params, is_v51_model=self.config.model.is_vr5 ) return wav def _execute(X_mag_pad: np.ndarray, roi_size: int) -> np.ndarray: X_dataset = [] patches = (X_mag_pad.shape[2] - 2 * self.model.offset) // roi_size for i in tqdm(range(patches), desc=_i18n("processing") + str(add_text), unit=_i18n("patches")): start = i * roi_size X_mag_window = X_mag_pad[:, :, start : start + window_size] X_dataset.append(X_mag_window) X_dataset = np.asarray(X_dataset) self.model.eval() with torch.no_grad(): mask = [] for i in tqdm(range(0, patches, batch_size), desc=_i18n("processing") + str(add_text), unit=_i18n("chunks")): X_batch = X_dataset[i : i + batch_size] X_batch = torch.from_numpy(X_batch).to(self.device) pred = self.model.predict_mask(X_batch) if not pred.size()[3] > 0: raise ValueError( _i18n("window_size_error") ) pred = pred.detach().cpu().numpy() pred = np.concatenate(pred, axis=2) mask.append(pred) if len(mask) == 0: raise ValueError( _i18n("window_size_error") ) mask = np.concatenate(mask, axis=2) return mask def postprocess( mask: np.ndarray, X_mag: np.ndarray, X_phase: np.ndarray ) -> Tuple[np.ndarray, np.ndarray]: is_non_accom_stem = False for stem in NON_ACCOM_STEMS: if stem == primary_stem.lower(): is_non_accom_stem = True mask = spec_utils.adjust_aggr(mask, is_non_accom_stem, aggressiveness) if enable_post_process: mask = spec_utils.merge_artifacts( mask, thres=post_process_threshold ) y_spec = mask * X_mag * np.exp(1.0j * X_phase) v_spec = (1 - mask) * X_mag * np.exp(1.0j * X_phase) return y_spec, v_spec X_mag, X_phase = spec_utils.preprocess(X_spec) n_frame = X_mag.shape[2] pad_l, pad_r, roi_size = spec_utils.make_padding( n_frame, window_size, self.model.offset ) X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant") X_mag_pad /= X_mag_pad.max() mask = _execute(X_mag_pad, roi_size) mask = mask[:, :, :n_frame] y_spec, v_spec = postprocess(mask, X_mag, X_phase) y_spec = np.nan_to_num(y_spec, nan=0.0, posinf=0.0, neginf=0.0) v_spec = np.nan_to_num(v_spec, nan=0.0, posinf=0.0, neginf=0.0) primary_stem_array = spec_to_wav(y_spec, high_end_process, input_high_end, input_high_end_h) primary_stem_array = librosa.resample( primary_stem_array, orig_sr=model_sample_rate, target_sr=sr, ).T secondary_stem_array = spec_to_wav(v_spec, high_end_process, input_high_end, input_high_end_h) secondary_stem_array = librosa.resample( secondary_stem_array, orig_sr=model_sample_rate, target_sr=sr, ).T self.output_arrays = { primary_stem: primary_stem_array, secondary_stem: secondary_stem_array, } del X_spec, X_mag, X_phase, X_mag_pad, mask del y_spec, v_spec, primary_stem_array, secondary_stem_array elif self.model_type == "htdemucs": mix = torch.tensor(self.input_mix, dtype=torch.float32) segment_sec: int = self.add_params.get("demucs_segment", 10) denoise: bool = self.add_params.get("demucs_denoise", False) num_overlap = self.add_params.get("demucs_overlap", 2) batch_size: int = self.add_params.get("demucs_batch_size", 1) if self.add_params.get("demucs_override_segment", False): chunk_size = self.config.training.samplerate * segment_sec else: chunk_size = getattr(self.config.training, "segment", 10) * self.config.training.samplerate num_instruments = len(self.instruments) step = chunk_size // num_overlap fade_size = chunk_size // 10 windowing_array = _getWindowingArray(chunk_size, fade_size) use_amp = getattr(self.config.training, "use_amp", True) with torch.inference_mode(): req_shape = (num_instruments,) + mix.shape result = torch.zeros(req_shape, dtype=torch.float32) counter = torch.zeros(req_shape, dtype=torch.float32) i = 0 batch_data = [] batch_locations = [] denoise_str = " "+_i18n("denoise") if denoise else "" with tqdm(total=mix.shape[1], desc=_i18n("processing") + denoise_str + str(add_text), unit=_i18n("samples")) as progress_bar: while i < mix.shape[1]: part = mix[:, i : i + chunk_size].to(self.device) chunk_len = part.shape[-1] pad_mode = "reflect" if chunk_len > chunk_size // 2 else "constant" part = nn.functional.pad( part, (0, chunk_size - chunk_len), mode=pad_mode, value=0 ) batch_data.append(part) batch_locations.append((i, chunk_len)) i += step if len(batch_data) >= batch_size or i >= mix.shape[1]: arr = torch.stack(batch_data, dim=0) if denoise: x1 = self.model(arr) x2 = self.model(-arr) x = (x1 + -x2) * 0.5 else: x = self.model(arr) window = windowing_array.clone() if i - step == 0: window[:fade_size] = 1 elif i >= mix.shape[1]: window[-fade_size:] = 1 for j, (start, seg_len) in enumerate(batch_locations): result[..., start : start + seg_len] += ( x[j, ..., :seg_len].cpu() * window[..., :seg_len] ) counter[..., start : start + seg_len] += window[..., :seg_len] batch_data.clear() batch_locations.clear() progress_bar.update(step) estimated_sources = result / counter estimated_sources = estimated_sources.detach().cpu().numpy() np.nan_to_num(estimated_sources, copy=False, nan=0.0) if num_instruments <= 1: self.output_arrays = estimated_sources else: instruments = self.instruments self.output_arrays = {k: v for k, v in zip(instruments, estimated_sources)} del mix, result, counter, batch_data, batch_locations if denoise: del estimated_sources, x, x1, x2 else: del estimated_sources, x elif self.model_type == "medley_vox": import pyloudnorm as pyln from models.medley_vox.loudness_utils import loudnorm, db2linear if self.add_params.get("mvox_override_segment", False): segment_sec: int = self.add_params.get("mvox_segment", self.config.model.seq_dur) else: segment_sec = self.config.model.seq_dur overlap: int = self.add_params.get("mvox_overlap", 2) stems: List[str] = self.instruments if self.input_mix.ndim == 1: self.input_mix = np.expand_dims(self.input_mix, axis=0) num_channels = 1 elif self.input_mix.ndim == 2: if self.input_mix.shape[0] <= self.input_mix.shape[1]: num_channels = self.input_mix.shape[0] else: self.input_mix = self.input_mix.T num_channels = self.input_mix.shape[0] samplerate = self.config.model.sample_rate chunk_size = int(samplerate * segment_sec) step = chunk_size // overlap fade_size = chunk_size // 10 meter = pyln.Meter(samplerate) try: if num_channels > 1: mix_for_loudnorm = self.input_mix.T else: mix_for_loudnorm = self.input_mix[0] mixture_d, adjusted_gain = loudnorm(mix_for_loudnorm, -24.0, meter) if num_channels > 1: if isinstance(mixture_d, np.ndarray) and mixture_d.ndim == 2: mixture_d = mixture_d.T else: mixture_d = np.tile(mixture_d, (num_channels, 1)) else: if mixture_d.ndim == 1: mixture_d = mixture_d.reshape(1, -1) except Exception as e: print(_i18n("loudnorm_error", error=str(e))) mixture_d = mix.copy() rms = np.sqrt(np.mean(mix**2)) target_rms = 0.1 if rms > 0: adjusted_gain = 20 * np.log10(target_rms / rms) mixture_d = mix * (target_rms / rms) else: adjusted_gain = 0 length_init = mixture_d.shape[1] windowing_array = _getWindowingArray(chunk_size, fade_size).to(self.device) result_stems = {stem: np.zeros((num_channels, length_init), dtype=np.float32) for stem in stems} mix_tensor = torch.tensor(mixture_d, dtype=torch.float32).to(self.device) counters = {stem: torch.zeros((num_channels, length_init), dtype=torch.float32, device=self.device) for stem in stems} i = 0 with tqdm(total=length_init, desc=_i18n("processing") + str(add_text), unit=_i18n("samples")) as progress_bar: while i < length_init: end_idx = min(i + chunk_size, length_init) chunk = mix_tensor[:, i:end_idx] cur_chunk_len = chunk.shape[1] chunk_results = torch.zeros((num_channels, 2, cur_chunk_len), dtype=torch.float32, device=self.device) for ch in range(num_channels): channel_chunk = chunk[ch:ch+1, :] if cur_chunk_len < chunk_size: pad_len = chunk_size - cur_chunk_len channel_chunk = torch.nn.functional.pad( channel_chunk, (0, pad_len), mode='constant', value=0 ) channel_chunk = channel_chunk.unsqueeze(0) with torch.no_grad(): out_chunk = self.model.separate(channel_chunk) chunk_results[ch, :, :cur_chunk_len] = out_chunk[0, :, :cur_chunk_len].cpu() window = windowing_array[:cur_chunk_len].clone() if i == 0: window[:fade_size] = 1 if end_idx >= length_init: window[-fade_size:] = 1 for stem_idx, stem in enumerate(stems): result_stems[stem][:, i:end_idx] += chunk_results[:, stem_idx, :].cpu().numpy() * window.cpu().numpy() counters[stem][:, i:end_idx] += window i += step progress_bar.update(step) for stem in stems: counters_np = counters[stem].detach().cpu().numpy() mask = counters_np > 0 result_stems[stem][mask] /= counters_np[mask] result_stems[stem] = result_stems[stem] * db2linear(-adjusted_gain) self.output_arrays = result_stems del mix_tensor, mixture_d, counters del result_stems, chunk_results del meter else: mix = torch.tensor(self.input_mix, dtype=torch.float32).to(self.device) segment: int = self.add_params.get("mdxc_segment_size", 256) if hasattr(self.config, "model"): if hasattr(self.config.model, "stft_hop_length"): hop_length = self.config.model.stft_hop_length elif hasattr(self.config.model, "hop_size"): hop_length = self.config.model.hop_size elif hasattr(self.config.model, "hop_length"): hop_length = self.config.model.hop_length if hasattr(self.config, "audio"): if hasattr(self.config.audio, "hop_length"): hop_length = self.config.audio.hop_length if hasattr(self.config, "kwargs"): if hasattr(self.config.kwargs, "hop_length"): hop_length = self.config.kwargs.hop_length if self.add_params.get("mdxc_override_segment", False): chunk_size = int(hop_length) * (int(segment) - 1) else: chunk_size = self.config.audio.chunk_size instruments = self.prefer_target_instrument() num_instruments = len(instruments) denoise: bool = self.add_params.get("mdxc_denoise", False) num_overlap: int = self.add_params.get("mdxc_overlap", 2) fade_size = chunk_size // 10 step = chunk_size // num_overlap border = chunk_size - step length_init = mix.shape[-1] windowing_array = _getWindowingArray(chunk_size, fade_size) if length_init > 2 * border and border > 0: mix = nn.functional.pad(mix, (border, border), mode="reflect") batch_size: int = self.add_params.get("mdxc_batch_size", 1) use_amp = getattr(self.config.training, "use_amp", True) with torch.inference_mode(), get_autocast_context(self.device.type, use_amp): req_shape = (num_instruments,) + mix.shape result = torch.zeros(req_shape, dtype=torch.float32) counter = torch.zeros(req_shape, dtype=torch.float32) i = 0 batch_data = [] batch_locations = [] denoise_str = " "+_i18n("denoise") if denoise else "" with tqdm(total=mix.shape[1], desc=_i18n("processing") + denoise_str + str(add_text), unit=_i18n("samples")) as progress_bar: while i < mix.shape[1]: part = mix[:, i : i + chunk_size].to(self.device) chunk_len = part.shape[-1] pad_mode = "reflect" if chunk_len > chunk_size // 2 else "constant" part = nn.functional.pad( part, (0, chunk_size - chunk_len), mode=pad_mode, value=0 ) batch_data.append(part) batch_locations.append((i, chunk_len)) i += step if len(batch_data) >= batch_size or i >= mix.shape[1]: arr = torch.stack(batch_data, dim=0) if denoise: x1 = self.model(arr) x2 = self.model(-arr) x = (x1 + -x2) * 0.5 else: x = self.model(arr) window = windowing_array.clone() if i - step == 0: window[:fade_size] = 1 elif i >= mix.shape[1]: window[-fade_size:] = 1 for j, (start, seg_len) in enumerate(batch_locations): result[..., start : start + seg_len] += ( x[j, ..., :seg_len].cpu() * window[..., :seg_len] ) counter[..., start : start + seg_len] += window[..., :seg_len] batch_data.clear() batch_locations.clear() progress_bar.update(step) estimated_sources = result / counter estimated_sources = estimated_sources.detach().cpu().numpy() np.nan_to_num(estimated_sources, copy=False, nan=0.0) if length_init > 2 * border and border > 0: estimated_sources = estimated_sources[..., border:-border] self.output_arrays = {k: v for k, v in zip(instruments, estimated_sources)} del mix, result, counter, batch_data, batch_locations if denoise: del estimated_sources, x, x1, x2 else: del estimated_sources, x self.add_second_stem() return def add_second_stem(self): if self.target_instrument: second_stem = [instrument for instrument in self.instruments if instrument != self.target_instrument][0] self.output_arrays[second_stem] = subtractor(self.input_mix, self.output_arrays[self.target_instrument], self.sample_rate, self.sample_rate, spectrogram=self.use_spec_invert)[0] print(_i18n("added_second_stem") + ": " + second_stem) else: return def delete_unselected_stems(self, selected_stems: list): if selected_stems: output_keys = list(self.output_arrays.keys()) deleted_keys = [] for stem in output_keys: if stem not in selected_stems: self.output_arrays[stem] = None del self.output_arrays[stem] deleted_keys.append(stem) print(_i18n("deleted_stems") + f": " + ",".join(deleted_keys)) else: return def extract_instrumental(self, extract_instrumental: bool, return_: bool = False): if extract_instrumental: if self.output_arrays: output_keys = [key_ for key_ in self.output_arrays] self.output_arrays["invert"] = self.input_mix.copy() for stem in output_keys: self.output_arrays["invert"] = subtractor(self.output_arrays["invert"], self.output_arrays[stem], self.sample_rate, self.sample_rate, spectrogram=self.use_spec_invert)[0] if return_: return self.output_arrays["invert"] def write(self, template: str, format_return: str = "name_stems_list"): results = [] writed_stems = [] model_name = self.ckpt_path.stem print(_i18n("format_return") + ": " + _i18n(format_return)) for stem, array in tqdm(self.output_arrays.items(), desc=_i18n("writing"), unit=_i18n('files')): custom_name = Namer.template( template, STEM=stem, MODEL=model_name, NAME=Namer.short_input_name_template(template, STEM=stem, MODEL=model_name, NAME=self.input_file_name) ) writed_stems.append([stem, write(Namer.iter(self.output_dir / f"{custom_name}.{self.output_format}"), array, self.sample_rate)]) if writed_stems: match format_return: case "name_stems_list": results = [self.input_file_name, writed_stems] case "stems_list": results = writed_stems case "stems_list_append_self": self.output_files_list.append(writed_stems) case "name_stems_list_append_self": self.output_files_list.append([self.input_file_name, writed_stems]) return results def clear_mix(self): self.input_file_name = None self.input_mix = None self.output_arrays.clear() def clear_outputs(self): self.output_files_list.clear() def get_outputs(self): return self.output_files_list def _process(self, i: int, total: int, path: str, template: str, selected_stems: list = [], extract_instrumental: bool = True): template = Namer.sanitize(template) template = Namer.dedup_template(template, keys=["NAME", "MODEL", "STEM"]) template = Namer.short(template, length=40) self.clear_mix() self.load_mix(path) try: self.demix(f" | {i}/{total} {_i18n('files')}") except Exception as e: self.clear_mix() raise DemixError(_i18n("demix_error", error=e)) from e self.delete_unselected_stems(selected_stems) self.extract_instrumental(extract_instrumental) self.write(template, "name_stems_list_append_self") self.clear_mix() def _process_array_ensemble(self, i: int, total: int, array: np.ndarray, sr: int, primary_stem: str | None = None, invert: bool = False): self.clear_mix() self.load_array(array, sr) try: self.demix(f" | {i}/{total} {_i18n('models')} | {self.ckpt_path.stem}") except Exception as e: self.clear_mix() raise DemixError(_i18n("demix_error", error=e)) from e self.delete_unselected_stems([primary_stem]) if invert: result = self.extract_instrumental(True, return_=True) else: result = self.output_arrays[primary_stem] return result, self.sample_rate def load_model(self, model_type: str, ckpt: str | Path, conf: str | Path): self.clear_model() self.load_config(model_type=model_type, conf=conf) self.load_model_instance() self.load_checkpoint(ckpt=ckpt) def inference(self, input: str | list, /, *inputs, template: str = "NAME_MDOEL_STEM", selected_stems: list = [], extract_instrumental: bool = False): self.clear_outputs() all_inputs = [] if isinstance(input, list): all_inputs.extend(input) else: all_inputs.append(input) if inputs: all_inputs.extend(inputs) total = len(all_inputs) for i, input_file in enumerate(all_inputs, start=1): try: self._process(i, total, input_file, template=template, selected_stems=selected_stems, extract_instrumental=extract_instrumental) except Exception as e: traceback.print_exc() return self.get_outputs() class ModelManager: def __init__(self): self.info = {} self.info_url = "https://huggingface.co/noblebarkrr/mvsepless_resources/resolve/main/models.json?download=true" self.info_path = Path(BASE_DIR) / "models.json" self.load_info() self.cache_dir = Path(BASE_DIR) / "separation_cache" self.cache_dir.mkdir(parents=True, exist_ok=True) # Убираем lambda-функции, заменяем на обычные методы def get_all_models(self): return [mn for mn in self.info] def get_stems(self, model_name): return [stem for stem in self.info.get(model_name, {}).get("stems", [])] def get_target_instrument(self, model_name): return self.info.get(model_name, {}).get("target_instrument", None) def get_model_type(self, model_name): return self.info.get(model_name, {}).get("model_type", "") def get_links(self, model_name): return (self.info.get(model_name, {}).get("checkpoint_url", None), self.info.get(model_name, {}).get("config_url", None)) def generate_local_paths(self, model_name): return (self.cache_dir / f"{model_name}.ckpt", self.cache_dir / f"{model_name}_config.yaml") def check_installed(self, model_name): return [path.exists() for path in self.generate_local_paths(model_name)] def check_installed2(self, model_name): return all(self.check_installed(model_name)) def load_info(self): self.info = json.loads(self.info_path.read_text("utf-8")) def show_info(self, limit: int = None, stem: str = None, only_installed: bool = False): models = [] if stem: models = [ model for model in self.get_all_models() if (stem in self.get_stems(model) or stem.lower() in self.get_stems(model) or stem.upper() in self.get_stems(model) or stem.capitalize() in self.get_stems(model) or stem.title() in self.get_stems(model)) ] else: models = self.get_all_models() if only_installed: models = [model for model in models if self.check_installed2(model)] if limit: models = models[:limit] console = rich.console.Console() table = rich.table.Table(title=_i18n("model_info"), show_lines=True) table.add_column(_i18n("model_name"), no_wrap=True) table.add_column(_i18n("output_stems")) table.add_section() table.add_row(_i18n("table_model_info_installed_legend"), _i18n("table_model_info_target_instrument_legend")) table.add_section() if models: for model_ in models: target_instrument = self.get_target_instrument(model_) stems = self.get_stems(model_) if target_instrument: for i, stem in enumerate(stems): if stem == target_instrument: stems[i] = f"[green]{stem}[/]" stems_str = ", ".join(stems) table.add_row(f"[green]{model_}[/]" if self.check_installed2(model_) else model_, stems_str) else: table.add_row(_i18n("na"), _i18n("na")) console.print(table) def update_info(self): dw_file(self.info_url, self.info_path) print(_i18n("model_info_updated")) def download(self, model_name: str): status = "" urls = self.get_links(model_name) local_paths = self.generate_local_paths(model_name) local_exists = self.check_installed(model_name) for url, local_path, exists in zip(urls, local_paths, local_exists): if not exists: dw_file(url, local_path) if all(local_exists): status = _i18n("model_already_downloaded") else: status = _i18n("model_downloaded") print(status) return status class Ensembler: def __init__(self): self.arrays = [] self.srs = [] def add_array(self, y: np.ndarray, sr: int): self.arrays.append(y) self.srs.append(sr) def get_arrays(self): return self.arrays def get_srs(self): return self.srs def clear(self): self.arrays.clear() class Separator(ModelManager): def __init__(self): super().__init__() self.mssi = MSSI() def unload_model(self): self.mssi.clear_model() @hf_spaces_gpu # (duration=120) Для спейса LongQuota / длинная квота на HuggingFace ZeroGPU (по умолчанию 60 секунд) def separate_base( self, input_valid_files: list[str | Path], model_name: str, template: str, checkpoint: str | Path, config: str | Path, selected_stems: list, extract_instrumental: bool ): self.mssi.clear_model() self.mssi.load_model(self.get_model_type(model_name), checkpoint, config) self.mssi.print_instruments() results = self.mssi.inference(input_valid_files, template=template, selected_stems=selected_stems, extract_instrumental=extract_instrumental) self.mssi.clear_model() return results def separate( self, input_files: list[str | Path], output_dir: str | Path = Path("."), output_format: str = output_formats[0], template: str = "NAME_(STEM)_MODEL", model_name: str = "bs_6stem", extract_instrumental: bool = False, use_spec_invert: bool = False, selected_stems: list = [], add_params: dict = {} ): if not output_dir: output_dir = "" input_valid_files = get_audio_files_from_list(input_files, only_files=False) if not input_valid_files: raise PathsNotSpecified(_i18n("paths_not_specified")) self.mssi.settings(output_dir=output_dir, output_format=output_format, use_spec_invert=use_spec_invert) self.mssi.set_add_params(**add_params) self.download(model_name) checkpoint, config = self.generate_local_paths(model_name) results = self.separate_base(input_valid_files, model_name, template, checkpoint, config, selected_stems, extract_instrumental) return results @hf_spaces_gpu # (duration=120) Для спейса LongQuota / длинная квота на HuggingFace ZeroGPU (по умолчанию 60 секунд) def custom_separate( self, input_files: list, output_dir: str | Path = Path("."), output_format: str = output_formats[0], template: str = "NAME_(STEM)_MODEL", model_type: str = "bs_roformer", ckpt: str = "model.ckpt", conf: str = "conf.ckpt", extract_instrumental: bool = False, use_spec_invert: bool = False, selected_stems: list = [], add_params: dict = {} ): if not output_dir: output_dir = "" input_valid_files = get_audio_files_from_list(input_files, only_files=False) if not input_valid_files: raise PathsNotSpecified(_i18n("paths_not_specified")) checkpoint, config = Path(ckpt), Path(conf) model_name = checkpoint.stem self.mssi.settings(output_dir=output_dir, output_format=output_format, use_spec_invert=use_spec_invert) self.mssi.set_add_params(**add_params) self.mssi.clear_model() self.mssi.load_model(model_type, checkpoint, config) self.previous_model_name = model_name self.mssi.print_instruments() results = self.mssi.inference(input_valid_files, template=template, selected_stems=selected_stems, extract_instrumental=extract_instrumental) self.mssi.clear_model() return results def print_flow(self, flow): """Print current ensemble flow in a formatted table (like show_info)""" if not flow: return console = rich.console.Console() table = rich.table.Table(title="", show_lines=True) table.add_column("#", style="cyan", no_wrap=True) table.add_column(_i18n("model_name")) table.add_column(_i18n("primary_stem")) table.add_column(_i18n("invert")) table.add_column(_i18n("weights"), justify="right") for idx, (model_name, primary_stem, invert, weight) in enumerate(flow, start=1): invert_str = _i18n("yes") if invert else _i18n("no") table.add_row( str(idx), model_name, primary_stem, invert_str, f"{weight:.2f}" if isinstance(weight, (int, float)) else str(weight) ) console.print(table) @hf_spaces_gpu # (duration=120) Для спейса LongQuota / длинная квота на HuggingFace ZeroGPU (по умолчанию 60 секунд) def auto_ensemble_base( self, model_name: str, checkpoint: str | Path, config: str | Path, i: int, model_count: int, mix: np.ndarray, orig_sr: int, primary_stem: str, invert: bool ): self.mssi.clear_model() self.mssi.load_model(self.get_model_type(model_name), checkpoint, config) self.mssi.print_instruments() output, model_sr = self.mssi._process_array_ensemble(i, model_count, mix, orig_sr, primary_stem, invert) self.mssi.clear_model() return output, model_sr def auto_ensemble( self, input_file: str | Path, output_dir: str | Path = Path("."), flow: list[list[str | bool | int | float]] = [], template: str = "NAME_TYPE_COUNT", etype: str = ensemble_types[0], output_format: str = output_formats[0], use_spec_invert: bool = False, save_primary_stems: bool = False ) -> tuple[str, str, list[str]]: if not output_dir: output_dir = "" if not input_file: raise PathNotSpecified(_i18n("path_not_specified")) input_file = Path(input_file) if not input_file.exists(): raise PathNotExist(_i18n("path_not_exist")) if not check(input_file): raise FileIsNotAudio(_i18n("file_is_not_audio", path=input_file)) self.print_flow(flow) if not flow: return None output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) model_count = len(flow) print(_i18n("ensemble_type")+": "+etype) print(_i18n("ensemble_models_count")+": "+str(model_count)) mix, orig_sr = read(input_file, sr=44100) template = Namer.sanitize(template) template = Namer.dedup_template(template, keys=["NAME", "TYPE", "COUNT"]) template = Namer.short(template, length=40) invert_key = "_invert" custom_name = Namer.template( template, TYPE=etype, COUNT=model_count, NAME=Namer.short_input_name_template(template, TYPE=etype, COUNT=model_count, NAME=input_file.stem) ) auto_ensembler = Ensembler() weights = [] saved_primary_stems = [] self.mssi.set_add_params(**{"demucs_denoise": True, "mdx_denoise": True}) for i, (model_name, primary_stem, invert, weight) in enumerate(flow, start=1): try: self.download(model_name) checkpoint, config = self.generate_local_paths(model_name) output, model_sr = self.auto_ensemble_base(model_name, checkpoint, config, i, model_count, mix, orig_sr, primary_stem, invert) auto_ensembler.add_array(output, model_sr) weights.append(weight) if save_primary_stems: primary_stem_file_name = primary_stem + (invert_key if invert else "") saved_primary_stems.append(write(Namer.iter(output_dir / model_name / f"{primary_stem_file_name}.flac"), output, model_sr)) except Exception as e: print(_i18n("error_occured_separation")+": "+str(e)) gr.Warning(message=""+f'{_i18n("error_occured_separation")}'.replace("\n", "
")+": "+str(e)+"", title="") continue extracted_primary_stems = auto_ensembler.get_arrays() srs = auto_ensembler.get_srs() output_array, sr_ = ensemble(extracted_primary_stems, srs, etype, weights) extracted_primary_stems = None auto_ensembler.clear() auto_ensembler, output = None, None del auto_ensembler, output inverted_array, i_sr = subtractor(mix, output_array, orig_sr, sr_, spectrogram=use_spec_invert) return write(Namer.iter(output_dir / f"{custom_name}.{output_format}"), output_array, sr_), write(Namer.iter(output_dir / f"{Namer.short(custom_name+invert_key)}.{output_format}"), inverted_array, i_sr), saved_primary_stems def manual_ensemble( self, input_files: list[str | Path], output_dir: str | Path = Path("."), weights: list[float] | None = None, template: str = "ensembled_TYPE_COUNT", etype: str = ensemble_types[0], output_format: str = output_formats[0], ) -> str: if not output_dir: output_dir = "" input_valid_files = get_audio_files_from_list(input_files, only_files=True) if not input_valid_files: raise PathsNotSpecified(_i18n("paths_not_specified")) output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) arrays, srs = multiread(input_valid_files) model_count = len(srs) results, max_sr = ensemble(arrays, srs, etype, weights) template = Namer.sanitize(template) template = Namer.dedup_template(template, keys=["TYPE", "COUNT"]) template = Namer.short(template, length=40) custom_name = Namer.template( template, TYPE=etype, COUNT=model_count ) return write(Namer.iter(output_dir / f"{custom_name}.{output_format}"), results, max_sr) def subtract(self, audio1: str | Path, audio2: str | Path, output_dir: str | Path = Path("."), output_format: str = output_formats[0], use_spec_invert: bool = False, template: str = "invert_TYPE_NAME"): if not output_dir: output_dir = "" output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) if not audio1 or not audio2: raise PathsNotSpecified(_i18n("paths_not_specified")) audio1, audio2 = Path(audio1), Path(audio2) if not audio1.exists() or not audio2.exists(): raise PathsNotExist(_i18n("paths_not_exist")) if not check(audio1) or not check(audio2): raise FilesIsNotAudio(_i18n("files_is_not_audio")) template = Namer.sanitize(template) template = Namer.dedup_template(template, keys=["NAME", "TYPE"]) template = Namer.short(template, length=40) invert_type_key = ("spectrogram" if use_spec_invert else "waveform") custom_name = Namer.template( template, TYPE=invert_type_key, NAME=Namer.short_input_name_template(template, TYPE=invert_type_key, NAME=audio1.stem) ) y1, sr1 = read(audio1) y2, sr2 = read(audio2) inverted, min_sr = subtractor(y1, y2, sr1, sr2, spectrogram=use_spec_invert) return write(Namer.iter(output_dir / f"{custom_name}.{output_format}"), inverted, min_sr) if __name__ == "__main__": separator = Separator() args = parse_separator_args(add_params_args) if args.mode == "separate": separator.separate( input_files=args.input, output_dir=args.output_dir, output_format=args.output_format, template=args.template, model_name=args.model_name, extract_instrumental=args.extract_instrumental, use_spec_invert=args.use_spec_invert, selected_stems=args.selected_stems, add_params=get_add_params(args) ) elif args.mode == "custom_separate": separator.custom_separate( input_files=args.input, output_dir=args.output_dir, output_format=args.output_format, template=args.template, model_type=args.model_type, ckpt=args.checkpoint_path, conf=args.config_path, extract_instrumental=args.extract_instrumental, use_spec_invert=args.use_spec_invert, selected_stems=args.selected_stems, add_params=get_add_params(args) ) elif args.mode == "auto_ensemble": if args.preset: flow = json.loads(Path(args.preset).read_text("utf-8")) elif args.flow: flow = [] for params in args.flow: list_values_param = params.split(":") if len(list_values_param) == 4: flow.append([str(list_values_param[0]), str(list_values_param[1]), tobool(list_values_param[2]), float(list_values_param[3])]) else: raise ValueError() separator.auto_ensemble( input_file=args.input, output_dir=args.output_dir, flow=flow, template=args.template, etype=args.ensemble_type, output_format=args.output_format, use_spec_invert=args.use_spec_invert, save_primary_stems=args.save_primary_stems ) elif args.mode == "manual_ensemble": separator.manual_ensemble( input_files=args.input, output_dir=args.output_dir, weights=args.weights, template=args.template, etype=args.ensemble_type, output_format=args.output_format ) elif args.mode == "subtract": separator.subtract( audio1=args.input_1, audio2=args.input_2, output_dir=args.output_dir, output_format=args.output_format, use_spec_invert=args.spec_invert, template=args.template ) elif args.mode == "info": if args.update: separator.update_info() elif args.download: separator.download(args.model_name) elif args.clear_cache: separator.cache_dir.unlink(missing_ok=True) else: separator.show_info(args.limit, args.stem, args.only_installed)