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 \
    && 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 dependencies
RUN pip3 install --no-cache-dir numpy && \
    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 && \
    pip3 install --no-cache-dir colorlog && \
    pip3 install --no-cache-dir -r requirements.txt && \
    pip3 install -e .

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

# Rename gradio_demo.py to app.py
RUN mv gradio_demo.py 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"]