cog.yaml

build:
  gpu: true
  cuda: "12.1"
  system_packages:
    - ffmpeg
  python_version: "3.10.6"
  python_packages:
    - torch
    - torchvision
    - torchaudio
    - torchsde
    - einops
    - transformers>=4.28.1
    - tokenizers>=0.13.3
    - sentencepiece
    - safetensors>=0.3.0
    - aiohttp
    - accelerate
    - pyyaml
    - Pillow
    - scipy
    - tqdm
    - psutil
    - spandrel
    - soundfile
    - kornia>=0.7.1
    - websocket-client==1.6.3

    # fix for pydantic issues in cog
    # https://github.com/replicate/cog/issues/1623
    - albumentations==1.4.3

    - numpy>=1.26.4
    - opencv-python-headless
    - imageio-ffmpeg>=0.5.1
    - lmdb>=1.4.1
    - rich>=13.7.1
    - ultralytics
    - tyro==0.8.5
    - dill
  run:
    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.1/pget_linux_x86_64" && chmod +x /usr/local/bin/pget
predict: "predict.py:Predictor"