# Uses the exact same verified environment as your Dev Container
FROM nvcr.io/nvidia/pytorch:25.02-py3

# Install API and Flux-specific dependencies
# Note: We skip the torch versions in your list because the NVIDIA base 
# already includes the optimized sm_120 version we need.
RUN pip3 install --no-cache-dir \
    fastapi \
    uvicorn \
    pydantic \
    transformers==4.46.1 \
    diffusers==0.35.2 \
    accelerate==0.30.1 \
    sentencepiece \
    protobuf

# Set up the application
WORKDIR /app
COPY api.py /app/api.py

# Ensure GPU visibility (matching your devcontainer.json)
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

CMD ["python3", "api.py"]