FROM python:3.9-slim

WORKDIR /app

# Install dependencies
RUN apt-get update && apt-get install -y \
    build-essential \
    python3-dev \
    libfreetype6-dev \
    pkg-config \
    libpng-dev \
    cmake \
    git \
    && rm -rf /var/lib/apt/lists/*

# Clone diffvg repository
RUN git clone https://github.com/BachiLi/diffvg.git && \
    cd diffvg && \
    git submodule update --init --recursive && \
    python setup.py install

# Install PyTorch and other dependencies
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
RUN pip install transformers pillow svgwrite svgpathtools numpy tqdm scikit-image matplotlib
RUN pip install fastapi uvicorn

# Copy model files
COPY . /app/model

# Copy handler files
COPY handler_template.py /app/
COPY diffsketcher_handler.py /app/
COPY api.py /app/

# Set environment variables
ENV MODEL_ID="jree423/diffsketcher"
ENV TASK="text-to-image"

# Run the API
CMD ["python", "api.py"]