FROM nvcr.io/nvidia/pytorch:23.04-py3

    ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && \
apt-get install -y --no-install-recommends software-properties-common && \
add-apt-repository ppa:deadsnakes/ppa -y && \
apt-get update && \
apt-get install -y --no-install-recommends \
    python3.10 \
    python3.10-distutils \
    python3.10-dev && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10 && \
apt-get clean && rm -rf /var/lib/apt/lists/*

# Update GPG key of NVIDIA Docker Images 
# (See https://developer.nvidia.com/blog/updating-the-cuda-linux-gpg-repository-key/ for more detail)
RUN rm -f /etc/apt/sources.list.d/cuda.list \
 && apt-get update && apt-get install -y --no-install-recommends \
    wget \
 && wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-keyring_1.0-1_all.deb \
 && dpkg -i cuda-keyring_1.0-1_all.deb \
 && rm -f cuda-keyring_1.0-1_all.deb


RUN apt-get clean && \
    apt-get -y update && \
    apt-get install -y --no-install-recommends \
    nano curl git zip htop unzip tmux ca-certificates sudo bzip2 libx11-6 lua5.3 \
    build-essential vim gcc g++ make openssl ffmpeg libssl-dev libbz2-dev libboost-all-dev \
    libreadline-dev libsqlite3-dev python3-tk tk-dev python-tk libfreetype6-dev \
    libffi-dev libaio-dev openmpi-bin openmpi-doc libopenmpi-dev liblzma-dev libncurses-dev libsndfile1 \
    # add basic apt packages
    && apt-get clean \ 
    && rm -rf /var/lib/apt/lists/*


RUN apt update && apt install -y --no-install-recommends apt-transport-https gnupg && \
    curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor >bazel-archive-keyring.gpg && \
    mv bazel-archive-keyring.gpg /usr/share/keyrings && \
    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/bazel-archive-keyring.gpg] https://storage.googleapis.com/bazel-apt stable jdk1.8" > /etc/apt/sources.list.d/bazel.list && \
    apt update && apt install -y --no-install-recommends bazel && \
    apt install -y --no-install-recommends bazel-5.3.2 \
    && apt-get clean \ 
    && rm -rf /var/lib/apt/lists/*

RUN apt-get update & \
    python -m pip install --upgrade pip
# RUN pip install git+https://github.com/facebookresearch/audiocraft.git
RUN pip install alias-free-torch==0.0.6 auraloss==0.4.0 descript-audio-codec einops einops-exts==0.0.4 ema-pytorch==0.2.3
RUN pip install gradio huggingface_hub importlib-resources==5.12.0 mpi4py opencv-python timm==0.9.12 gdown networkx
RUN pip install pandas==2.0.2 pedalboard==0.7.4 prefigure==0.0.9 pytorch_lightning PyWavelets safetensors omegaconf moviepy annoy matplotlib
RUN pip install sentencepiece==0.1.99 s3fs torchmetrics==0.11.4 tqdm nitrous-ema fvcore iopath submitit scikit-image h5py
RUN pip install wandb==0.15.4 webdataset==0.2.48 x-transformers diffusers ninja av ffmpeg-python yt-dlp seaborn prettytable hdf5plugin
RUN pip install librosa==0.9.2 scikit_learn==1.0.2 scipy==1.7.3 soundfile==0.11.0 termcolor==2.1.0 hydra-colorlog hydra-core openpyxl==3.1.2
RUN pip install pytest==7.4.3 pylint==3.0.2 pyloudnorm==0.1.1 pandarallel==1.6.5 julius==0.2.7 mutagen==1.47.0 GitPython==3.1.40 boto3==1.34.60 rich==13.7.1 dictdiffer==0.9.0 torchopenl3==1.0.1
RUN pip install laion-clap python-dotenv tensordict==0.6.0 vllm deepspeed accelerate 
RUN pip install flash-attn==2.5.1.post1
RUN pip install numpy==1.23.5
RUN pip uninstall torch -y
RUN pip install --extra-index-url https://download.pytorch.org/whl/cu121 \
    torch==2.5.1+cu121 torchvision==0.20.1+cu121 torchaudio==2.5.1+cu121
RUN pip install xformers==0.0.29 
RUN pip install passt@git+https://github.com/hkchengrex/passt_hear21.git
RUN pip install imagebind@git+https://github.com/hkchengrex/ImageBind.git
RUN pip install mauve-text decord
# RUN pip install \
#     --extra-index-url=https://pypi.nvidia.com \
#     "cudf-cu12==25.4.*" "dask-cudf-cu12==25.4.*" "cuml-cu12==25.4.*" \
#     "cugraph-cu12==25.4.*" "nx-cugraph-cu12==25.4.*" "cuspatial-cu12==25.4.*" \
#     "cuproj-cu12==25.4.*" "cuxfilter-cu12==25.4.*" "cucim-cu12==25.4.*" \
#     "pylibraft-cu12==25.4.*" "raft-dask-cu12==25.4.*" "cuvs-cu12==25.4.*" "nx-cugraph-cu12==25.4.*"
# RUN pip install numpy==1.23.5
# RUN git clone https://github.com/yukara-ikemiya/friendly-stable-audio-tools.git
# RUN cd friendly-stable-audio-tools && \
#     pip install .
# RUN pip install .
# # To avoid Accelerate import error
RUN pip uninstall -y transformer-engine
RUN pip uninstall -y apex