Dockerfile.rocm

# Use ROCm base image with Python
FROM rocm/dev-ubuntu-22.04:7.2-complete

# Set the working directory in the container
WORKDIR /workdir

# Install necessary packages
RUN apt-get update && apt-get install -y --no-install-recommends \
    ffmpeg \
    python3 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

RUN python3 -m pip install --upgrade pip

# Install PyTorch with ROCm support
RUN --mount=type=cache,target=/root/.cache \
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2

# Install audio-separator with GPU support
RUN --mount=type=cache,target=/root/.cache \
    pip3 install "audio-separator[gpu]" onnxruntime-rocm

# Default environment variables for AMD RX 6600 series (gfx1032)
# Override these for other GPUs:
#   RX 7900: HSA_OVERRIDE_GFX_VERSION=11.0.0, PYTORCH_ROCM_ARCH=gfx1100
#   MI250X:  HSA_OVERRIDE_GFX_VERSION=9.4.2, PYTORCH_ROCM_ARCH=gfx90a
ARG HSA_OVERRIDE_GFX_VERSION=10.3.2
ARG PYTORCH_ROCM_ARCH=gfx1030
ENV HSA_OVERRIDE_GFX_VERSION=${HSA_OVERRIDE_GFX_VERSION}
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}

# Run audio-separator when the container launches
ENTRYPOINT ["audio-separator"]