name: Release Build
on:
  workflow_dispatch:
permissions:
  contents: write
jobs:
  release:
    name: Tag Main Branch and Create Release
    runs-on: ubuntu-latest
    if: github.repository == 'nunchaku-tech/nunchaku'
    outputs:
      tag_name: ${{ steps.set-tag.outputs.tag_name }}
    steps:
      - name: Checkout main branch
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
          ref: main
      - name: Extract version from __version__.py
        id: version
        run: |
          version=$(grep '__version__' nunchaku/__version__.py | sed -E 's/.*"([^"]+)".*/\1/')
          echo "version=$version" >> "$GITHUB_OUTPUT"
      - name: Check if tag exists
        id: check-tag
        run: |
          tag_name="v${{ steps.version.outputs.version }}"
          if git rev-parse "$tag_name" >/dev/null 2>&1; then
            echo "Tag $tag_name already exists."
            echo "exists=true" >> "$GITHUB_OUTPUT"
          else
            echo "Tag $tag_name does not exist."
            echo "exists=false" >> "$GITHUB_OUTPUT"
          fi
      - name: Create and push tag if not exists
        if: steps.check-tag.outputs.exists == 'false'
        run: |
          tag_name="v${{ steps.version.outputs.version }}"
          git config user.name "github-actions"
          git config user.email "github-actions@users.noreply.github.com"
          git tag "$tag_name"
          git push origin "$tag_name"
      - name: Set tag_name output
        id: set-tag
        run: |
          echo "tag_name=v${{ steps.version.outputs.version }}" >> "$GITHUB_OUTPUT"
  linux-wheels:
    name: Build the linux release wheels
    runs-on: [self-hosted, linux-build]
    needs: release
    strategy:
      matrix:
        python: ["3.10", "3.11", "3.12"]
        torch: ["2.5", "2.6", "2.7", "2.8", "2.9"]
    steps:
      - name: Checkout to the tag
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
          ref: ${{ needs.release.outputs.tag_name }}
          submodules: true
      - name: Show current commit
        run: git log -1 --oneline
      - name: Build wheels
        run: |
          if [[ "${{ matrix.torch }}" == "2.7" || "${{ matrix.torch }}" == "2.8" ]]; then
            cuda_version="12.8"
          else
            cuda_version="12.4"
          fi
          if [[ "${{ matrix.torch }}" == "2.9" ]]; then
            bash scripts/build_linux_wheel_torch_nightly.sh ${{ matrix.python }} ${{ matrix.torch }} 12.8
          else
            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
          fi
      - name: Upload wheels to GitHub Release
        uses: softprops/action-gh-release@v2
        with:
          files: dist/*.whl
          name: Nunchaku ${{ needs.release.outputs.tag_name }}
          tag_name: ${{ needs.release.outputs.tag_name }}
          prerelease: false
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.12'
      - name: Upload wheels to HuggingFace Hub
        env:
          HF_TOKEN: ${{ secrets.HF_WHEEL_UPLOAD_TOKEN }}
        run: |
          pip install huggingface_hub
          for whl in dist/*.whl; do
            hf upload nunchaku-tech/nunchaku "$whl"
          done
      - name: Clean up
        if: always()
        run: bash scripts/linux_cleanup.sh
  windows-wheels:
    name: Build the windows release wheels
    runs-on: [self-hosted, windows-build]
    needs: release
    strategy:
      matrix:
        python: ["3.10", "3.11", "3.12"]
        torch: ["2.5", "2.6", "2.7", "2.8", "2.9"]
    steps:
      - name: Checkout to the tag
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
          ref: ${{ needs.release.outputs.tag_name }}
          submodules: true
      - name: Show current commit
        run: git log -1 --oneline
      - name: Build wheels
        shell: cmd
        run: |
          SET TORCH_VERSION=${{ matrix.torch }}
          SET PYTHON_VERSION=${{ matrix.python }}
          IF "%TORCH_VERSION%"=="2.7" (
            SET CUDA_VERSION=12.8
          ) ELSE IF "%TORCH_VERSION%"=="2.8" (
            SET CUDA_VERSION=12.8
          ) ELSE (
            SET CUDA_VERSION=12.4
          )
          call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
          IF "%TORCH_VERSION%"=="2.9" (
            call scripts\build_windows_wheel_torch_nightly.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
          ) ELSE (
            call scripts\build_windows_wheel.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
          )
      - name: Upload wheels to GitHub Release
        uses: softprops/action-gh-release@v2
        with:
          files: dist/*.whl
          name: Nunchaku ${{ needs.release.outputs.tag_name }}
          tag_name: ${{ needs.release.outputs.tag_name }}
          prerelease: false
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
      - name: Upload wheels to HuggingFace Hub
        shell: cmd
        env:
          HF_TOKEN: ${{ secrets.HF_WHEEL_UPLOAD_TOKEN }}
        run: |
          call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
          pip install huggingface_hub
          for %%w in (dist\*.whl) do (
            hf upload nunchaku-tech/nunchaku "%%w"
          )