AMD GPU support for stem separation

[judilsteve]

Currently, the project appears to use two different python packages for stem separation:

When the user selects demucs, it calls out to https://pypi.org/project/demucs/

When the user selects UVR (the default), it calls out to https://pypi.org/project/audio-separator/

Both paths could potentially be consolidated to use https://pypi.org/project/uvr-headless-runner/. uvr-headless-runner is set up as a command-line tool, but also provides a convenient python API.

Aside from code consolidation, this package has the benefit of supporting GPU acceleration on AMD GPUs via DirectML (when running on Windows).

[rzru]

hi @judilsteve, excellent feedback, thanks! I'll take a look whether it would be straightforward to swap the packages.

[judilsteve]

btw, this kicks ass

Only thing I can suggest feature-wise is a "simple" UI where I can just drag/drop individual files to have them processed, or select them from an embedded directory tree.

The "point it at your library directory and let it scan everything" doesn't really work for me, as my music library lives on a network share and contains ~15k songs, so it takes forever to enumerate them all and scrape the tags. If I waited for it to complete, I'd end up with an absolutely enormous scrolling list of uncategorised songs sorted alphabetically by title, which isn't particularly useful.

[rzru]

@judilsteve noted as well, thanks for the suggestion! indeed there would be a lot of enhancements to the whole directory management process

[judilsteve]

I had a crack at this, but uvr-headless-runner has a lot of hard-pinned deps. I forked it and tried to relax some of the constraints here:

https://github.com/judilsteve/uvr-headless-runner#

but it just turned into a hideous yak-shaving exercise. Got to a point where audio-separator>=0.25.0 was trying to pull rotary-embedding-torch>=0.6.1,<0.7.0 but uvr-headless-runner wants rotary-embedding-torch>=0.8.0,<0.9 and have given up for now

[judilsteve]

Ah, I should probably just remove audio-separator since uvr-headless-runner would replace it.

Was just giving this a quick look on my lunch break from work, might have a proper go at it on the weekend if I decide I'm up for some self-inflicted suffering.

[judilsteve]

Had a bit more of a go at this, not sure it's still worth implementing at this point. First off, DirectML is supposedly abandoned and has been superseded by something called "WindowsML", but WindowsML only supports ONNX runtime and everything I've seen so far seems to run on pytorch. I'm not sure if models are generally portable across the different inference runtimes/engines/whatever they're called.

DirectML still works on my machine for now (I'm able to use the DirectML build of UVR GUI, I can see it utilising my GPU, and it runs much faster than the CUDA/CPU build), but the torch-directml package that makes this possible is abandoned, has no source code available that I can see, and only supports up to python 3.10 (which is very close to EoL).

Finally, uvr-headless-runner doesn't seem to be packaged properly. After installing the package, none of the entrypoints are importable inside the venv. It seems you were expected to use the code directly (e.g. python -m entrypoint.py) rather than downloading it as a package from pypi. So I'd have to fix that up in my fork, and I just don't have the willpower for that right now.

For the broader issue of these packages constantly fighting over which version of pytorch/scipy/whatever to install, one way to go might be to hoist the code that orchestrates the pipeline of separate vocals from instrumentals -> align/transcribe lyrics up into the rust part of the project. Then you could have separate venvs for the two steps, which should hopefully avoid a lot of dependency hell and allow you to keep whisperx up to date even if the other packages demanded ancient python versions and pypi packages, or vice-versa.

The more complicated solution would be to stop relying on packages that wrap these models up into neat little functions that take wav files as input and write wav files for output, and instead wrangle with the models and pytorch/what-have-you directly. This is a lot of code that I don't know how to write, but it would mean you weren't beholden to anyone else's poor package-management practices, and might also allow you to massively improve performance by pipelining things (as soon as the splitter model has processed a chunk of audio, you can feed that chunk into the transcription/alignment model, instead of waiting for the entire splitting process to finish before transcription/alignment can begin).

[rzru]

@judilsteve that's a lot of exploration, wow, thanks! I also don't have enough expertise to dive deeper into this, so I think it can remain parked for now, but that's very nice of you to dive deep into this and give me so much insight!

[judilsteve]

Turns out there is audio-separator[dml] which uses directml, although it will by default resolve an old version of torch-directml that didn't work on my machine, so I had to manually install the latest release (which has py3.12 wheels despite the metadata only listing 3.10 support) to fix that.

However, torch-directml wants torch~=2.4 while whisperx wants torch~=2.8, so the two can't co-exist in the same venv. I have made a fork that uses two separate venvs (one for audio-separator[dml] and another for everything else), and then I use a disgusting subprocess.run to run the audio-separator step in the other venv, with some truly heinous envvars to inject the location of the other venv into the parent python process, and then equally heinous parsing of stdout to plumb the output filenames back up to the parent process.

I also had to change the model from mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt. Roformer models make DirectML complain about "Invalid or unsupported data type ComplexFloat", so I chose UVR_MDXNET_KARA_2.onnx as it seemed to be the best-performing (by SDR) karaoke (lead vocals isolated, backing vocals kept in instrumental) model available in audio-separator --list-models. It sounds fine to me and now I can split a song into stems in <10s. Forced alignment/transcription of lyrics still takes a while, but if I skip language detection by hard-coding English then I can analyse most songs in under a minute using a 7800XT and 5900X.

For anyone curious, the fork is here: https://github.com/judilsteve/nightingale/tree/filthy_directml
It works on my machine and I suspect nobody else's. I wouldn't consider merging this, it's an atrocity. I made it for myself only, but someone might find it useful.