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 \
    && 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
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10

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

# Set working directory to MMAudio
WORKDIR /code/MMAudio

# Install base dependencies first
RUN pip3 install --no-cache-dir \
    setuptools \
    wheel \
    numpy==1.24.3 \
    colorlog \
    typing_extensions

# Install PyTorch and related packages
RUN pip3 install --no-cache-dir \
    torch==2.1.2 \
    torchvision==0.16.2 \
    torchaudio==2.1.2 \
    --index-url https://download.pytorch.org/whl/cu118

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

# Install MMAudio
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"]