rocm docker image

[fatmalama]

I've made a docker image based on the pytorch rocm image for use with (compatible) AMD GPUs

you can check it over here https://hub.docker.com/r/mikelpr/stemgen-rocm

here's the Containerfile used to build it:

FROM docker.io/rocm/pytorch:latest

RUN bash -xec 'apt-get update && apt-get upgrade -y && apt-get install -y ffmpeg libboost-python1.74-dev libboost-python1.74.0 cmake libutfcpp-dev && apt-get clean'
WORKDIR /usr/local/src
RUN bash -xec 'wget -O taglib.tar.gz https://github.com/taglib/taglib/releases/download/v2.0.1/taglib-2.0.1.tar.gz && \
    tar xf taglib.tar.gz && cd taglib-2.0.1 && \
    cmake -DCMAKE_INSTALL_PREFIX=/usr -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=ON . && \
    make -j && make install && cd .. && rm -rf taglib-2.0.1 taglib.tar.gz'

RUN bash -xec 'pip install "git+https://github.com/acolombier/stemgen.git@v0.4.0#egg=stemgen" \
        --index-url "https://download.pytorch.org/whl/pytorch-triton-rocm" \
        --extra-index-url https://pypi.org/simple'

ENTRYPOINT [ "/opt/conda/envs/py_3.12/bin/stemgen" ]
CMD []
VOLUME [ "/data" ]
WORKDIR /data

[fatmalama]

I don't know why I went with pytorch-triton-rocm instead of rocm6.4 but I just checked that it works with https://download.pytorch.org/whl/rocm6.4. also was kinda exploring using rocm-portable from RushingAlien as it's much much smaller than AMD's build but I've arrived to see that a rust rewrite is taking place so I'll hold on until that materializes and see lol

[acolombier]

Hi @mikelpr , thanks for your contribution!

As of the coming 0.5.0, stemgen is been refactored into Rust, leveraging the ONNX version of Demucs.
I am very keen to get ROCM support out of the box, but I believe it will require to find the dependencies that needs to be added to the image (example with CUDA). I don't have any hardware to test, but I might try to at least get the DLLs in place. Feel free to go ahead and try that if you are interested!

[fatmalama]

I'll look into how I can contribute ROCm support into the rust version :)