!pip install -U demucs
# or for local development, if you have a clone of Demucs
# pip install -e .

# You can use the `demucs` command line to separate tracks
!demucs test.mp3

# You can also load directly the pretrained models,
# for instance for the MDX 2021 winning model of Track A:
from demucs import pretrained
model = pretrained.get_model('mdx')

# Because `model` is a bag of 4 models, you cannot directly call it on your data,
# but the `apply_model` will know what to do of it.
import torch
from demucs.apply import apply_model
x = torch.randn(1, 2, 44100 * 10)  # ten seconds of white noise for the demo
out = apply_model(model, x)[0]     # shape is [S, C, T] with S the number of sources

# So let see, where is all the white noise content is going ?
for name, source in zip(model.sources, out):
    print(name, source.std() / x.std())
# The outputs are quite weird to be fair, not what I would have expected.

# now let's take a single model from the bag, and let's test it on a pure cosine
freq = 440  # in Hz
sr = model.samplerate
t = torch.arange(10 * sr).float() / sr
x = torch.cos(2 * 3.1416 * freq * t).expand(1, 2, -1)
sub_model = model.models[3]
out = sub_model(x)[0]

# Same question where does it go?
for name, source in zip(model.sources, out):
    print(name, source.std() / x.std())
    
# Well now it makes much more sense, all the energy is going
# in the `other` source.
# Feel free to try lower pitch (try 80 Hz) to see what happens !

# For training or more fun, refer to the Demucs README on our repo
# https://github.com/facebookresearch/demucs/tree/main/demucs