FROM python:3.10-slim

# Update package list and install git
RUN apt-get update && \
    apt-get install -y git && \
    rm -rf /var/lib/apt/lists/*  # Clean up to reduce image size

WORKDIR /app

# Clone the octo repository
RUN git clone https://github.com/octo-models/octo.git
WORKDIR /app/octo

# Install dependencies
RUN pip3 install -e .
RUN pip3 install -r requirements.txt
RUN pip3 install --upgrade "jax[cuda11_pip]==0.4.20" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
RUN pip3 install scipy==1.10.1

# Install FastAPI and Uvicorn for the API
RUN pip3 install fastapi uvicorn
ENV HF_HOME=/app/octo/.cache
# Copy and run the model initialization script to cache the model
COPY init_model.py /app/octo
RUN python init_model.py

# Copy the original main.py and the API app.py
COPY main.py /app/octo
COPY app.py /app/octo

# Expose port 8000 for the API
EXPOSE 7860

# Run the API with Uvicorn
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]