#!/bin/bash

set -e

. ./venv/bin/activate

git clone -b main --single-branch https://github.com/facebookresearch/audiocraft /opt/audiocraft
cd /opt/audiocraft
git checkout 72cb16f9fb239e9cf03f7bd997198c7d7a67a01c

# pinned requirements
cat > requirements.txt <<EOF
av==11.0.0
demucs==4.0.1
einops==0.7.0
encodec==0.1.1
flashy==0.0.2
gradio==3.48.0
gradio_client==0.6.1
hydra-core==1.3.2
julius==0.2.7
librosa==0.10.1
num2words==0.5.12
numpy==1.25.2
protobuf==4.24.4
sentencepiece==0.1.99
spacy==3.6.1
spacy-legacy==3.0.12
spacy-loggers==1.0.5
torch==2.1.0+cu121
torchaudio==2.1.0+cu121
torchmetrics==1.2.0
tqdm==4.66.1
transformers==4.34.0
xformers<0.0.23
torchvision==0.16.0
torchtext==0.16.0
EOF

pip install -e . --no-cache-dir

pip install git+https://github.com/CPJKU/madmom@0551aa8f48d71a367d92b5d3a347a0cf7cd97cc9#egg=madmom

export TORCH_HOME=/opt/audiocraft-models
export HF_HOME=/opt/audiocraft-models
python <<EOF
from audiocraft.models import MusicGen, AudioGen
MusicGen.get_pretrained("facebook/musicgen-large")
MusicGen.get_pretrained("facebook/musicgen-stereo-large")
EOF