name: Diffusion CI Ground Truth Generation (NPU)

on:
  workflow_dispatch:
    inputs:
      ref:
        description: 'Git ref to checkout'
        required: false
        default: ''
        type: string
      case_ids:
        description: 'Specific case IDs to run (space-separated, optional)'
        required: false
        default: ''
        type: string
      output_name:
        description: 'Custom local output/artifact folder name. Leave empty to use defaults.'
        required: false
        default: ''
        type: string
      publish_target_dir:
        description: 'Remote target directory in sgl-project/ci-data-diffusion. Leave empty to use diffusion-ci/consistency_gt/sglang_generated/ascend.'
        required: false
        default: ''
        type: string


concurrency:
  group: diffusion-ci-gt-gen-${{ github.ref }}-${{ inputs.output_name || inputs.case_ids || 'default' }}
  cancel-in-progress: true

permissions:
  contents: write
  actions: read

env:
  SGLANG_IS_IN_CI: true
  OUTPUT_NAME: ${{ inputs.output_name || 'diffusion-ci-outputs' }}
  PUBLISH_TARGET_DIR: ${{ inputs.publish_target_dir || 'diffusion-ci/consistency_gt/sglang_generated/ascend' }}

jobs:

  compute-diffusion-partitions:
    if: github.repository == 'sgl-project/sglang'
    runs-on: ubuntu-latest
    outputs:
      matrix-1npu: ${{ steps.compute.outputs.matrix-1npu }}
      matrix-2npu: ${{ steps.compute.outputs.matrix-2npu }}
      partition-count-1npu: ${{ steps.compute.outputs['partition-count-1npu'] }}
      partition-count-2npu: ${{ steps.compute.outputs['partition-count-2npu'] }}
      plan-1npu: ${{ steps.compute.outputs.plan-1npu }}
      plan-2npu: ${{ steps.compute.outputs.plan-2npu }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        with:
          ref: ${{ inputs.ref || github.ref }}

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.10'

      - name: Verify write access to sgl-project/ci-data-diffusion
        env:
          GH_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }}
        run: |
          if [ "$(gh api repos/sgl-project/ci-data-diffusion --jq .permissions.push)" != "true" ]; then
            echo "::error::GH_PAT_FOR_NIGHTLY_CI_DATA lacks write (push) access to sgl-project/ci-data-diffusion. Grant it before running GT generation."
            exit 1
          fi
          echo "OK: write access to sgl-project/ci-data-diffusion confirmed"

      - name: Compute partitions
        env:
          USE_NPU_CONFIGS: true
        id: compute
        run: |
          python scripts/ci/utils/diffusion/compute_diffusion_partitions.py \
            --min-time 1200 \
            --target-time 1800 \
            --max-time 2400 \
            --max-partitions 10 \
            --parametrized-only

  multimodal-diffusion-gen-1npu:
    needs: compute-diffusion-partitions
    if: |
      needs.compute-diffusion-partitions.result == 'success' &&
      needs.compute-diffusion-partitions.outputs.matrix-1npu != '{"include":[]}'
    runs-on: linux-aarch64-a3-800t-2
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.compute-diffusion-partitions.outputs.matrix-1npu) }}
    timeout-minutes: 150
    container:
      image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0-a3-ubuntu22.04-py3.11
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        with:
          ref: ${{ inputs.ref || github.ref }}

      - name: Mark repository safe
        run: |
          git config --system --add safe.directory ${GITHUB_WORKSPACE}

      - name: Install dependencies
        env:
          TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu"
          PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple"
          UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple"
          GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/"
          RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082"
        run: |
          # speed up by using infra cache services
          CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
          sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
          pip config set global.index-url http://${CACHING_URL}/pypi/simple
          pip config set global.trusted-host "${CACHING_URL}"

          bash scripts/ci/npu/npu_ci_install_dependency.sh a3 diffusion

      - name: Generate outputs
        env:
          PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-1npu }}
          SGLANG_USE_MODELSCOPE: true
          SGLANG_IS_IN_CI: true
          HF_ENDPOINT: https://hf-mirror.com
          TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
          PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
          STREAMS_PER_DEVICE: 32
        run: |
          cd python
          python -m sglang.multimodal_gen.test.scripts.gen_diffusion_ci_outputs \
            --suite 1-npu \
            --partition-id ${{ matrix.part }} \
            --total-partitions ${{ needs.compute-diffusion-partitions.outputs['partition-count-1npu'] }} \
            --partition-plan-json "$PARTITION_PLAN_JSON" \
            --out-dir ./${{ env.OUTPUT_NAME }} \
            ${{ inputs.case_ids != '' && format('--case-ids {0}', inputs.case_ids) || '' }}

      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
          name: ${{ env.OUTPUT_NAME }}-1npu-part${{ matrix.part }}
          path: python/${{ env.OUTPUT_NAME }}
          retention-days: 7

      - name: Publish GT images to sgl-project/ci-data-diffusion
        env:
          GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }}
        run: |
          python scripts/ci/utils/diffusion/publish_diffusion_gt.py \
            --source-dir python/${{ env.OUTPUT_NAME }} \
            --target-dir "${{ env.PUBLISH_TARGET_DIR }}"

  multimodal-diffusion-gen-2npu:
    needs: compute-diffusion-partitions
    if: |
      needs.compute-diffusion-partitions.result == 'success' &&
      needs.compute-diffusion-partitions.outputs.matrix-2npu != '{"include":[]}'
    runs-on: linux-aarch64-a3-800t-16
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.compute-diffusion-partitions.outputs.matrix-2npu) }}
    timeout-minutes: 150
    container:
      image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0-a3-ubuntu22.04-py3.11
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        with:
          ref: ${{ inputs.ref || github.ref }}

      - name: Mark repository safe
        run: |
          git config --system --add safe.directory ${GITHUB_WORKSPACE}

      - name: Install dependencies
        env:
          TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu"
          PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple"
          UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple"
          GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/"
          RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082"
        run: |
          # speed up by using infra cache services
          CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
          sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
          pip config set global.index-url http://${CACHING_URL}/pypi/simple
          pip config set global.trusted-host "${CACHING_URL}"

          bash scripts/ci/npu/npu_ci_install_dependency.sh a3 diffusion

      - name: Generate outputs
        env:
          PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-2npu }}
          SGLANG_USE_MODELSCOPE: true
          SGLANG_IS_IN_CI: true
          HF_ENDPOINT: https://hf-mirror.com
          TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
          PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
          STREAMS_PER_DEVICE: 32
        run: |
          cd python
          python -m sglang.multimodal_gen.test.scripts.gen_diffusion_ci_outputs \
            --suite 2-npu \
            --partition-id ${{ matrix.part }} \
            --total-partitions ${{ needs.compute-diffusion-partitions.outputs['partition-count-2npu'] }} \
            --partition-plan-json "$PARTITION_PLAN_JSON" \
            --out-dir ./${{ env.OUTPUT_NAME }} \
            ${{ inputs.case_ids != '' && format('--case-ids {0}', inputs.case_ids) || '' }}

      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
          name: ${{ env.OUTPUT_NAME }}-2npu-part${{ matrix.part }}
          path: python/${{ env.OUTPUT_NAME }}
          retention-days: 7

      - name: Publish GT images to sgl-project/ci-data-diffusion
        env:
          GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }}
        run: |
          python scripts/ci/utils/diffusion/publish_diffusion_gt.py \
            --source-dir python/${{ env.OUTPUT_NAME }} \
            --target-dir "${{ env.PUBLISH_TARGET_DIR }}"