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

WORKDIR /app

# System deps
RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*

# Pre-baked ai-toolkit (numpy 2.5 / dctorch / torchdiffeq / torchsde / clip forks, 3.14-compatible)
# Copied from local checkout so HF doesn't have to clone the submodule / run pip install -e .
COPY ai-toolkit /app/ai-toolkit

# Make ai-toolkit importable. Upstream ai-toolkit ships without setup.py / pyproject.toml,
# so pip install -e . would fail. We add it to PYTHONPATH instead.
ENV PYTHONPATH=/app:/app/ai-toolkit
ENV PYTHONUNBUFFERED=1
ENV HF_HOME=/app/hf_cache
ENV TRANSFORMERS_CACHE=/app/hf_cache

# Install runtime deps
RUN pip install --no-cache-dir huggingface_hub hf_transfer

# Copy training files
COPY . /app/

# Pre-download FLUX base + training adapter at build time so they're in the image cache
RUN python -c "import os; os.environ['HF_HUB_ENABLE_HF_TRANSFER']='1'; from huggingface_hub import snapshot_download; snapshot_download('Niansuh/FLUX.1-schnell'); snapshot_download('ostris/FLUX.1-schnell-training-adapter')"

CMD ["python", "/app/train_cloud.py"]