action.yml

name: Setup Cog
description: Install Cog so you can run, test, and push machine learning models to Replicate.
author: replicate
branding:
  icon: settings
  color: gray-dark

inputs:
  token:
    required: true
    description: "If present, this token is used to authenticate to Replicate using `cog login` so you can push your model."
  cog-version:
    description: "The version of Cog to install"
    required: false
    default: "v0.9.8"
  install-cuda:
    description: "Install NVIDIA CUDA Toolkit to build GPU-bound containers."
    required: false
    default: 'false'
    options:
      - true
      - false

runs:
  using: "composite"
  steps:
      - name: Set up Docker buildx
        uses: docker/setup-buildx-action@d70bba72b1f3fd22344832f00baa16ece964efeb # v3.3.0
        with:
          driver: "docker"

      - name: Install NVIDIA CUDA Toolkit
        if: inputs.install-cuda == true || inputs.install-cuda == 'true'
        uses: Jimver/cuda-toolkit@dc0ca7bb29c5a92f7a963d3d5c93f8d59765136a # v0.2.14

      - name: Setup Cog
        shell: bash
        run: |
          sudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/download/${{ inputs.cog-version }}/cog_Linux_x86_64
          sudo chmod +x /usr/local/bin/cog
          cog --version

      - name: Log in to Replicate
        shell: bash
        if: inputs.token != ''
        env:
          REPLICATE_CLI_AUTH_TOKEN: ${{ inputs.token }}
        run: |
          echo "Logging in to Replicate..." >&2
          echo "$REPLICATE_CLI_AUTH_TOKEN" | cog login --token-stdin