openunmix_torchhub_music_source_separation_quickstart.py

1import torch
2import torchaudio
3import stempeg
4
5# 1. Load the model from TorchHub
6# Options for targets: 'vocals', 'drums', 'bass', 'other'
7separator = torch.hub.load('sigsep/open-unmix-pytorch', 'umxhq', device='cpu')
8
9# 2. Load audio file
10# Replace 'path/to/audio.wav' with your actual file path
11audio, sample_rate = torchaudio.load('path/to/audio.wav')
12
13# 3. Perform separation
14# The model expects input shape (nb_samples, nb_channels, nb_timesteps)
15# We add a batch dimension using [None]
16estimates = separator(audio[None, ...])
17
18# 4. Access the results
19# estimates is a tensor of shape (nb_samples, nb_targets, nb_channels, nb_timesteps)
20vocals = estimates[0, 0]
21drums = estimates[0, 1]
22bass = estimates[0, 2]
23other = estimates[0, 3]
24
25# 5. Save the results (example: vocals)
26torchaudio.save('vocals.wav', vocals, sample_rate)