FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime

ENV DEBIAN_FRONTEND=noninteractive

# GDAL required by rasterio and geopandas (terratorch deps)
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgdal-dev gdal-bin git build-essential \
    && rm -rf /var/lib/apt/lists/*

# torchvision 0.19.0 matches torch 2.4.0 in the base image
RUN pip install --no-cache-dir \
    "torchvision==0.19.0" \
    --extra-index-url https://download.pytorch.org/whl/cu124

# GDAL Python binding must match system GDAL version
RUN pip install --no-cache-dir "GDAL==$(gdal-config --version)"

# terratorch 1.2.7 with pinned deps:
# - torchgeo>=0.7.2: Python 3.11 image allows it; 0.7.2 has SENTINEL2_ALL_SOFTCON
#   which terratorch's torchgeo_resnet.py requires at import time
# - diffusers==0.30.0: stable API used by terramind tokenizer scheduler
# - transformers<5: terratorch 1.2.x built against 4.x API
RUN pip install --no-cache-dir \
    "torchgeo==0.7.2" \
    "terratorch==1.2.7" \
    "diffusers==0.30.0" \
    "transformers>=4.40.0,<5.0" \
    "huggingface_hub>=0.20.0"

WORKDIR /app
COPY handler.py /app/handler.py
COPY prithvi_mae.py /app/prithvi_mae.py
COPY config.json /app/config.json