openunmix_pytorch_hub_music_stem_separation.py

1import torch
2import torchaudio
3from openunmix import utils
4
5# 1. Load a pre-trained model (e.g., 'umxhq') from PyTorch Hub
6# This will automatically download the model weights on the first run
7separator = torch.hub.load('sigsep/open-unmix-pytorch', 'umxhq')
8
9# 2. Load an audio file
10# Replace 'path/to/audio.wav' with your actual file path
11audio_path = 'path/to/audio.wav'
12audio, sample_rate = torchaudio.load(audio_path)
13
14# 3. Perform separation
15# The model expects input shape [batch, channels, samples]
16# 'estimates' is a dictionary containing the separated tensors for 
17# 'vocals', 'drums', 'bass', and 'other'
18estimates = separator(audio)
19
20# 4. Access a specific stem (e.g., vocals)
21vocals = estimates['vocals']
22
23# Optional: Save the output
24torchaudio.save('vocals.wav', vocals[0], sample_rate)