Skip to content

feat: add spleeter_pytorch — TF-free 4/5-stem inference for Python 3.14#954

Open
salvatoretrimarchi wants to merge 1 commit into
deezer:masterfrom
salvatoretrimarchi:master
Open

feat: add spleeter_pytorch — TF-free 4/5-stem inference for Python 3.14#954
salvatoretrimarchi wants to merge 1 commit into
deezer:masterfrom
salvatoretrimarchi:master

Conversation

@salvatoretrimarchi

Copy link
Copy Markdown

PyTorch re-implementation of the Spleeter U-Net separation pipeline, replacing TensorFlow (unsupported on Python 3.14) with pure PyTorch inference and a hand-written TF1 SSTable checkpoint reader.

Modules added:

  • unet.py — U-Net architecture ported from Deezer TF model
    (ELU activations for 5-stems config)
  • checkpoint_reader.py — Pure-Python LevelDB SSTable parser for reading
    TF1 .index + .data checkpoint files without TF
  • convert_weights.py — Converts Spleeter 4-stem or 5-stem TF1 checkpoint
    weights → PyTorch .pt files (one per instrument)
  • separator.py — Complete separation pipeline: STFT (librosa) →
    UNet per instrument → Wiener masking → iSTFT
  • init.py — Package entry point

Features:

  • Supports 4-stem (vocals/drums/bass/other) and 5-stem (+ piano) configs
  • Auto-detects available .pt files; loads only those present
  • GPU-accelerated when torch.cuda.is_available()
  • No TensorFlow, no Docker, no Python version constraints beyond PyTorch

Pull request title

  • I read contributing guideline
  • I didn't find a similar pull request already open.
  • My PR is related to Spleeter only, not a derivative product (such as Webapplication, or GUI provided by others)

Description

A few sentences describing the overall goals of the pull request's commits.

How this patch was tested

You tested it, right?

  • I implemented unit test whicn ran successfully using poetry run pytest tests/
  • Code has been formatted using poetry run black spleeter
  • Imports has been formatted using `poetry run isort spleeter``

Documentation link and external references

Please provide any info that may help us better understand your code.

PyTorch re-implementation of the Spleeter U-Net separation pipeline,
replacing TensorFlow (unsupported on Python 3.14) with pure PyTorch
inference and a hand-written TF1 SSTable checkpoint reader.

Modules added:
- unet.py              — U-Net architecture ported from Deezer TF model
                         (ELU activations for 5-stems config)
- checkpoint_reader.py — Pure-Python LevelDB SSTable parser for reading
                         TF1 .index + .data checkpoint files without TF
- convert_weights.py   — Converts Spleeter 4-stem or 5-stem TF1 checkpoint
                         weights → PyTorch .pt files (one per instrument)
- separator.py         — Complete separation pipeline: STFT (librosa) →
                         UNet per instrument → Wiener masking → iSTFT
- __init__.py          — Package entry point

Features:
- Supports 4-stem (vocals/drums/bass/other) and 5-stem (+ piano) configs
- Auto-detects available .pt files; loads only those present
- GPU-accelerated when torch.cuda.is_available()
- No TensorFlow, no Docker, no Python version constraints beyond PyTorch

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant