# Use an official PyTorch image with CUDA support as the base image
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime

# Set the working directory
WORKDIR /app

# Copy the requirements files
COPY requirements.txt /app/
# Copy the FastAPI app code
COPY main.py /app/
RUN pip install --no-cache-dir -r requirements.txt

# Install dependencies
RUN apt-get update && apt-get install -y \
    libgl1-mesa-glx \
    libegl1-mesa \
    libxrandr2 \
    libxrandr2 \
    libxss1 \
    libxcursor1 \
    libxcomposite1 \
    libasound2 \
    libxi6 \
    libxtst6 \
    libglib2.0-0 \
    git \
    wget \
    && rm -rf /var/lib/apt/lists/*

# Clone the dust3r repository
RUN git clone --recursive https://github.com/tur-learning/dust3r /app/dust3r
WORKDIR /app/dust3r
RUN pip install --no-cache-dir -r requirements.txt
# RUN pip install --no-cache-dir -r requirements_optional.txt
RUN pip install opencv-python==4.8.0.74

# Download the model weights
RUN mkdir -p /app/dust3r/checkpoints && \
    wget https://huggingface.co/camenduru/dust3r/resolve/main/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth -P /app/dust3r/checkpoints

# Expose the port FastAPI is running on
EXPOSE 7860

WORKDIR /app
# Command to run the FastAPI app
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]