FROM docker.io/nvidia/cuda:13.0.2-runtime-ubuntu24.04

ARG DEBIAN_FRONTEND=noninteractive
ARG COMFYUI_COMMIT=b33e2b55cae074eca5aec96283cceac19aa249ba
ARG PYTORCH_CUDA_INDEX=https://download.pytorch.org/whl/cu130

ENV PYTHONUNBUFFERED=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PIP_NO_CACHE_DIR=1 \
    VIRTUAL_ENV=/opt/venv \
    PATH=/opt/venv/bin:$PATH \
    XDG_CACHE_HOME=/data/cache \
    HF_HOME=/data/cache/huggingface \
    TORCH_HOME=/data/cache/torch \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=compute,utility \
    H3_MODEL=minimax_h3_fl2va_pruned_int8_convrot.safetensors

RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        build-essential ca-certificates curl ffmpeg git \
        libglib2.0-0 libgl1 libgomp1 \
        python3 python3-dev python3-pip python3-venv \
    && rm -rf /var/lib/apt/lists/* \
    && python3 -m venv /opt/venv \
    && python -m pip install --upgrade pip setuptools wheel

RUN python -m pip install \
        torch==2.12.1 torchvision==0.27.1 torchaudio==2.12.1 \
        --index-url "${PYTORCH_CUDA_INDEX}" \
    && python - <<'PY'
import torch
if not torch.version.cuda:
    raise SystemExit("Installed PyTorch does not include CUDA support")
print(f"PyTorch {torch.__version__}; CUDA {torch.version.cuda}")
PY

WORKDIR /opt
RUN git clone https://github.com/Comfy-Org/ComfyUI.git /opt/ComfyUI \
    && cd /opt/ComfyUI \
    && git checkout --detach "${COMFYUI_COMMIT}" \
    && python -m pip install -r requirements.txt \
    && rm -rf .git

COPY server.py /opt/h3-lab/server.py
COPY static /opt/h3-lab/static
COPY start.sh /opt/h3-lab/start.sh
RUN chmod 0755 /opt/h3-lab/start.sh \
    && mkdir -p /data/input /data/results /data/user /data/cache /opt/ComfyUI/custom_nodes

WORKDIR /opt/ComfyUI
EXPOSE 7862 8188
ENTRYPOINT ["/opt/h3-lab/start.sh"]
