# Use official PyTorch image with CUDA 12.2 for GPU support
FROM pytorch/pytorch:2.0.1-cuda12.2-cudnn8-runtime

# Set working directory
WORKDIR /app

# Copy handler and code
COPY handler.py ./ 
COPY wan ./wan  # your custom code folder if any

# Copy requirements
COPY requirements.txt ./requirements.txt

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

# Upgrade pip
RUN pip install --upgrade pip

RUN pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

# Install python dependencies
RUN pip install -r requirements.txt

# HF inference toolkit (if needed, you can install or add it manually)
RUN pip install huggingface_hub

# Expose port (if needed, usually 8080)
EXPOSE 8080

# Command to run your inference server
CMD ["python", "handler.py"]