# Stage 1: Base image with common dependencies
FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04

# Prevents prompts from packages asking for user input during installation
ENV DEBIAN_FRONTEND=noninteractive
# Prefer binary wheels over source distributions for faster pip installations
ENV PIP_PREFER_BINARY=1
# Ensures output from python is printed immediately to the terminal without buffering
ENV PYTHONUNBUFFERED=1 

RUN chmod 1777 /tmp

# Install Python, git and other necessary tools
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget

# Clean up to reduce image size
RUN apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/*

# Install ComfyUI dependencies
RUN python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# Install runpod
RUN pip3 install runpod requests huggingface-hub

COPY src/comfyui /comfyui
RUN rm -rdf /comfyui/models && mkdir /comfyui/models/
COPY models/checkpoints /comfyui/models/checkpoints
COPY models/clip /comfyui/models/clip
COPY models/clip_vision /comfyui/models/clip_vision
COPY models/configs /comfyui/models/configs
COPY models/diffusion_models /comfyui/models/diffusion_models
COPY models/facedetection /comfyui/models/facedetection
COPY models/facerestore_models /comfyui/models/facerestore_models
COPY models/insightface /comfyui/models/insightface
COPY models/inspyrenet /comfyui/models/inspyrenet
COPY models/pulid /comfyui/models/pulid
COPY models/rembg /comfyui/models/rembg
COPY models/roop /comfyui/models/roop
COPY models/ultralytics /comfyui/models/ultralytics
COPY models/upscale_models /comfyui/models/upscale_models
COPY models/vae /comfyui/models/vae

RUN pip3 install --upgrade -r /comfyui/requirements.txt

RUN find /comfyui/custom_nodes -name 'requirements.txt' -exec bash -c 'pip install -r <(grep -v -E "^torch|^torchsde|^torchvision|^transformers|^safetensors|^cuml-cu11|^xformers|^pytorch_lightning" {})' \;
RUN cd /comfyui && find ./custom_nodes -type f -name 'install.py' -exec python3 {} \;

COPY src/patch /tmp/patch
RUN cp /tmp/patch/face_analysis.py $(python3 -c "import site; print(site.getsitepackages()[0])")/insightface/app/face_analysis.py

# 卡通人脸检测
RUN huggingface-cli download "deepghs/anime_face_detection" face_detect_v1.4_s/model.onnx

# Support for the network volume
ADD src/extra_model_paths.yaml /comfyui/
WORKDIR /

# Add the start and the handler
ADD src/start.sh src/rp_handler.py test_input.json /
RUN chmod +x /start.sh

# Start the container
CMD /start.sh