import os import logging import librosa import numpy as np import soundfile as sf import torch from stqdm import stqdm import streamlit as st from pydub import AudioSegment from app.service.vocal_remover import nets if os.environ.get("LIMIT_CPU", False): torch.set_num_threads(1) def merge_artifacts(y_mask, thres=0.05, min_range=64, fade_size=32): if min_range < fade_size * 2: raise ValueError("min_range must be >= fade_size * 2") idx = np.where(y_mask.min(axis=(0, 1)) > thres)[0] start_idx = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0]) end_idx = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1]) artifact_idx = np.where(end_idx - start_idx > min_range)[0] weight = np.zeros_like(y_mask) if len(artifact_idx) > 0: start_idx = start_idx[artifact_idx] end_idx = end_idx[artifact_idx] old_e = None for s, e in zip(start_idx, end_idx): if old_e is not None and s - old_e < fade_size: s = old_e - fade_size * 2 if s != 0: weight[:, :, s : s + fade_size] = np.linspace(0, 1, fade_size) else: s -= fade_size if e != y_mask.shape[2]: weight[:, :, e - fade_size : e] = np.linspace(1, 0, fade_size) else: e += fade_size weight[:, :, s + fade_size : e - fade_size] = 1 old_e = e v_mask = 1 - y_mask y_mask += weight * v_mask return y_mask def make_padding(width, cropsize, offset): left = offset roi_size = cropsize - offset * 2 if roi_size == 0: roi_size = cropsize right = roi_size - (width % roi_size) + left return left, right, roi_size def wave_to_spectrogram(wave, hop_length, n_fft): wave_left = np.asfortranarray(wave[0]) wave_right = np.asfortranarray(wave[1]) spec_left = librosa.stft(wave_left, n_fft=n_fft, hop_length=hop_length) spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length) spec = np.asfortranarray([spec_left, spec_right]) return spec def spectrogram_to_wave(spec, hop_length=1024): if spec.ndim == 2: wave = librosa.istft(spec, hop_length=hop_length) elif spec.ndim == 3: spec_left = np.asfortranarray(spec[0]) spec_right = np.asfortranarray(spec[1]) wave_left = librosa.istft(spec_left, hop_length=hop_length) wave_right = librosa.istft(spec_right, hop_length=hop_length) wave = np.asfortranarray([wave_left, wave_right]) return wave class Separator(object): def __init__(self, model, device, batchsize, cropsize, postprocess=False, progress_bar=None): self.model = model self.offset = model.offset self.device = device self.batchsize = batchsize self.cropsize = cropsize self.postprocess = postprocess self.progress_bar = progress_bar def _separate(self, X_mag_pad, roi_size): X_dataset = [] patches = (X_mag_pad.shape[2] - 2 * self.offset) // roi_size for i in range(patches): start = i * roi_size X_mag_crop = X_mag_pad[:, :, start : start + self.cropsize] X_dataset.append(X_mag_crop) X_dataset = np.asarray(X_dataset) self.model.eval() with torch.no_grad(): mask = [] # To reduce the overhead, dataloader is not used. for i in stqdm( range(0, patches, self.batchsize), st_container=self.progress_bar, gui=False, ): X_batch = X_dataset[i : i + self.batchsize] X_batch = torch.from_numpy(X_batch).to(self.device) pred = self.model.predict_mask(X_batch) pred = pred.detach().cpu().numpy() pred = np.concatenate(pred, axis=2) mask.append(pred) mask = np.concatenate(mask, axis=2) return mask def _preprocess(self, X_spec): X_mag = np.abs(X_spec) X_phase = np.angle(X_spec) return X_mag, X_phase def _postprocess(self, mask, X_mag, X_phase): if self.postprocess: mask = merge_artifacts(mask) 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 def separate(self, X_spec): X_mag, X_phase = self._preprocess(X_spec) n_frame = X_mag.shape[2] pad_l, pad_r, roi_size = make_padding(n_frame, self.cropsize, self.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 = self._separate(X_mag_pad, roi_size) mask = mask[:, :, :n_frame] y_spec, v_spec = self._postprocess(mask, X_mag, X_phase) return y_spec, v_spec @st.cache_resource(show_spinner=False) def load_model(pretrained_model, n_fft=2048): model = nets.CascadedNet(n_fft, 32, 128) if torch.cuda.is_available(): device = torch.device("cuda:0") model.to(device) # elif torch.backends.mps.is_available() and torch.backends.mps.is_built(): # device = torch.device("mps") # model.to(device) else: device = torch.device("cpu") model.load_state_dict(torch.load(pretrained_model, map_location=device)) return model, device # @st.cache_data(show_spinner=False) def separate( input, model, device, output_dir, batchsize=4, cropsize=256, postprocess=False, hop_length=1024, n_fft=2048, sr=44100, progress_bar=None, only_no_vocals=False, ): X, sr = librosa.load(input, sr=sr, mono=False, dtype=np.float32, res_type="kaiser_fast") basename = os.path.splitext(os.path.basename(input))[0] if X.ndim == 1: # mono to stereo X = np.asarray([X, X]) X_spec = wave_to_spectrogram(X, hop_length, n_fft) with torch.no_grad(): sp = Separator(model, device, batchsize, cropsize, postprocess, progress_bar=progress_bar) y_spec, v_spec = sp.separate(X_spec) base_dir = f"{output_dir}/vocal_remover/{basename}" os.makedirs(base_dir, exist_ok=True) wave = spectrogram_to_wave(y_spec, hop_length=hop_length) try: sf.write(f"{base_dir}/no_vocals.mp3", wave.T, sr) except Exception: logging.error("Failed to write no_vocals.mp3, trying pydub...") pydub_write(wave, f"{base_dir}/no_vocals.mp3", sr) if only_no_vocals: return wave = spectrogram_to_wave(v_spec, hop_length=hop_length) try: sf.write(f"{base_dir}/vocals.mp3", wave.T, sr) except Exception: logging.error("Failed to write vocals.mp3, trying pydub...") pydub_write(wave, f"{base_dir}/vocals.mp3", sr) def pydub_write(wave, output_path, frame_rate, audio_format="mp3"): # Ensure the wave data is in the right format for pydub (mono and 16-bit depth) wave_16bit = (wave * 32767).astype(np.int16) audio_segment = AudioSegment( wave_16bit.tobytes(), frame_rate=frame_rate, sample_width=wave_16bit.dtype.itemsize, channels=1, ) audio_segment.export(output_path, format=audio_format)