cog.yaml

build:
  cuda: "11.4.3"
  gpu: true
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
#    - "libcudnn8=8.1.0.77-1+cuda11.2"
  python_version: "3.8"
  python_packages:
    - "Flask==2.1.0"
    - "Flask-Cors==3.0.9"
    - "Flask-RESTful==0.3.8"
    - "flax==0.5.0"
  run:
    - "git clone https://github.com/borisdayma/dalle-mini.git && cd dalle-mini && pip install -e . && cd .."
    - "git clone https://github.com/patil-suraj/vqgan-jax.git && cd vqgan-jax && pip install -e . && cd .."
    - "git clone https://github.com/huggingface/transformers.git && cd transformers && pip install -e . && cd .."
    - "pip install --upgrade \"jax[cuda]\" -f https://storage.googleapis.com/jax-releases/jax_releases.html"
    - "pip install wandb"
    - "wandb login --anonymously"
    - "wandb artifact get dalle-mini/dalle-mini/mini-1:v0"
    - "git config --global --add safe.directory /src"
predict: "predict.py:Predictor"
image: "r8.im/nicholascelestin/dalle-mega"