HS-TasNet

Implementation of HS-TasNet, "Real-time Low-latency Music Source Separation using Hybrid Spectrogram-TasNet", proposed by the research team at L-Acoustics

Install

$ pip install HS-TasNet

Usage

import torch
from hs_tasnet import HSTasNet

model = HSTasNet()

audio = torch.randn(1, 2, 204800) # ~5 seconds of stereo

separated_audios, _ = model(audio)

assert separated_audios.shape == (1, 4, 2, 204800) # second dimension is the separated tracks

With the Trainer

# model

from hs_tasnet import HSTasNet, Trainer

model = HSTasNet()

# trainer

trainer = Trainer(
    model,
    dataset = None,               # add your in-house Dataset
    concat_musdb_dataset = True,  # concat the musdb dataset automatically
    batch_size = 2,
    max_steps = 2,
    cpu = True,
)

trainer()

# after much training
# inferencing

model.sounddevice_stream(
    duration_seconds = 2,
    return_reduced_sources = [0, 2]
)

# or from the exponentially smoothed model (in the trainer)

trainer.ema_model.sounddevice_stream(...)

# or you can load from a specific checkpoint

model.load('./checkpoints/path.to.desired.ckpt.pt')
model.sounddevice_stream(...)

# to load an HS-TasNet from any of the saved checkpoints, without having to save its hyperparameters, just run

model = HSTasNet.init_and_load_from('./checkpoints/path.to.desired.ckpt.pt')

Training script

First make sure dependencies are there by running

$ sh scripts/install.sh

Then make sure uv is installed

$ pip install uv

Finally run the following to train a newly initialized model on a small subset of MusDB, and make sure the loss goes down

$ uv run train.py

For distributed training, you just need to run accelerate config first, courtesy of accelerate from 🤗 but single machine is fine too

Experiment tracking

To enable online experiment monitoring / tracking, you need to have wandb installed and logged in

$ pip install wandb && wandb login

Then

$ uv run train.py --use-wandb

To wipe the previous checkpoints and evaluated results, append --clear-folders

Alternative RNNs

The architecture defaults to using PyTorch's LSTM (or GRU), but you can easily substitute it for any other module by passing an rnn_klass to the HSTasNet constructor, as long as it adheres to a specific interface (read alternative_rnns.py)

For example, to use the minGRU architecture:

import torch
from hs_tasnet import HSTasNet
from hs_tasnet.alternative_rnns import minGRUWrapper

model = HSTasNet(rnn_klass = minGRUWrapper)

audio = torch.randn(1, 2, 204800)
separated_audios, _ = model(audio)

Test

$ uv pip install '.[test]' --system

Then

$ pytest tests

Citations

@misc{venkatesh2024realtimelowlatencymusicsource,
    title    = {Real-time Low-latency Music Source Separation using Hybrid Spectrogram-TasNet},
    author   = {Satvik Venkatesh and Arthur Benilov and Philip Coleman and Frederic Roskam},
    year     = {2024},
    eprint   = {2402.17701},
    archivePrefix = {arXiv},
    primaryClass = {eess.AS},
    url      = {https://arxiv.org/abs/2402.17701},
}

@inproceedings{Feng2024WereRA,
    title   = {Were RNNs All We Needed?},
    author  = {Leo Feng and Frederick Tung and Mohamed Osama Ahmed and Yoshua Bengio and Hossein Hajimirsadegh},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273025630}
}

Metadata

Release files for HS-TasNet 0.3.3

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for HS-TasNet 0.3.3
File	Size	Uploaded
hs_tasnet-0.3.3.tar.gz	19.7 kB	Aug 25, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for HS-TasNet 0.3.3
File	Interpreter	ABI	Platform	Reset
hs_tasnet-0.3.3-py3-none-any.whl 20.1 kB Aug 25, 2026	Python 3	none	any	Details

Total release size: 39.8 kB

Release files / hs_tasnet-0.3.3.tar.gz

Download URL	hs_tasnet-0.3.3.tar.gz
Size	19.7 kB
Tags	Source
SHA-256 checksum How to use checksums	`d59d6010031e810ac09dd2b21f1e10486de7ddc77cbdab0b20e5fc6028ab898c`
BLAKE2b-256 checksum How to use checksums	`46459469a25f452c97cceefe0e4a9c110c961dc4fdde9bbd0e281e6e8cbcaf7e`
Upload date	Aug 25, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / hs_tasnet-0.3.3-py3-none-any.whl

Download URL	hs_tasnet-0.3.3-py3-none-any.whl
Size	20.1 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`41f3706fd07ebc638e86e4400e0c062b21629336c1d7f4031321e83073dc4bbd`
BLAKE2b-256 checksum How to use checksums	`9010995a57d1c67aa5211e41053a73927a41ecbd7badedb4c600f0943dbb7268`
Upload date	Aug 25, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`