FROM python:3.12-slim

ENV PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    UV_HTTP_RETRIES=8 \
    UV_HTTP_TIMEOUT=180

WORKDIR /app

RUN apt-get update && \
    apt-get install -y --no-install-recommends libgl1 libglib2.0-0 && \
    rm -rf /var/lib/apt/lists/*

COPY core /app/core
COPY runtime /app/runtime

RUN pip install --upgrade pip setuptools wheel uv

RUN pip install ./core ./runtime

RUN printf 'pynput>=0.0.0; sys_platform == "never"\nevdev>=0.0.0; sys_platform == "never"\ntorch==2.8.0\ntorchvision==0.23.0\ntransformers==5.3.0\n' > /tmp/overrides.txt

RUN uv pip install --system \
        --override /tmp/overrides.txt \
        --extra-index-url https://download.pytorch.org/whl/cu129 \
        --index-strategy unsafe-best-match \
        torch==2.8.0 torchvision==0.23.0

RUN uv pip install --system \
        --override /tmp/overrides.txt \
        --extra-index-url https://download.pytorch.org/whl/cu129 \
        --index-strategy unsafe-best-match \
        'transformers[vision]==5.3.0'

RUN uv pip install --system \
        --override /tmp/overrides.txt \
        --extra-index-url https://download.pytorch.org/whl/cu129 \
        --index-strategy unsafe-best-match \
        'lerobot[feetech,smolvla]==0.5.1' \
        accelerate numpy pillow opencv-python-headless safetensors scipy && \
    python -c "import torch, torchvision; from transformers.models.auto.processing_auto import AutoProcessor; print(torch.__version__, torchvision.__version__)"

# Patch the installed LeRobot diffusion policy to accept heterogeneous camera resolutions
# (same patch used to train these policies). Must run after lerobot is installed and before
# the runtime imports it. Fails the build loudly on lerobot version drift.
COPY diffusion_resolution_patch.py /app/diffusion_resolution_patch.py
RUN python /app/diffusion_resolution_patch.py

# Install MolmoAct2 as a 3rd-party lerobot plugin (it only ships in lerobot `main`, not 0.5.1).
# Vendors policies/molmoact2 (+ hf_model) from a pinned lerobot commit and registers it.
# Must run after lerobot is installed; fails the build loudly on lerobot version drift.
COPY install_molmoact2_plugin.py /app/install_molmoact2_plugin.py
RUN python /app/install_molmoact2_plugin.py

COPY lerobot_eval_runtime.py /app/lerobot_eval_runtime.py

CMD ["armnet-runtime", "/app/lerobot_eval_runtime.py"]
