# Use NVIDIA CUDA 12.8.1 image based on Ubuntu 22.04
FROM nvidia/cuda:12.8.1-devel-ubuntu22.04

# Set environment variables for a cleaner setup
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;8.9;9.0,12.0"
# CUDA Architecture Guide:
# ----------------------
# 6.0, 6.1: Pascal (GTX 1080, 1070, 1060)
# 7.0: Volta (V100)
# 7.2: Turing (T4)
# 7.5: Turing (RTX 2080 Ti, 2080, 2070, 2060)
# 8.0: Ampere (A100)
# 8.6: Ampere Consumer (RTX 3090, 3080, 3070, 3060)
# 8.9: Ada Loverlace (RTX 4090, 4080, 4070)
# 9.0: Hopper (H200, H100, GH200
# 12.0: Blacwell (GeForce RTX 5090, RTX 5080, RTX 5070 Ti, RTX 5070, RTX 5060 Ti, RTX 5060, \
#   RTX PRO 6000 Blackwell, RTX PRO 5000 Blackwell, RTX PRO 4500 Blackwell, RTX PRO 4000 Blackwell


# Install python 3.12 and base tools
RUN apt-get update && apt-get install -y \
    software-properties-common && \
    add-apt-repository ppa:deadsnakes/ppa && \
    apt-get update && apt-get install -y \
    python3.12 \
    python3.12-dev \
    python3.12-venv \
    tmux \
    git \
    wget \
    curl \
    gnupg2 \
    build-essential \
    ninja-build \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# Set Python 3.12 as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 && \
    update-alternatives --set python3 /usr/bin/python3.12

# Install pip for Python 3.12
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12

# Install PyTorch with CUDA 12.8 support
RUN python3 -m pip install --pre --no-cache-dir torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/nightly/cu128

# Clone the repositories and install requirements
WORKDIR /app
RUN git clone https://github.com/comfyanonymous/ComfyUI.git . && \
    python3 -m pip install --no-cache-dir -r requirements.txt && \
    python3 -m pip install --no-cache-dir deepdiff

# Custom nodes setup
WORKDIR /app/custom_nodes
RUN git clone https://github.com/crystian/ComfyUI-Crystools.git && \
    cd ComfyUI-Crystools && \
    python3 -m pip install --no-cache-dir -r requirements.txt && \
    cd .. && \
    git clone https://github.com/nonnonstop/comfyui-faster-loading.git && \
    git clone https://github.com/ltdrdata/ComfyUI-Manager.git

# Expose the necessary port for the application
EXPOSE 8188/tcp

# Final launch command
WORKDIR /app
ENTRYPOINT ["python3", "main.py", "--listen", "--cuda-device", "0", "--disable-cuda-malloc", "--port", "8188"]