# Use the official ROCm base image
FROM ubuntu:26.04 AS reqs
# Set build variables
ENV PIP_BREAK_SYSTEM_PACKAGES=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive
# Install libraries
RUN apt update -y && \
    apt upgrade -y && \
    apt install -y --no-install-recommends \
    ca-certificates \
    rocm \
    hipcc \
    libhipblaslt-dev \
    libomp-22-dev \
    git \
    cmake \
    libnuma-dev \
    build-essential \
    zlib1g-dev \
    libncurses5-dev \
    libopenblas-dev \
    libpthreadpool-dev \
    libgoogle-glog-dev \
    libgtest-dev \
    libbenchmark-dev \
    libgdbm-dev \
    libnss3-dev \
    libssl-dev \
    libreadline-dev \
    libffi-dev \
    libsqlite3-dev \
    libnuma-dev \
    libelf-dev \
    libpciaccess-dev \
    protobuf-compiler \
    libprotobuf-dev \
    wget \
    libbz2-dev \
    liblzma-dev \
    && rm -rf /var/lib/apt/lists/*


# Build python
FROM reqs AS python
WORKDIR /build/python
RUN wget https://www.python.org/ftp/python/3.11.15/Python-3.11.15.tgz && \
    tar -xf Python-3.11.15.tgz && \
    cd Python-3.11.15 && \
    ./configure --enable-optimizations --enable-shared LDFLAGS="-Wl,-rpath /usr/local/lib" && \
    make -j 80 && \
    make install && \
    ldconfig /usr/local/lib && \
    wget https://bootstrap.pypa.io/get-pip.py && \
    python3 get-pip.py && \
    ln -s /usr/local/bin/pip3 /bin/pip && \
    ln -s /usr/local/bin/pip3 /bin/pip3 && \
    ln -s /usr/local/bin/python3 /bin/python && \
    ln -s /usr/local/bin/python3 /bin/python3.11


# Build pytorch
FROM python AS pytorch
# Setup Envronment
ENV PYTORCH_ROCM_ARCH=1100
ENV USE_ROCM=1
ENV USE_CUDA=0
ENV USE_AOTRITON=0
ENV USE_SYSTEM_PTHREADPOOL=1
ENV USE_SYSTEM_PROTOC=1
ENV USE_ONNX=OFF
ENV BUILD_CUSTOM_PROTOBUF=OFF
ENV USE_XNNPACK=OFF
ENV USE_XNNPACK=OFF
ENV USE_QNNPACK=OFF
ENV USE_PYTORCH_QNNPACK=OFF
ENV USE_NNPACK=OFF
ENV USE_KINETO=OFF
ENV USE_OPENMP=OFF
ENV USE_TENSORPIPE=OFF
ENV ROCM_PATH=/usr
ENV MAX_JOBS=12
ENV DEVICE_LIB_PATH=/usr/lib/llvm-21/lib/clang/21/amdgcn/bitcode
ENV HIPFLAGS="--rocm-device-lib-path=$DEVICE_LIB_PATH"
ENV CMAKE_ARGS="-DCMAKE_HIP_COMPILER_WORKS=1 \
-DCMAKE_CXX_COMPILER_WORKS=1 \
-DCMAKE_C_COMPILER_WORKS=1 \
-DCMAKE_HIP_COMPILER_ROCM_LIB_PATH=$DEVICE_LIB_PATH \
-DCMAKE_HIP_FLAGS=--rocm-device-lib-path=$DEVICE_LIB_PATH \
-DCMAKE_C_COMPILER=/usr/lib/llvm-21/bin/clang \
-DCMAKE_CXX_COMPILER=/usr/lib/llvm-21/bin/clang++ \
-DCMAKE_HIP_COMPILER=/usr/lib/llvm-21/bin/clang++ \
-DPython3_ROOT_DIR=/usr/local \
-DPYTHON_EXECUTABLE=/usr/local/bin/python3.11 \
-DPYTHON_LIBRARY=/usr/local/lib/libpython3.11.so \
-DPYTHON_INCLUDE_DIR=/usr/local/include/python3.11"
# Pull and build pytorch
WORKDIR /build/pytorch
RUN git clone https://github.com/pytorch/pytorch.git && \
    cd pytorch && \
    git submodule update --init --recursive && \
    make triton && \
    python tools/amd_build/build_amd.py && \
    pip install -r requirements.txt && \
    mkdir -p /opt/rocm/bin \
    ln -s /usr/bin/hipcc /opt/rocm/bin/hipcc \
    ln -s /usr/lib/llvm-21/bin/clang++ /usr/bin/clang++ \
    python3 setup.py build && \
    python3 setup.py install


# Build ComfyUI
FROM pytorch AS comfyui
WORKDIR /app
RUN git clone https://github.com/comfyanonymous/ComfyUI.git .
RUN pip3 install -r requirements.txt
# Setup Manager
WORKDIR /app/custom_nodes
RUN git clone https://github.com/Comfy-Org/ComfyUI-Manager.git
# Setup for run
WORKDIR /app
EXPOSE 8188
# Launch
CMD ["python3", "main.py", "--listen", "0.0.0.0"]