FROM nvidia/cuda:13.1.0-cudnn-devel-ubuntu24.04    

ENV PIP_BREAK_SYSTEM_PACKAGES=1
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    ca-certificates \
    git \
    python3 \
    python3-pip \
    libgl1 \
    libglx-mesa0 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# Create User Permissions
ARG DOCKER_USER=default_user
RUN useradd -m "$DOCKER_USER"

# Clone vladmandic's stable diffusion 
RUN git clone --depth 1 https://github.com/comfyanonymous/ComfyUI

WORKDIR /ComfyUI/custom_nodes/
RUN \   
  git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager \
  && git clone --depth 1 https://github.com/city96/ComfyUI-GGUF \
  && git clone --depth 1 https://github.com/SeanScripts/ComfyUI-Unload-Model.git \
  && git clone --depth 1 https://github.com/ssitu/ComfyUI_UltimateSDUpscale --recursive \
  && git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Impact-Pack \
  && git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Impact-Subpack \
  && git clone --depth 1 https://github.com/talesofai/comfyui-browser \
  && git clone --depth 1 https://github.com/SEkINVR/ComfyUI-SaveAs.git

WORKDIR /ComfyUI

RUN pip install --no-cache-dir --break-system-packages -r requirements.txt \
    && pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-Manager/requirements.txt \
    && pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-GGUF/requirements.txt \
    && pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-Impact-Pack/requirements.txt \
    && pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-Impact-Subpack/requirements.txt \
    && pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/comfyui-browser/requirements.txt \
    && pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-SaveAs/requirements.txt \
    && pip install --no-cache-dir --break-system-packages gitpython psutil

# Install prebuilt flash-attn
RUN pip install --break-system-packages https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.22/flash_attn-2.8.1+cu130torch2.9-cp312-cp312-linux_x86_64.whl

# Set ownership and move models directory
RUN chown -R $DOCKER_USER:$DOCKER_USER /ComfyUI && \
    mv /ComfyUI/models /ComfyUI/models.orig

# Copy start script
COPY ./start.sh /usr/local/bin/
RUN chmod +x /usr/local/bin/start.sh

# Switch to non-root user
USER $DOCKER_USER
WORKDIR /ComfyUI

CMD ["/usr/local/bin/start.sh"]