ml-models/.github/workflows/export.yaml

on:
  workflow_call:
    secrets:
      HF_TOKEN:
        required: true
    inputs:
      model-name:
        required: true
        type: string
      model-source:
        required: true
        type: string
      runner:
        required: true
        type: string
      hf-name:
        required: false
        type: string
      environment:
        type: string
        default: ''
      compile-all:
        type: boolean
        default: true
      previous-ref:
        type: string
        default: ''

jobs:
  export:
    runs-on: ${{ inputs.runner }}
    outputs:
      changed: ${{ steps.needed.outputs.compile }}
    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          fetch-depth: 0 # the compile scope below diffs this model's renderings against the previous ref

      - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
        with:
          enable-cache: true
          cache-dependency-glob: uv.lock

      - run: uv run --frozen --extra export immich-model onnx export "$MODEL_NAME" "$MODEL_SOURCE" --hf-model-name "$HF_MODEL_NAME"
        env:
          MODEL_NAME: ${{ inputs.model-name }}
          MODEL_SOURCE: ${{ inputs.model-source }}
          HF_MODEL_NAME: ${{ inputs.hf-name || inputs.model-name }}
          HF_TOKEN: ${{ secrets.HF_TOKEN }}

      # Here rather than in its own job because this is where an exported model exists.
      - id: check
        run: uv run --frozen --extra export immich-model check model models

      # An exporter change stales every rendering at once; hand back this run's rather than making the
      # author re-export the catalog

      - if: ${{ failure() && steps.check.outcome == 'failure' }}
        run: uv run --frozen --extra export immich-model check model models --bless
      - if: ${{ failure() && steps.check.outcome == 'failure' }}
        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
        with:
          name: renderings-${{ inputs.model-name }}
          path: |
            ci/graphs/${{ inputs.model-name }}
            ci/rewrites/${{ inputs.model-name }}

      # Skip compilation if the renderings are the same, being what the binary is built from
      - id: needed
        env:
          MODEL_NAME: ${{ inputs.model-name }}
          COMPILE_ALL: ${{ inputs.compile-all }}
          PREVIOUS_REF: ${{ inputs.previous-ref }}
        run: |
          git add -N ci # a model exported for the first time renders to a path git is not yet tracking
          if [ "$COMPILE_ALL" = 'true' ] || ! git diff --quiet "$PREVIOUS_REF" -- ci/plans.txt \
              "ci/graphs/$MODEL_NAME" "ci/rewrites/$MODEL_NAME"; then
            echo "compile=yes" >> "$GITHUB_OUTPUT"
          fi

      # RKNN binaries are built from the same fused ONNX every other backend uses
      - if: ${{ steps.needed.outputs.compile }}
        run: uv run --frozen --extra rknn immich-model rknn compile "$MODEL_NAME"
        env:
          MODEL_NAME: ${{ inputs.model-name }}

      - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
        env:
          ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY: 32
        with:
          name: ${{ inputs.model-name }}
          path: models/
          compression-level: 0

  # When doing benchmarking, that goes here
  # Download the exported artifact, and upload an artifact with the benchmark result to be used below

  upload:
    if: ${{ inputs.environment && needs.export.outputs.changed }}
    runs-on: ${{ inputs.runner == 'ubuntu-latest' && 'ubuntu-latest' || 'pokedex-general' }}
    needs: export
    environment: ${{ inputs.environment }}
    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1

      - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
        with:
          enable-cache: true
          cache-dependency-glob: uv.lock
      - run: uv sync --frozen --extra export

      - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
        with:
          name: ${{ inputs.model-name }}
          merge-multiple: 'true' # We don't actually have multiple, but this also sets the subfolder naming behaviour
          path: models/

      - run: |
          uv run --frozen --extra export immich-model upload "$MODEL_NAME" \
            --hf-organization "$ORGANIZATION" --hf-branch v2
        env:
          MODEL_NAME: ${{ inputs.model-name }}
          ORGANIZATION: ${{ vars.HF_ORGANIZATION }}
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
          HF_ENDPOINT: https://huggingface.co