FROM python:3.11-slim

WORKDIR /app

# 1. ติดตั้ง System Dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    git git-lfs curl ffmpeg libgl1 libglib2.0-0 wget \
    && rm -rf /var/lib/apt/lists/* \
    && git lfs install

# 2. ติดตั้ง PyTorch 2.5.1 (สำคัญมาก! เพื่อแก้ Error comfy-kitchen ใน 2.6.0)
RUN pip install --no-cache-dir \
    --extra-index-url https://download.pytorch.org/whl/cpu \
    "torch==2.5.1" "torchvision==0.20.1" "torchaudio==2.5.1"

# 3. Clone ComfyUI และติดตั้ง Requirements
RUN git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git /app/ComfyUI \
    && pip install --no-cache-dir -r /app/ComfyUI/requirements.txt

# 4. ติดตั้ง Custom Nodes ที่จำเป็น
RUN cd /app/ComfyUI/custom_nodes && \
    git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager.git && \
    git clone --depth 1 https://github.com/city96/ComfyUI-GGUF.git && \
    git clone --depth 1 https://github.com/crystian/ComfyUI-Crystools.git && \
    git clone --depth 1 https://github.com/rgthree/rgthree-comfy.git && \
    git clone --depth 1 https://github.com/ty0x2333/ComfyUI-Dev-Utils.git && \
    git clone --depth 1 https://github.com/XLabs-AI/x-flux-comfyui.git && \
    git clone --depth 1 https://github.com/willmiao/ComfyUI-Lora-Manager.git && \
    git clone --depth 1 https://github.com/kijai/ComfyUI-KJNodes.git && \
    git clone --depth 1 https://github.com/lquesada/ComfyUI-Inpaint-CropAndStitch.git && \
    git clone --depth 1 https://github.com/gseth/ControlAltAI-Nodes.git && \
    git clone --depth 1 https://github.com/ssitu/ComfyUI_UltimateSDUpscale --recursive && \
    git clone --depth 1 https://github.com/nicolab28/ComfyUI-ClipProj.git && \
    git clone --depth 1 https://github.com/ethanfel/ComfyUI-Krea2TextEncoder.git && \
    git clone --depth 1 https://github.com/LAOGOU-666/Comfyui-Memory_Cleanup.git && \
    git clone --depth 1 https://github.com/KLL535/ComfyUI_SimpleButcher.git

# ติดตั้ง requirements ของ custom nodes
RUN cd /app/ComfyUI/custom_nodes/ComfyUI-GGUF && pip install --no-cache-dir -r requirements.txt || true
RUN cd /app/ComfyUI/custom_nodes/ComfyUI-Crystools && pip install --no-cache-dir -r requirements.txt || true
RUN cd /app/ComfyUI/custom_nodes/rgthree-comfy && pip install --no-cache-dir -r requirements.txt || true
RUN cd /app/ComfyUI/custom_nodes/ComfyUI-KJNodes && pip install --no-cache-dir -r requirements.txt || true

# 5. ติดตั้ง Python Packages เพิ่มเติม (เพิ่ม color-matcher สำหรับ Node ColorMatch)
RUN pip install --no-cache-dir \
    huggingface_hub psutil websocket-client opencv-python-headless \
    einops transformers safetensors color-matcher

# 6. สร้างโครงสร้างโฟลเดอร์
RUN mkdir -p /app/ComfyUI/models/checkpoints \
    /app/ComfyUI/models/diffusion_models \
    /app/ComfyUI/models/unet \
    /app/ComfyUI/models/text_encoders \
    /app/ComfyUI/models/llm \
    /app/ComfyUI/models/vae \
    /app/ComfyUI/models/loras \
    /app/ComfyUI/models/controlnet \
    /app/ComfyUI/models/upscale_models \
    /app/ComfyUI/output

# 7. ดาวน์โหลดโมเดลทั้งหมดที่ Workflow ต้องการ (ใช้ wget เพื่อป้องกัน RAM พุ่ง)
RUN wget -c https://huggingface.co/lilcheaty/Krea2-INT8-ConvRot/resolve/main/Krea2-Turbo-int8-ConvRot.safetensors -P /app/ComfyUI/models/diffusion_models/
RUN wget -c https://huggingface.co/Comfy-Org/flux2-klein-4B/resolve/main/split_files/diffusion_models/flux-2-klein-4b-nvfp4.safetensors -P /app/ComfyUI/models/diffusion_models/
RUN wget -c https://huggingface.co/DreamFast/qwen3-4b-heretic/resolve/main/comfyui/qwen3-4b-heretic_nvfp4.safetensors -P /app/ComfyUI/models/text_encoders/
RUN wget -c https://huggingface.co/Comfy-Org/Krea-2/resolve/main/text_encoders/qwen3vl_4b_fp8_scaled.safetensors -P /app/ComfyUI/models/text_encoders/
RUN wget -c https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/vae/ae.safetensors -P /app/ComfyUI/models/vae/
RUN wget -c https://huggingface.co/Comfy-Org/flux2-dev/resolve/main/split_files/vae/flux2-vae.safetensors -P /app/ComfyUI/models/vae/
RUN wget -c https://huggingface.co/Comfy-Org/Krea-2/resolve/main/vae/qwen_image_vae.safetensors -P /app/ComfyUI/models/vae/
RUN wget -c https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/photo/4xNomosWebPhoto_RealPLKSR.pth -P /app/ComfyUI/models/upscale_models/

# 8. Health Check
HEALTHCHECK --interval=30s --timeout=30s --start-period=10s --retries=3 \
    CMD curl -f http://localhost:7860/ || exit 1

# 9. คำสั่งเริ่มต้น
CMD ["python", "-u", "/app/ComfyUI/main.py", "--cpu", "--listen", "0.0.0.0", "--port", "7860", "--output-directory", "/app/ComfyUI/output", "--disable-smart-memory"]