| import os |
| import sys |
| import warnings |
|
|
| import numpy as np |
| import torch |
| import torchaudio |
|
|
| sys.path.append(os.path.dirname(os.path.abspath(__file__))) |
|
|
| from utils.audio_utils import denormalize_audio, normalize_audio |
| from utils.model_utils import demix, load_start_checkpoint |
| from utils.settings import get_model_from_config |
|
|
| warnings.filterwarnings("ignore") |
|
|
|
|
| class Separator: |
| def __init__( |
| self, |
| config_path: str, |
| checkpoint_path: str, |
| model_type: str = "mel_band_roformer", |
| device: str = "auto", |
| ): |
| if device == "auto": |
| device = "cuda:0" if torch.cuda.is_available() else "cpu" |
| self.device = torch.device(device) |
| self.model_type = model_type |
|
|
| torch.backends.cudnn.benchmark = True |
| self.model, self.config = get_model_from_config(model_type, config_path) |
|
|
| if "model_type" in self.config.training: |
| self.model_type = self.config.training.model_type |
|
|
| from argparse import Namespace |
|
|
| fake_args = Namespace( |
| model_type=self.model_type, |
| config_path=config_path, |
| start_check_point=checkpoint_path, |
| device="auto", |
| output_dir="./output", |
| use_tta=False, |
| extract_instrumental=True, |
| pcm_type="FLOAT", |
| lora_checkpoint_loralib="", |
| draw_spectro=False, |
| ) |
| ckpt = torch.load(checkpoint_path, weights_only=False, map_location="cpu") |
| load_start_checkpoint(fake_args, self.model, ckpt, type_="inference") |
|
|
| self.model = self.model.to(self.device) |
| self.model.eval() |
| self.sample_rate = getattr(self.config.audio, "sample_rate", 44100) |
|
|
| def separate(self, wav: torch.Tensor, sr: int): |
| """ |
| Args: |
| wav: Waveform returned by torchaudio.load, shape (channels, samples) |
| sr: Sample rate |
| Returns: |
| vocal_wav: np.ndarray, shape (channels, samples) |
| inst_wav: np.ndarray, shape (channels, samples) |
| sr: int, output sample rate |
| """ |
| |
| if sr != self.sample_rate: |
| wav = torchaudio.transforms.Resample(sr, self.sample_rate)(wav) |
| sr = self.sample_rate |
|
|
| mix = wav.numpy() |
|
|
| |
| if mix.shape[0] == 1 and getattr(self.config.audio, "num_channels", 1) == 2: |
| mix = np.concatenate([mix, mix], axis=0) |
|
|
| mix_orig = mix.copy() |
|
|
| |
| norm_params = None |
| if getattr(self.config.inference, "normalize", False): |
| mix, norm_params = normalize_audio(mix) |
|
|
| |
| waveforms = demix( |
| self.config, |
| self.model, |
| mix, |
| self.device, |
| model_type=self.model_type, |
| pbar=True, |
| ) |
|
|
| |
| vocal_wav = waveforms.get("vocals", list(waveforms.values())[0]) |
| if norm_params is not None: |
| vocal_wav = denormalize_audio(vocal_wav, norm_params) |
|
|
| |
| inst_wav = mix_orig - vocal_wav |
|
|
| return vocal_wav, inst_wav, sr |
|
|
|
|
| |
| if __name__ == "__main__": |
| sep = Separator( |
| config_path="ckpts/config_vocals_mel_band_roformer_kj.yaml", |
| checkpoint_path="ckpts/MelBandRoformer.ckpt", |
| device="cuda:0", |
| ) |
|
|
| wav, sr = torchaudio.load("path/to/input.mp3") |
| vocal_wav, inst_wav, sr = sep.separate(wav, sr) |
|
|
| torchaudio.save("output_vocals.wav", torch.from_numpy(vocal_wav), sr) |
| torchaudio.save("output_instrumental.wav", torch.from_numpy(inst_wav), sr) |
|
|
