FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu20.04

WORKDIR /app

# Set environment variables
ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive \
    TRANSFORMERS_CACHE=/app/.cache/transformers \
    HF_HOME=/app/.cache/huggingface \
    TORCH_HOME=/app/.cache/torch \
    HF_DATASETS_CACHE=/app/.cache/datasets

# Install basic dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.8 \
    python3.8-dev \
    python3-pip \
    python3-setuptools \
    git \
    wget \
    ca-certificates \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Create symbolic link for python
RUN ln -sf /usr/bin/python3.8 /usr/bin/python

# Upgrade pip
RUN pip install --no-cache-dir --upgrade pip

# Create cache directories
RUN mkdir -p /app/.cache/transformers \
    /app/.cache/huggingface \
    /app/.cache/torch \
    /app/.cache/datasets

# Clone LAVIS repository to temp directory for installation
RUN git clone https://github.com/salesforce/LAVIS.git /tmp/LAVIS \
    && cd /tmp/LAVIS \
    && sed -i '/open3d/d' requirements.txt \
    && pip install --no-cache-dir -e . \
    && cd / \
    && cp -r /tmp/LAVIS/lavis /app/LAVIS/ \
    && rm -rf /tmp/LAVIS

# Install PyTorch with CUDA support
RUN pip install --no-cache-dir torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1 --extra-index-url https://download.pytorch.org/whl/cu118

# Copy requirements files and install dependencies
COPY requirements_lavis.txt requirements_emo.txt ./
RUN pip install --no-cache-dir -r requirements_lavis.txt -r requirements_emo.txt

# Copy model and application files
COPY app.py blip2_vicuna_instruct.py ./
COPY static/ ./static/
COPY templates/ ./templates/
COPY LAVIS/ ./LAVIS/

# Create directory for model weights (to be mounted or downloaded at runtime)
RUN mkdir -p ./LAVIS/lavis/weight/vicuna-7b-2/

# Set up a volume for persistent cache
VOLUME /app/.cache

# Set the default command to run the Flask app
CMD ["python", "app.py"]