| | 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 |
| |
|
