|
|
import torchaudio |
|
|
import os |
|
|
import torch |
|
|
from third_party.demucs.models.pretrained import get_model_from_yaml |
|
|
|
|
|
|
|
|
class Separator(torch.nn.Module): |
|
|
def __init__(self, dm_model_path='third_party/demucs/ckpt/htdemucs.pth', dm_config_path='third_party/demucs/ckpt/htdemucs.yaml', gpu_id=0) -> None: |
|
|
super().__init__() |
|
|
if torch.cuda.is_available() and gpu_id < torch.cuda.device_count(): |
|
|
self.device = torch.device(f"cuda:{gpu_id}") |
|
|
else: |
|
|
self.device = torch.device("cpu") |
|
|
self.demucs_model = self.init_demucs_model(dm_model_path, dm_config_path) |
|
|
|
|
|
def init_demucs_model(self, model_path, config_path): |
|
|
model = get_model_from_yaml(config_path, model_path) |
|
|
model.to(self.device) |
|
|
model.eval() |
|
|
return model |
|
|
|
|
|
def load_audio(self, f): |
|
|
a, fs = torchaudio.load(f) |
|
|
if (fs != 48000): |
|
|
a = torchaudio.functional.resample(a, fs, 48000) |
|
|
if a.shape[-1] >= 48000*10: |
|
|
a = a[..., :48000*10] |
|
|
else: |
|
|
a = torch.cat([a, a], -1) |
|
|
return a[:, 0:48000*10] |
|
|
|
|
|
def run(self, audio_path, output_dir='tmp', ext=".flac"): |
|
|
os.makedirs(output_dir, exist_ok=True) |
|
|
name, _ = os.path.splitext(os.path.split(audio_path)[-1]) |
|
|
output_paths = [] |
|
|
|
|
|
for stem in self.demucs_model.sources: |
|
|
output_path = os.path.join(output_dir, f"{name}_{stem}{ext}") |
|
|
if os.path.exists(output_path): |
|
|
output_paths.append(output_path) |
|
|
if len(output_paths) == 1: |
|
|
vocal_path = output_paths[0] |
|
|
else: |
|
|
drums_path, bass_path, other_path, vocal_path = self.demucs_model.separate(audio_path, output_dir, device=self.device) |
|
|
for path in [drums_path, bass_path, other_path]: |
|
|
os.remove(path) |
|
|
full_audio = self.load_audio(audio_path) |
|
|
vocal_audio = self.load_audio(vocal_path) |
|
|
bgm_audio = full_audio - vocal_audio |
|
|
return full_audio, vocal_audio, bgm_audio |
|
|
|
