FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 as builder

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# Install system dependencies as root
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    git-lfs \
    wget \
    curl \
    ca-certificates \
    build-essential \
    libsndfile1 \
    && rm -rf /var/lib/apt/lists/* \
    && git lfs install --system

# Create non-root user first
RUN useradd -m -u 1000 user

# Set environment variables
ENV CONDA_DIR=/home/user/conda
ENV PATH=$CONDA_DIR/bin:$PATH
ENV HOME=/home/user
ENV PATH=/home/user/.local/bin:$PATH

# Switch to user immediately
USER user
WORKDIR $HOME

# Install Miniconda in user's home directory
RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
    /bin/bash ~/miniconda.sh -b -p $CONDA_DIR && \
    rm ~/miniconda.sh && \
    conda clean -a -y

# Create conda environment
RUN conda create -n yue python=3.8 -y
SHELL ["conda", "run", "-n", "yue", "/bin/bash", "-c"]

# Install PyTorch with CUDA support
RUN conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -y

# Install flash-attention
RUN pip install flash-attn --no-build-isolation

# First clone repositories and create directories
WORKDIR $HOME/app

RUN git lfs install && \
    git clone --recursive https://github.com/multimodal-art-projection/YuE.git && \
    cd YuE && \
    git lfs pull && \
    cd inference && \
    git lfs checkout && \
    cd .. && \
    pip install -r requirements.txt
    
WORKDIR $HOME/app/YuE/inference
RUN git clone https://huggingface.co/m-a-p/xcodec_mini_infer

# Create output directory
RUN mkdir -p $HOME/app/output

# Copy only the necessary additional files
COPY --chown=user:user run.sh $HOME/app/
COPY --chown=user:user wrapper.py $HOME/app/YuE/inference/
RUN chmod +x $HOME/app/run.sh

# Set final working directory
WORKDIR $HOME/app

# ENTRYPOINT ["./run.sh"]