Dockerfile

# ComfyUI with Wan Animate for RunPod Serverless
FROM --platform=linux/amd64 runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV COMFYUI_VERSION="0.3.16"
ENV TORCH_CUDA_ARCH_LIST="8.6+PTX;8.9+PTX;9.0+PTX"

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    wget \
    curl \
    ffmpeg \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    vim \
    rsync \
    unzip \
    tar \
    && rm -rf /var/lib/apt/lists/*

# Create working directory
WORKDIR /comfyui

# Clone ComfyUI with specific version
RUN git clone https://github.com/comfyanonymous/ComfyUI.git . && \
    git checkout tags/v${COMFYUI_VERSION} && \
    pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir -r requirements.txt

# Install additional Python packages optimized for serverless
RUN pip install --no-cache-dir \
    runpod>=1.3.0 \
    onnxruntime-gpu>=1.18.0 \
    opencv-python-headless>=4.9.0.80 \
    mediapipe>=0.10.10 \
    insightface>=0.7.3 \
    segment-anything-2>=0.1.0 \
    torchmetrics>=1.3.0 \
    einops>=0.8.0

# Create model directories structure
RUN mkdir -p \
    models/checkpoints \
    models/clip \
    models/unet \
    models/vae \
    models/clip_vision \
    models/loras \
    models/controlnet \
    models/upscale_models \
    models/sam2 \
    models/animatediff \
    models/animatediff_motion_modules \
    models/facerestore_models \
    models/insightface \
    models/segment_anything \
    custom_nodes/ComfyUI-Frame-Interpolation/ckpts/rife \
    custom_nodes/comfyui_controlnet_aux/ckpts/onnx \
    input \
    output \
    temp

# Install Custom Nodes with error handling and optimization
WORKDIR /comfyui/custom_nodes

# 1. ComfyUI Manager (essential for node management)
RUN set -e && \
    git clone https://github.com/ltdrdata/ComfyUI-Manager.git && \
    cd ComfyUI-Manager && \
    pip install --no-cache-dir -e . || echo "Manager install partially failed"

# 2. Video Helper Suite (VHS) - REQUIRED for video processing
RUN set -e && \
    git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git && \
    cd ComfyUI-VideoHelperSuite && \
    pip install --no-cache-dir -r requirements.txt && \
    pip install --no-cache-dir -e . || echo "VHS install failed"

# 3. ComfyUI Advanced ControlNet
RUN set -e && \
    git clone https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet.git && \
    cd ComfyUI-Advanced-ControlNet && \
    pip install --no-cache-dir -r requirements.txt 2>/dev/null || echo "ControlNet deps failed"

# 4. rgthree nodes
RUN set -e && \
    git clone https://github.com/rgthree/rgthree-comfy.git && \
    cd rgthree-comfy && \
    pip install --no-cache-dir -e . || echo "rgthree install failed"

# 5. ComfyUI ControlNet Aux (for DWPreprocessor, ViTPose)
RUN set -e && \
    git clone https://github.com/Fannovel16/comfyui_controlnet_aux.git && \
    cd comfyui_controlnet_aux && \
    pip install --no-cache-dir -r requirements.txt && \
    pip install --no-cache-dir -e . || echo "ControlNet Aux install failed"

# 6. ComfyUI Impact Pack (for face detection and enhancement)
RUN set -e && \
    git clone https://github.com/ltdrdata/ComfyUI-Impact-Pack.git && \
    cd ComfyUI-Impact-Pack && \
    pip install --no-cache-dir -r requirements.txt 2>/dev/null && \
    pip install --no-cache-dir -e . || echo "Impact Pack install failed"

# 7. ComfyUI Essentials
RUN set -e && \
    git clone https://github.com/cubiq/ComfyUI_essentials.git && \
    cd ComfyUI_essentials && \
    pip install --no-cache-dir -e . || echo "Essentials install failed"

# 8. WAS Node Suite
RUN set -e && \
    git clone https://github.com/WASasquatch/was-node-suite-comfyui.git && \
    cd was-node-suite-comfyui && \
    pip install --no-cache-dir -r requirements.txt 2>/dev/null && \
    pip install --no-cache-dir -e . || echo "WAS Suite install failed"

# 9. ComfyUI Segment Anything 2
RUN set -e && \
    git clone https://github.com/kijai/ComfyUI-segment-anything-2.git && \
    cd ComfyUI-segment-anything-2 && \
    pip install --no-cache-dir -r requirements.txt && \
    pip install --no-cache-dir -e . || echo "SAM2 install failed"

# 10. ComfyUI Custom Nodes for Animation (essential for Wan Animate)
RUN set -e && \
    git clone https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet.git && \
    cd ComfyUI-Advanced-ControlNet && \
    pip install --no-cache-dir -e . || echo "Advanced ControlNet failed"

# 11. ComfyUI LayerStyle (for advanced compositing)
RUN set -e && \
    git clone https://github.com/colorfulcloud/comfyui-layerstyle.git && \
    cd comfyui-layerstyle && \
    pip install --no-cache-dir -e . || echo "LayerStyle install failed"

# 12. FaceFace Detailer (for face enhancement)
RUN set -e && \
    git clone https://github.com/LEv145/ComfyUI-faceface.git && \
    cd ComfyUI-faceface && \
    pip install --no-cache-dir -e . || echo "FaceFace install failed"

# 13. ComfyUI IP Adapter (for image prompting)
RUN set -e && \
    git clone https://github.com/TencentARC/comfyui_ipadapter_plus.git && \
    cd comfyui_ipadapter_plus && \
    pip install --no-cache-dir -e . || echo "IP Adapter install failed"

WORKDIR /comfyui

# Create startup script for serverless optimization
COPY startup.sh /comfyui/startup.sh
RUN chmod +x /comfyui/startup.sh

# Create RunPod handler
COPY handler.py /comfyui/handler.py

# Add model downloader script
COPY download_models.py /comfyui/download_models.py
RUN chmod +x /comfyui/download_models.py

# Download all required models
RUN python3 /comfyui/download_models.py

# Expose port
EXPOSE 8188

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
    CMD curl -f http://localhost:8188/ || exit 1

# Set entrypoint
ENTRYPOINT ["/comfyui/startup.sh"]