FROM python:3.10-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \
    build-essential \
    git \
    && rm -rf /var/lib/apt/lists/*

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

# Install FastAPI and other dependencies
RUN pip install --no-cache-dir fastapi==0.104.1 \
    uvicorn==0.23.2 \
    python-multipart==0.0.6 \
    pillow==10.0.1 \
    safetensors==0.4.0 \
    transformers==4.34.0 \
    diffusers==0.23.0 \
    accelerate==0.23.0

# Copy application code
COPY . .

# Create model directory
RUN mkdir -p models

# Expose the port the app will run on
EXPOSE 7860

# Start the FastAPI application
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]