name: pr_test

on:
  workflow_dispatch:

jobs:
  set-up-test-env:
    runs-on: self-hosted

    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
          submodules: true

      - name: Set up Python
        run: |
          which python
          echo "Setting up Python with Conda"
          conda create -n test_env python=3.11 -y

      - name: Install dependencies
        run: |
          source $(conda info --base)/etc/profile.d/conda.sh
          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
          which python
          conda install -c conda-forge gxx=11 gcc=11
          echo "Installing dependencies"
          pip install torch torchvision torchaudio
          pip install ninja wheel diffusers transformers accelerate sentencepiece protobuf huggingface_hub
          NUNCHAKU_INSTALL_MODE=ALL python setup.py develop
          pip install -r tests/requirements.txt

  test-flux-memory:
    needs: set-up-test-env
    runs-on: self-hosted

    steps:
      - name: Run memory test
        run: |
          which python
          source $(conda info --base)/etc/profile.d/conda.sh
          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
          which python
          pytest -v -x tests/flux/test_flux_memory.py

  clean-up:
    if: always()
    needs: [set-up-test-env, test-flux-memory]
    runs-on: self-hosted

    steps:
      - name: Clean up
        run: |
          cd ..
          rm -rf *