# ROMA inference image for Quadro RTX 6000 (Turing sm_75, x86_64, 24GB ×4 = 96GB).
#
# Base: CUDA 12.4 runtime on Ubuntu 22.04. CUDA 12.4 matches the repo's torch 2.6.0+cu124, and runs
# fine on the host's 570.x / CUDA-12.8 driver (drivers are backward compatible). We install the
# OFFICIAL PyTorch cu124 wheels (via download.pytorch.org), which include Turing (sm_75) kernels --
# NGC datacenter images may omit sm_75 and produce "no kernel image is available" at runtime, so we
# deliberately do NOT use an NGC base here.
#
# Turing notes (handled in the demo scripts, not here): FlashAttention-2 is unsupported on sm_75 so
# the demos default to attn_implementation=sdpa; bf16 is not accelerated so they default to fp16.
ARG BASE_IMAGE=nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
FROM ${BASE_IMAGE}

ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_INDEX=https://pypi.org/simple
ENV TORCH_INDEX=https://download.pytorch.org/whl/cu124
ENV HF_HUB_ENABLE_HF_TRANSFER=1
# The gradio demos call demo.launch() without a server name; bind to all interfaces so the UI is
# reachable from the host (the demos read this env, added in the gradio edits).
ENV GRADIO_SERVER_NAME=0.0.0.0
ENV PYTHONUNBUFFERED=1

ARG TRANSFORMERS_REF=roma_patch

# System deps: Python 3.10, ffmpeg (decord/av/moviepy), git (to pip-install the transformers fork).
RUN apt-get update && apt-get install -y --no-install-recommends \
        python3.10 python3.10-dev python3-pip \
        ffmpeg git ca-certificates && \
    ln -sf /usr/bin/python3.10 /usr/bin/python && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Install the runtime deps. The +cu124 torch/vision/audio pins resolve from the PyTorch index; the
# transformers@roma_patch fork is fetched from git (already pinned in the requirements file).
COPY requirements-rtx6000.txt /app/requirements-rtx6000.txt
RUN python -m pip install --upgrade pip && \
    python -m pip install \
        --index-url "$PIP_INDEX" \
        --extra-index-url "$TORCH_INDEX" \
        -r /app/requirements-rtx6000.txt

# Register the llamafactory package editable from the local checkout (deps already satisfied above).
COPY . /app
RUN python -m pip install --no-deps -e .

# Build-time sanity import (GPU is not visible during build, so we don't assert cuda here).
RUN python -c "import torch, transformers, bitsandbytes; print('torch', torch.__version__, '| transformers', transformers.__version__)"

ENV GRADIO_SERVER_PORT=7860
EXPOSE 7860
ENV API_PORT=8000
EXPOSE 8000

# Persist the checkpoint, HF cache and demo media across container runs.
VOLUME [ "/app/whole_model", "/root/.cache/huggingface", "/app/demo_media" ]

CMD ["bash"]