import os import gc import sys import torch import librosa import numpy as np import soundfile as sf from pydub import AudioSegment sys.path.append(os.getcwd()) from main.library import opencl from main.library.uvr5_lib.spec_utils import normalize class CommonSeparator: VOCAL_STEM = "Vocals" INST_STEM = "Instrumental" OTHER_STEM = "Other" BASS_STEM = "Bass" DRUM_STEM = "Drums" GUITAR_STEM = "Guitar" PIANO_STEM = "Piano" PRIMARY_STEM = "Primary Stem" SECONDARY_STEM = "Secondary Stem" LEAD_VOCAL_STEM = "lead_only" BV_VOCAL_STEM = "backing_only" NO_STEM = "No " STEM_PAIR_MAPPER = {VOCAL_STEM: INST_STEM, INST_STEM: VOCAL_STEM, LEAD_VOCAL_STEM: BV_VOCAL_STEM, BV_VOCAL_STEM: LEAD_VOCAL_STEM, PRIMARY_STEM: SECONDARY_STEM} def __init__(self, config): self.logger = config.get("logger") self.torch_device = config.get("torch_device") self.torch_device_cpu = config.get("torch_device_cpu") self.torch_device_mps = config.get("torch_device_mps") self.onnx_execution_provider = config.get("onnx_execution_provider") self.model_name = config.get("model_name") self.model_path = config.get("model_path") self.model_data = config.get("model_data") self.output_dir = config.get("output_dir") self.output_format = config.get("output_format") self.output_bitrate = config.get("output_bitrate") self.normalization_threshold = config.get("normalization_threshold") self.enable_denoise = config.get("enable_denoise") self.output_single_stem = config.get("output_single_stem") self.invert_using_spec = config.get("invert_using_spec") self.sample_rate = config.get("sample_rate") self.primary_stem_name = None self.secondary_stem_name = None if "training" in self.model_data and "instruments" in self.model_data["training"]: instruments = self.model_data["training"]["instruments"] if instruments: self.primary_stem_name = instruments[0] self.secondary_stem_name = instruments[1] if len(instruments) > 1 else self.secondary_stem(self.primary_stem_name) if self.primary_stem_name is None: self.primary_stem_name = self.model_data.get("primary_stem", "Vocals") self.secondary_stem_name = self.secondary_stem(self.primary_stem_name) self.is_karaoke = self.model_data.get("is_karaoke", False) self.is_bv_model = self.model_data.get("is_bv_model", False) self.bv_model_rebalance = self.model_data.get("is_bv_model_rebalanced", 0) self.audio_file_path = None self.audio_file_base = None self.primary_source = None self.secondary_source = None self.primary_stem_output_path = None self.secondary_stem_output_path = None self.cached_sources_map = {} def secondary_stem(self, primary_stem): primary_stem = primary_stem if primary_stem else self.NO_STEM return self.STEM_PAIR_MAPPER[primary_stem] if primary_stem in self.STEM_PAIR_MAPPER else primary_stem.replace(self.NO_STEM, "") if self.NO_STEM in primary_stem else f"{self.NO_STEM}{primary_stem}" def separate(self, audio_file_path): pass def final_process(self, stem_path, source, stem_name): self.write_audio(stem_path, source) return {stem_name: source} def cached_sources_clear(self): self.cached_sources_map = {} def cached_source_callback(self, model_architecture, model_name=None): model, sources = None, None mapper = self.cached_sources_map[model_architecture] for key, value in mapper.items(): if model_name in key: model = key sources = value return model, sources def cached_model_source_holder(self, model_architecture, sources, model_name=None): self.cached_sources_map[model_architecture] = {**self.cached_sources_map.get(model_architecture, {}), **{model_name: sources}} def prepare_mix(self, mix): if not isinstance(mix, np.ndarray): mix, _ = librosa.load(mix, mono=False, sr=self.sample_rate) else: mix = mix.T if mix.ndim == 1: mix = np.asfortranarray([mix, mix]) return mix def write_audio(self, stem_path, stem_source): duration_seconds = librosa.get_duration(y=librosa.load(self.audio_file_path, sr=None)[0]) duration_hours = duration_seconds / 3600 if duration_hours >= 1: self.write_audio_soundfile(stem_path, stem_source) else: self.write_audio_pydub(stem_path, stem_source) def write_audio_pydub(self, stem_path, stem_source): stem_source = normalize(wave=stem_source, max_peak=self.normalization_threshold) if np.max(np.abs(stem_source)) < 1e-6: return if self.output_dir: os.makedirs(self.output_dir, exist_ok=True) stem_path = os.path.join(self.output_dir, stem_path) if stem_source.dtype != np.int16: stem_source = (stem_source * 32767).astype(np.int16) stem_source_interleaved = np.empty((2 * stem_source.shape[0],), dtype=np.int16) stem_source_interleaved[0::2] = stem_source[:, 0] stem_source_interleaved[1::2] = stem_source[:, 1] audio_segment = AudioSegment(stem_source_interleaved.tobytes(), frame_rate=self.sample_rate, sample_width=stem_source.dtype.itemsize, channels=2) file_format = stem_path.lower().split(".")[-1] if file_format == "m4a": file_format = "mp4" elif file_format == "mka": file_format = "matroska" audio_segment.export(stem_path, format=file_format, bitrate="320k" if file_format == "mp3" and self.output_bitrate is None else self.output_bitrate) def write_audio_soundfile(self, stem_path, stem_source): if stem_source.shape[1] == 2: if stem_source.flags["F_CONTIGUOUS"]: stem_source = np.ascontiguousarray(stem_source) else: stereo_interleaved = np.empty((2 * stem_source.shape[0],), dtype=np.int16) stereo_interleaved[0::2] = stem_source[:, 0] stereo_interleaved[1::2] = stem_source[:, 1] stem_source = stereo_interleaved sf.write(stem_path, stem_source, self.sample_rate) def clear_gpu_cache(self): gc.collect() if self.torch_device == torch.device("mps"): torch.mps.empty_cache() elif self.torch_device == torch.device("cuda"): torch.cuda.empty_cache() elif opencl.torch_available and self.torch_device == torch.device("ocl"): opencl.pytorch_ocl.empty_cache() def clear_file_specific_paths(self): self.audio_file_path = None self.audio_file_base = None self.primary_source = None self.secondary_source = None self.primary_stem_output_path = None self.secondary_stem_output_path = None