FROM pytorch/pytorch:2.2.2-cuda12.1-cudnn8-devel

WORKDIR /app

# Common system dependencies
RUN apt-get update && apt-get install -y \
    libportaudio2 \
    libportaudiocpp0 \
    portaudio19-dev \
    libasound-dev \
    libsndfile1-dev \
    kmod \
    && rm -rf /var/lib/apt/lists/*

# Conditional CUDA configuration for AMD64 only
ARG TARGETARCH
RUN if [ "$TARGETARCH" = "amd64" ]; then \
    apt-get update && apt-get install -y \
    ocl-icd-opencl-dev \
    libopenblas-dev \
    && rm -rf /var/lib/apt/lists/*; \
    fi

# Mac-specific audio fixes
RUN if [ "$TARGETARCH" = "arm64" ]; then \
    apt-get update && apt-get install -y \
    libomp5 \
    libopenblas0-pthread \
    && rm -rf /var/lib/apt/lists/*; \
    fi

# Update pip first
RUN pip install --upgrade pip setuptools wheel

# Install base requirements first
RUN pip install numpy==1.26.4

# Install appropriate PyTorch and TorchAudio versions
RUN if [ "$TARGETARCH" = "amd64" ]; then \
    pip install torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu121; \
    else \
    pip install torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cpu; \
    fi

# Install remaining requirements
COPY . .
RUN pip install --no-cache-dir -r requirements.txt

EXPOSE 7860
CMD ["python", "app.py"]