# Get cuda version from https://hub.docker.com/r/nvidia/cuda
FROM nvidia/cuda:12.9.1-devel-ubuntu20.04 AS build-deps

ENV TZ=Etc/UTC
ENV DEBIAN_FRONTEND=noninteractive

# Install build dependencies
RUN apt update -y && \
    apt install -y software-properties-common && \
    add-apt-repository ppa:ubuntu-toolchain-r/test && \
    apt update && \
    apt install -y \
    curl \
    gcc-13 \
    g++-13 \
    make \
    wget \
    libffi-dev \
    libssl-dev \
    libbz2-dev \
    zlib1g-dev \
    liblzma-dev \
    libreadline-dev \
    libsqlite3-dev \
    libopenblas-dev \
    libblas-dev \
    libjpeg-dev \
    libpng-dev \
    libavcodec-dev \
    libavformat-dev \
    libswscale-dev \
    libomp-dev \
    libnuma-dev \
    m4 \
    cmake \
    git \
    ninja-build \
    libprotobuf-dev \
    protobuf-compiler \
    libeigen3-dev \
    zip \
    unzip

# Build Python from source
FROM build-deps AS build-python
WORKDIR /usr/src

# Get release version from https://www.python.org/downloads/source/
ARG Python_Release=3.12.11

RUN wget https://www.python.org/ftp/python/${Python_Release}/Python-${Python_Release}.tgz && \
    tar xzf Python-${Python_Release}.tgz

WORKDIR /usr/src/Python-${Python_Release}

RUN ./configure --enable-optimizations && \
    make -j$(nproc) && \
    make altinstall

RUN Python_Version=$(echo "$Python_Release" | cut -d. -f1,2) && \
    ln -sf /usr/local/bin/python$Python_Version /usr/bin/python && \
    ln -sf /usr/local/bin/pip$Python_Version /usr/bin/pip

# Build pytorch from souce
FROM build-python AS build-pytorch
WORKDIR /usr/src

# Get release version from https://github.com/pytorch/pytorch/releases
ARG Pytorch_Release=2.7.1
# Set the number of jobs for the Pytorch_Build -> (Gb of RAM)/17 rounded down to the nearest whole number
ARG Pytorch_Jobs=4
# Set the CUDA arch to compile for
ARG CUDA_ARCH_LIST="7.5;8.0;8.6;8.9"

RUN git clone --recursive -b v${Pytorch_Release} https://github.com/pytorch/pytorch.git
WORKDIR /usr/src/pytorch

ENV USE_CUDA=1 \
    USE_CUDNN=1 \
    BUILD_CAFFE2=0 \
    USE_DISTRIBUTED=0 \
    USE_FBGEMM=0 \
    USE_MKLDNN=0 \
    USE_NCCL=0 \
    USE_QNNPACK=0 \
    USE_XNNPACK=0 \
    USE_OPENMP=0 \
    USE_TENSORPIPE=0 \
    USE_SYSTEM_EIGEN_INSTALL=0 \
    BUILD_TEST=0 \
    CMAKE_PREFIX_PATH="/usr/local/cuda" \
    CUDA_HOME="/usr/local/cuda" \
    USE_NVRTC=1

RUN python -m pip install numpy typing_extensions future sympy

RUN pip install -r requirements.txt

RUN git submodule sync && \
    git submodule update --init --recursive

RUN MAX_JOBS=${Pytorch_Jobs} && \
    TORCH_CUDA_ARCH_LIST=${CUDA_ARCH_LIST} && \
    python setup.py install

# Build ComfyUI from source
FROM build-pytorch AS publish
WORKDIR /app

# Get release version from https://github.com/comfyanonymous/ComfyUI/releases
ARG ComfyUI_Release=0.3.43

RUN wget https://github.com/comfyanonymous/ComfyUI/archive/refs/tags/v${ComfyUI_Release}.tar.gz && \
    tar xzf v${ComfyUI_Release}.tar.gz && \
    rm v${ComfyUI_Release}.tar.gz && \
    mv /app/ComfyUI-${ComfyUI_Release} /app/comfyui/
WORKDIR /app/comfyui

RUN python -m pip install --upgrade pip && \
    python -m pip install -r requirements.txt

EXPOSE 8188

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