| import torch, torchaudio |
| from diffsynth import load_model, ModelConfig |
| from diffsynth.models.demucs import HTDemucs |
|
|
| class AudioTrackSeparator(torch.nn.Module): |
| def __init__(self, torch_dtype=torch.float32, device="cuda", model_config=ModelConfig(model_id="DiffSynth-Studio/Demucs-Repackage", origin_file_pattern="model.safetensors")): |
| super().__init__() |
| model_config.download_if_necessary() |
| self.model = load_model(HTDemucs, model_config.path, torch_dtype=torch_dtype, device=device) |
|
|
| @torch.no_grad() |
| def __call__(self, audio, target_sample_rate=48000, **kwargs): |
| if isinstance(audio, str): |
| audio, sample_rate = torchaudio.load(audio) |
| else: |
| audio, sample_rate = audio |
| audio = audio.to(dtype=next(iter(self.model.parameters())).dtype, device=next(iter(self.model.parameters())).device) |
| vocals = self.model.extract_track(audio, sample_rate) |
| if target_sample_rate != 44100: |
| vocals = torchaudio.functional.resample(vocals, 44100, target_sample_rate) |
| return vocals |
|
|
