FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04

WORKDIR /code

# Install system dependencies
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3.10-distutils \
    python3-pip \
    git \
    ffmpeg \
    libsm6 \
    libxext6 \
    curl \
    libsndfile1 \
    && rm -rf /var/lib/apt/lists/*

# Ensure we're using Python 3.10
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1

# Install pip for Python 3.10 and upgrade it
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10 && \
    pip3 install --no-cache-dir --upgrade pip setuptools wheel

# Install numpy first, before anything else
RUN pip3 install --no-cache-dir numpy==1.24.3

# Install PyTorch and related packages first (as recommended in README)
RUN pip3 install --no-cache-dir \
    torch==2.6.0 \
    torchvision==0.21.0 \
    torchaudio==2.6.0 \
    --index-url https://download.pytorch.org/whl/cu118

# Install core dependencies
RUN pip3 install --no-cache-dir \
    colorlog==6.8.2 \
    torchdiffeq==0.2.3 \
    omegaconf>=2.3.0 \
    pandas>=2.0.0 \
    tensordict>=0.2.0 \
    hydra-core>=1.3.2 \
    tqdm>=4.65.0 \
    librosa>=0.10.1 \
    timm>=0.9.12 \
    requests>=2.31.0

# Clone MMAudio
RUN git clone https://github.com/hkchengrex/MMAudio.git

# Set working directory to MMAudio
WORKDIR /code/MMAudio

# Install other dependencies
RUN pip3 install --no-cache-dir -r requirements.txt

# Install MMAudio last (as recommended in README)
RUN pip3 install -e .

# Create output directory
RUN mkdir -p output/gradio && chmod 777 output/gradio

# Copy app.py (we'll use our own version instead of the one from the repo)
COPY app.py .

# Set environment variables for Hugging Face Spaces
ENV PYTHONUNBUFFERED=1
ENV GRADIO_SERVER_NAME=0.0.0.0
ENV GRADIO_SERVER_PORT=7860
ENV PYTHONPATH=/code/MMAudio

# Expose Gradio port
EXPOSE 7860

# Run the app
CMD ["python3", "app.py"]