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)

Or the Log-Depth Recurrent Modeling (ARGRC) architecture:

import torch
from hs_tasnet import HSTasNet
from hs_tasnet.alternative_rnns import ARGRCWrapper

model = HSTasNet(rnn_klass = ARGRCWrapper)

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}
}

@misc{wang2026logdepthrecurrentlanguagemodeling,
    title    = {Log-Depth Recurrent Language Modeling},
    author   = {Yiqin Wang and Nuri Cingillioglu and Charles Pert},
    year     = {2026},
    eprint   = {2609.28212},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url      = {https://arxiv.org/abs/2609.28212},
}

Metadata

Release files for HS-TasNet 0.3.4

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.4
File	Size	Uploaded
hs_tasnet-0.3.4.tar.gz	20.2 kB	Oct 13, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for HS-TasNet 0.3.4
File	Interpreter	ABI	Platform	Reset
hs_tasnet-0.3.4-py3-none-any.whl 20.7 kB Oct 13, 2026	Python 3	none	any	Details

Total release size: 40.9 kB

Release files / hs_tasnet-0.3.4.tar.gz

Download URL	hs_tasnet-0.3.4.tar.gz
Size	20.2 kB
Tags	Source
SHA-256 checksum How to use checksums	`bc9a9df3652d443aaa5bced132d3ad6c1159946bdd152434837bd13e01099922`
BLAKE2b-256 checksum How to use checksums	`a9521c740f9844a0d81dea6f053b4ddb6329124737b26aa5aaa95f1458700218`
Upload date	Oct 13, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

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

Download URL	hs_tasnet-0.3.4-py3-none-any.whl
Size	20.7 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`2adff5de88563a2e7aef5a519304037236cd61a11d496d419453230a6bded9c1`
BLAKE2b-256 checksum How to use checksums	`4ba69723ccdab1cfc5cd21f6ef1b81ee1fecf39c1ed3ad22d39f90add2d1d2ed`
Upload date	Oct 13, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`