# 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
# Set path variables
ENV ROCM_PATH=/usr
ENV HIP_PATH=/usr
# Install libraries
RUN apt update -y && \
    apt install -y --no-install-recommends \
    rocm \
    hipcc \
    libhipblaslt-dev \
    libomp-22-dev \
    git \
    cmake \
    libnuma-dev \
    build-essential \
    zlib1g-dev \
    libncurses5-dev \
    libgdbm-dev \
    libnss3-dev \
    libssl-dev \
    libreadline-dev \
    libffi-dev \
    libsqlite3-dev \
    wget \
    libbz2-dev \
    liblzma-dev \
    && rm -rf /var/lib/apt/lists/*

# Build python3.12
FROM reqs AS python312
WORKDIR /build/python312
RUN wget https://www.python.org/ftp/python/3.12.9/Python-3.12.9.tgz && \
    tar -xf Python-3.12.9.tgz && \
    cd Python-3.12.9 && \
    ./configure --enable-optimizations && \
    make -j 80 && \
    make install && \
    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.12

# Build pytorch
FROM python312 AS pytorch
ENV CMAKE_ARGS="$CMAKE_ARGS -DUSE_KINETO=OFF"
ENV PYTORCH_ROCM_ARCH=1100
ENV USE_ROCM=1
ENV USE_CUDA=0
ENV ROCM_PATH=/usr
ENV DEVICE_LIB_PATH=/usr/lib/llvm-21/lib/clang/21/amdgcn/bitcode
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++ \
-DUSE_KINETO=OFF \
-DUSE_OPENMP=OFF"

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 && \
    export HIPFLAGS="--rocm-device-lib-path=$DEVICE_LIB_PATH" && \
    export CMAKE_ARGS="-DCMAKE_HIP_FLAGS=--rocm-device-lib-path=$DEVICE_LIB_PATH -DCMAKE_HIP_COMPILER_ROCM_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++" && \
    python3 setup.py build && \
    python3 setup.py install

# Build ComfyUI
FROM pytorch AS comfyui
RUN git clone https://github.com/comfyanonymous/ComfyUI.git .

# Install non torch dependant requirements
RUN pip3 install -r requirements.txt

# Setup Manager
WORKDIR /app/custom_nodes
RUN git clone https://github.com/Comfy-Org/ComfyUI-Manager.git

WORKDIR /app
EXPOSE 8188

CMD ["python3", "main.py", "--listen", "0.0.0.0"]