cog.yaml

build:
  gpu: true
  cuda: "12.4"
  python_version: "3.10"
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
    - "ffmpeg"
  python_packages:
    - "torch==2.4.0"
    - "torchvision"
    - "transformers==4.42.3"
    - "opencv-python-headless<4.10"
    - "peft<0.14.0"
    - "timm==1.0.9"
    - "einops==0.8.0"
    - "sentencepiece==0.2.0"
    - "mmengine<1"
    - "accelerate"
    - "numpy<2"
  
  run:
    - FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation
    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64" && chmod +x /usr/local/bin/pget

predict: "predict.py:Predictor"