# =============================================================================
# HY-Motion Worker Image — ghcr.io/poly-hammer/ph-hy-motion
#
# Universal image for both self-hosted (Hatchet) and Modal workers.
# Weights are volume-mounted at runtime, NOT baked in.
#
# Self-hosted usage:
#   docker volume create hy-motion-weights
#   docker run --gpus all -v hy-motion-weights:/weights ph-hy-motion \
#       ph-worker download-weights --job hy-motion --all
#   docker run --gpus all -v hy-motion-weights:/weights ph-hy-motion \
#       ph-worker start --token <WORKER_TOKEN>
#
# Modal usage:
#   modal.Image.from_registry("ghcr.io/poly-hammer/ph-hy-motion:latest", ...)
#
# Build (from workers/self-hosted/ context):
#   docker compose build hy-motion
#
# Multi-CUDA build:
#   docker compose build hy-motion-cu121
#
# See workers/self-hosted/README.md for license notices.
# =============================================================================

# ---------- CUDA / PyTorch build arguments ----------
# Override these to build for different CUDA generations.
# Default: cu128 (CUDA 13.1 — Blackwell sm_120+, Hopper, Ampere)
ARG CUDA_TAG=13.1.1-cudnn-devel-ubuntu24.04
ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
ARG TORCH_VERSION=2.10.0
ARG TORCHVISION_VERSION=0.25.0

FROM nvidia/cuda:${CUDA_TAG}

LABEL org.opencontainers.image.source="https://github.com/poly-hammer/poly-hammer-portal"

ENV DEBIAN_FRONTEND=noninteractive

# ---------- System packages ----------
RUN apt-get update && apt-get install -y --no-install-recommends \
        python3 python3-venv python3-dev python3-pip \
        git git-lfs curl ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# ---------- Python venv (PEP 668 — Ubuntu 24.04 blocks system-wide pip) ----------
RUN python3 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# ---------- Python packages ----------
RUN pip install --no-cache-dir --upgrade pip setuptools wheel

# PyTorch — version and wheel index controlled by build args
ARG TORCH_INDEX_URL
ARG TORCH_VERSION
ARG TORCHVISION_VERSION
RUN python -m pip install --no-cache-dir \
    torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} \
    --extra-index-url ${TORCH_INDEX_URL}

# HY-Motion inference dependencies (CUDA-agnostic — shared across all variants)
COPY poly_hammer_worker/jobs/hy_motion/requirements.txt /tmp/requirements.txt
RUN python -m pip install --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt

# ---------- Install poly-hammer-worker (with jobs package) ----------
COPY . /opt/worker-src
RUN pip install --no-cache-dir /opt/worker-src \
    && python /opt/worker-src/scripts/patch_hatchet_logger.py \
    && rm -rf /opt/worker-src

# ---------- Clone HY-Motion repo ----------
ENV HY_MOTION_REPO_PATH=/opt/hy-motion
RUN git clone --depth 1 https://github.com/Tencent-Hunyuan/HY-Motion-1.0.git ${HY_MOTION_REPO_PATH}

# ---------- Runtime configuration ----------
ENV WEIGHTS_PATH=/weights
ENV HF_HUB_ENABLE_HF_TRANSFER=1
# High-performance HF Xet transfer backend — large weight files download
# noticeably faster on first run.
ENV HF_XET_HIGH_PERFORMANCE=1
ENV USE_HF_MODELS=0
ENV HF_HOME="${WEIGHTS_PATH}/huggingface"
ENV TRANSFORMERS_CACHE="${WEIGHTS_PATH}/huggingface"

ENV PH_JOB_CLASS="poly_hammer_worker.jobs.hy_motion.job.HyMotionAllJob"
ENV PYTHONPATH="${HY_MOTION_REPO_PATH}"

VOLUME /weights

# No ENTRYPOINT — Modal needs to run `python` directly inside the container.
# Self-hosted usage still works: `docker run ... ph-worker start --token X`
CMD ["ph-worker", "start"]
