# syntax=docker/dockerfile:1.14.0-labs

## OPTION 1 (pytorch container from docker hub)
## https://hub.docker.com/r/pytorch/pytorch/tags
FROM pytorch/pytorch:2.9.1-cuda13.0-cudnn9-devel
RUN addgroup --gid 1000 user && adduser --uid 1000 --gid 1000 user

## OPTION 2 (nvidia cuda container from docker hub)
## https://hub.docker.com/r/nvidia/cuda/tags
# FROM nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04

## OPTION 3 (ngc nvidia cuda container)
## https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/
# FROM nvcr.io/nvidia/pytorch:25.10-py3

## Setup ENV
ENV PIP_ROOT_USER_ACTION=ignore
ENV PIP_BREAK_SYSTEM_PACKAGES=1
## Example pip filter (prevents pip packages from being installed)
# ENV PIP_FILTER='/torchaudio$\|torchvision$/d'
## filter for x86_64
ENV PIP_FILTER='/xxxxxx$/d'
## Needed for aarch64 containers
# ENV PIP_FILTER='/onnxruntime-gpu/d'

## Install models and assets
## Do this first so we can avoid re-uploading at all costs
##
## Do as my separate smaller copy commands to make smaller layers
## Large layers can cause issues with uploads and caching
# COPY ./cache/models /usr/local/lib/models
COPY ./cache/assets /usr/local/lib/assets
COPY ./cache/models/checkpoints /usr/local/lib/models/checkpoints
COPY ./cache/models/clip /usr/local/lib/models/clip
COPY ./cache/models/clip_vision /usr/local/lib/models/clip_vision
COPY ./cache/models/controlnet /usr/local/lib/models/controlnet
COPY ./cache/models/diffusion_models /usr/local/lib/models/diffusion_models
COPY ./cache/models/loras /usr/local/lib/models/loras
COPY ./cache/models/style_models /usr/local/lib/models/style_models
COPY ./cache/models/text_encoders /usr/local/lib/models/text_encoders
COPY ./cache/models/unet /usr/local/lib/models/unet
COPY ./cache/models/vae /usr/local/lib/models/vae
COPY ./cache/models/prompt_generator /usr/local/lib/models/prompt_generator
COPY ./cache/models/model_patches /usr/local/lib/models/model_patches
COPY ./cache/models/deeplabv3 /usr/local/lib/models/deeplabv3
COPY ./cache/models/sam3 /usr/local/lib/models/sam3

## Add base libraries
RUN apt update && DEBIAN_FRONTEND=noninteractive TZ=US/Pacific apt install -y zsh git openssh-client wget bind9-host telnet iputils-ping ffmpeg libsm6 libxext6 libfreetype6-dev libcairo2-dev libjpeg-dev libgif-dev libopengl-dev python3 python3-pip

## Needed for OPTION 2 (nvidia cuda container) base container
## Dowload proper torch for the cuda included in the container
# RUN apt remove --yes python3-packaging
# RUN pip install --upgrade setuptools
# RUN pip install nvidia-pyindex
# RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130

## Install base pip libraries
RUN pip install setproctitle vastai accelerate websockets

## Configure github keys
## For accessing private repos (add private-repo-key.pub to github -> repo -> settings -> deploy keys)
ADD private-repo-key /etc/private-repo-key
RUN chmod 600 /etc/private-repo-key && echo "IdentityFile /etc/private-repo-key" >> /etc/ssh/ssh_config && echo "StrictHostKeyChecking no" >> /etc/ssh/ssh_config

## Install ComfyUI and map models and assets to already created directories
RUN git -c advice.detachedHead=false clone https://github.com/comfyanonymous/ComfyUI.git -b v0.8.2 /usr/local/src/comfyui
RUN cat /usr/local/src/comfyui/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin
RUN rm -rf /usr/local/src/comfyui/models && ln -sf /usr/local/lib/models /usr/local/src/comfyui/models && rm -rf /usr/local/src/comfyui/assets && ln -sf /usr/local/lib/assets /usr/local/src/comfyui/assets

## Some custom nodes will download into the diffusers cache
## We need to pre-download these models, but for now just link the cache to tmp so they can be downloaded dynamically
RUN mkdir -p /usr/local/src/comfyui/models/diffusers && ln -s /tmp /usr/local/src/comfyui/models/diffusers/.cache

## Install diffusers (required by some custom nodes)
RUN mkdir -p /mldata/huggingface/hub/ && chmod -R 777 /mldata
RUN pip install --src /usr/local/src -e git+https://github.com/huggingface/diffusers.git#egg=diffusers

## Install Hunyuan3D
## See Makefile for more information
RUN git clone https://github.com/kijai/ComfyUI-Hunyuan3DWrapper /usr/local/src/comfyui/custom_nodes/ComfyUI-Hunyuan3DWrapper
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-Hunyuan3DWrapper/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin
# RUN pip install pymeshlab
RUN pip install pymeshlab==2023.12.post3
# RUN cd /usr/local/src/comfyui/custom_nodes/ComfyUI-Hunyuan3DWrapper/hy3dgen/texgen/custom_rasterizer && python3 setup.py install
# RUN --device=nvidia.com/gpu cd /usr/local/src/comfyui/custom_nodes/ComfyUI-Hunyuan3DWrapper/hy3dgen/texgen/custom_rasterizer && python3 setup.py install

## Install comfy api wrapper
RUN pip install --src /usr/local/src -e git+https://github.com/deimos-deimos/comfy_api_simplified.git@main#egg=comfy_api_simplified

## Install KJNodes custom nodes
# 204f6d5aae73b10c0fe2fb26e61405fd6337bb77 is v1.2.3
# https://github.com/kijai/ComfyUI-KJNodes/commits/main/
RUN git clone https://github.com/kijai/ComfyUI-KJNodes.git /usr/local/src/comfyui/custom_nodes/ComfyUI-KJNodes
RUN cd /usr/local/src/comfyui/custom_nodes/ComfyUI-KJNodes/ && git checkout 204f6d5aae73b10c0fe2fb26e61405fd6337bb77
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-KJNodes/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install ask-gemini
RUN git clone https://github.com/Visionatrix/ComfyUI-Gemini /usr/local/src/comfyui/custom_nodes/ComfyUI-Gemini
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-Gemini/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install controlnet aux custom nodes
RUN git clone https://github.com/Fannovel16/comfyui_controlnet_aux/ /usr/local/src/comfyui/custom_nodes/comfyui_controlnet_aux
RUN cat /usr/local/src/comfyui/custom_nodes/comfyui_controlnet_aux/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin
COPY ./cache/models/controlnet/ControlNetHED.pth /usr/local/src/comfyui/custom_nodes/comfyui_controlnet_aux/ckpts/lllyasviel/Annotators/

## Install masquerade custom nodes
RUN git clone https://github.com/BadCafeCode/masquerade-nodes-comfyui /usr/local/src/comfyui/custom_nodes/masquerade-nodes-comfyui

## Install kaytool custom nodes and extensions
RUN git clone https://github.com/kk8bit/KayTool /usr/local/src/comfyui/custom_nodes/KayTool
RUN cat /usr/local/src/comfyui/custom_nodes/KayTool/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install LayerStyle custom nodes
RUN git clone https://github.com/chflame163/ComfyUI_LayerStyle /usr/local/src/comfyui/custom_nodes/ComfyUI_LayerStyle
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI_LayerStyle/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install ComfyUI_essentials
RUN git clone https://github.com/cubiq/ComfyUI_essentials.git /usr/local/src/comfyui/custom_nodes/ComfyUI_essentials
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI_essentials/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install efficiency-nodes-comfyui
RUN git clone https://github.com/jags111/efficiency-nodes-comfyui.git /usr/local/src/comfyui/custom_nodes/efficiency-nodes-comfyui
RUN cat /usr/local/src/comfyui/custom_nodes/efficiency-nodes-comfyui/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install RMBG
# RUN git clone https://github.com/reimage-ai/ComfyUI-RMBG /usr/local/src/comfyui/custom_nodes/ComfyUI-RMBG
RUN git clone git@github.com:reimage-ai/ComfyUI-RMBG.git /usr/local/src/comfyui/custom_nodes/ComfyUI-RMBG
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-RMBG/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install reimage custom nodes
# RUN git clone git@github.com:reimage-ai/comfy-custom-nodes.git /usr/local/src/comfyui/custom_nodes/comfy-custom-nodes
ADD ./comfy-custom-nodes /usr/local/src/comfyui/custom_nodes/comfy-custom-nodes
RUN cat /usr/local/src/comfyui/custom_nodes/comfy-custom-nodes/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install ReduxAdvanced
RUN git clone https://github.com/kaibioinfo/ComfyUI_AdvancedRefluxControl.git /usr/local/src/comfyui/custom_nodes/ComfyUI_AdvancedRefluxControl

## Install YANC
RUN git clone https://github.com/ALatentPlace/ComfyUI_yanc /usr/local/src/comfyui/custom_nodes/ComfyUI_yanc

## Install rgthree utility custom nodes
RUN git clone https://github.com/rgthree/rgthree-comfy /usr/local/src/comfyui/custom_nodes/rgthree-comfy

## Install impact pack
RUN git clone https://github.com/ltdrdata/ComfyUI-Impact-Pack /usr/local/src/comfyui/custom_nodes/ComfyUI-Impact-Pack
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-Impact-Pack/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install ComfyUI_Qwen2_5-VL-Instruct
RUN git clone git@github.com:reimage-ai/ComfyUI_Qwen2_5-VL-Instruct.git /usr/local/src/comfyui/custom_nodes/ComfyUI_Qwen2_5-VL-Instruct
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI_Qwen2_5-VL-Instruct/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install comfy_mtb
RUN git clone https://github.com/melMass/comfy_mtb /usr/local/src/comfyui/custom_nodes/comfy_mtb
RUN cat /usr/local/src/comfyui/custom_nodes/comfy_mtb/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin
RUN cd /usr/local/src/comfyui/custom_nodes/comfy_mtb && git submodule update --init --remote --recursive

## Install ComfyUI-Easy-User
RUN git clone https://github.com/yolain/ComfyUI-Easy-Use /usr/local/src/comfyui/custom_nodes/ComfyUI-Easy-Use
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-Easy-Use/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Install ComfyUI-GGUF
RUN git clone https://github.com/city96/ComfyUI-GGUF /usr/local/src/comfyui/custom_nodes/ComfyUI-GGUF
RUN cat /usr/local/src/comfyui/custom_nodes/ComfyUI-GGUF/requirements.txt | sed -e $PIP_FILTER | pip install -r /dev/stdin

## Now remove private key from the container
RUN rm -f /etc/private-repo-key

# Temporary fix for https://github.com/google-ai-edge/mediapipe/issues/6192
RUN pip install mediapipe==0.10.21

ADD ./src /usr/local/src/comfy-service/src

WORKDIR /workspace
ENV HOME=/tmp
ENTRYPOINT ["/bin/sh", "-c"]
