ComfyUI on ROCM Dockerfile

Dockerfile

FROM ubuntu:rolling AS rocm

# Initialize the image
# Modify to pre-install dev tools and ROCm packages
ARG ROCM_VERSION=6.3.2
ARG AMDGPU_VERSION=6.3.2

RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates curl gnupg  software-properties-common && \
  (curl -sL http://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor -o /etc/apt/trusted.gpg.d/amd-archive-keyring.gpg) && \
  sh -c 'echo deb [arch=amd64] http://repo.radeon.com/rocm/apt/$ROCM_VERSION/ jammy main > /etc/apt/sources.list.d/rocm.list' && \
  sh -c 'echo deb [arch=amd64] https://repo.radeon.com/amdgpu/$AMDGPU_VERSION/ubuntu jammy main > /etc/apt/sources.list.d/amdgpu.list' && \
  (printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | tee /etc/apt/preferences.d/rocm-pin-600 ) && \
  # add-apt-repository --yes ppa:deadsnakes/ppa && \
  apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
  sudo \
  libelf1 \
  libnuma-dev \
  build-essential \
  git \
  vim-nox \
  cmake-curses-gui \
  kmod \
  file \
  python3 \
  python3-pip \
  python3-venv \
  libgl1 \
  rocm-dev \
  rocm-hip-sdk rocm-hip-libraries && \
  apt-get clean && \
  rm -rf /var/lib/apt/lists/*

# Grant members of 'sudo' group passwordless privileges
# Comment out to require sudo
# COPY sudo-nopasswd /etc/sudoers.d/sudo-nopasswd

# # This is meant to be used as an interactive developer container
# # Create user rocm-user as member of sudo group
# # Append /opt/rocm/bin to the system PATH variable
# RUN useradd --create-home -G sudo,video --shell /bin/bash rocm-user
# #    sed --in-place=.rocm-backup 's|^\(PATH=.*\)"$|\1:/opt/rocm/bin"|' /etc/environment

# USER rocm-user
# WORKDIR /home/rocm-user
ENV PATH "${PATH}:/opt/rocm/bin"
ENV CC /opt/rocm/llvm/bin/clang 
ENV CXX /opt/rocm/llvm/bin/clang++

# Default to a login shell
CMD ["bash", "-l"]

FROM rocm AS torch

ENV TORCH_ARCH=nightly/rocm6.3
RUN update-alternatives --install /usr/bin/python python $(which python3) 9999 \
    && update-alternatives --install /usr/bin/pip pip $(which pip3) 9999

RUN pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/$TORCH_ARCH --break-system-packages

# Set necessary export for Stable Diffusion to work on RDNA 5700XT using ROCm
ENV HSA_OVERRIDE_GFX_VERSION=10.3.0
ENV PYTORCH_ROCM_ARCH=gfx1030
ENV GFX_ARCH=gfx1030

# RUN curl https://raw.githubusercontent.com/wiki/ROCm/pytorch/files/install_kdb_files_for_pytorch_wheels.sh | bash
RUN ( \
    TORCH_INSTALL_PATH=$(pip show torch | grep Location | cut -d" " -f 2); \
    EXTRACT_DIR=extract_miopen_dbs; \
    rm -rf $EXTRACT_DIR && mkdir -p "$EXTRACT_DIR" && cd "$EXTRACT_DIR"; \
    wget https://repo.radeon.com/rocm/apt/6.3.2/pool/main/m/miopen-hip-gfx1030kdb/miopen-hip-gfx1030kdb_3.3.0.60302-66~24.04_amd64.deb; \
    dpkg-deb -xv $(ls *deb) . ; \
    echo "Copying kdb files to ${TORCH_INSTALL_PATH}/torch/share" ;\
    cp -ra opt/rocm-*/share/miopen $TORCH_INSTALL_PATH/torch/share ;\
    cd .. && rm -rf "$EXTRACT_DIR" ;\
    echo "Successfully installed MIOpen kernel database files" \
    )

FROM torch AS comfy

RUN useradd -u 60091 --create-home -G sudo,video --shell /bin/bash rocm-user
RUN mkdir /sd && chown rocm-user:video /sd
USER rocm-user

RUN git clone --depth=1 https://github.com/comfyanonymous/ComfyUI /sd
WORKDIR /sd

# Expose configuration files
RUN mkdir /sd/config

# Setup venv
RUN python3 -m venv --system-site-packages venv
ENV VIRTUAL_ENV="/sd/venv"
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN venv/bin/pip install -r requirements.txt numba matplotlib opencv-python

# Install bitsandbytes
RUN ( \
  git clone --recursive https://github.com/ROCm/bitsandbytes /bitsandbytes; \
  cd /bitsandbytes; \
  git checkout rocm_enabled_multi_backend; \
  pushd bitsandbytes/triton; \
  wget https://raw.githubusercontent.com/triton-lang/kernels/refs/heads/main/kernels/matmul_perf_model.py; \
  sed -i 's/triton\.ops\.matmul_perf_model/.matmul_perf_model/g' int8_matmul_mixed_dequantize.py int8_matmul_rowwise_dequantize.py; \
  popd; \
  source /sd/venv/bin/activate; \
  pip install -r requirements-dev.txt; \
  cmake -DCOMPUTE_BACKEND=hip -DBNB_ROCM_ARCH="gfx1030" -S . \
  make \
  pip install . \
)

EXPOSE 8188

CMD ["venv/bin/python", "main.py", "--listen"]