#!/bin/bash
set -x

echo "=== Environment ==="
python3 -c "import torch; print(f'torch={torch.__version__}, cuda={torch.version.cuda}, abi={torch._C._GLIBCXX_USE_CXX11_ABI}')"

echo ""
echo "=== Method 1: pip install flash-attn with HF mirror index ==="
# The strangertoolshf repo on HF might have wheels
# Try using hf-mirror.com to access HuggingFace
HF_ENDPOINT=https://hf-mirror.com pip install flash-attn==2.8.3 \
    --no-build-isolation \
    --extra-index-url https://hf-mirror.com/strangertoolshf/flash_attention_2_wheelhouse/resolve/main/ \
    2>&1 | tail -20

python3 -c "import flash_attn; print(f'OK: {flash_attn.__version__}')" 2>&1 && exit 0

echo ""
echo "=== Method 2: Direct pip install from PyPI with build from source ==="
# flash-attn setup.py will try to download from GitHub
# But we can set environment variable to use a mirror
# Override the download URL in the setup.py by setting env
MAX_JOBS=16 pip install flash-attn==2.8.3 --no-build-isolation 2>&1 | tail -30

python3 -c "import flash_attn; print(f'OK: {flash_attn.__version__}')" 2>&1 && exit 0

echo ""
echo "=== Method 3: Build from HF source zip ==="
cd /tmp
rm -rf flash-attention-2.8.3*
wget -q --timeout=60 "https://hf-mirror.com/strangertoolshf/flash_attention_2_wheelhouse/resolve/main/flash-attention-2.8.3.zip" -O flash-attention-2.8.3.zip 2>&1
if [ -f flash-attention-2.8.3.zip ] && [ $(stat -c%s flash-attention-2.8.3.zip) -gt 100000 ]; then
    echo "Source zip downloaded, building..."
    unzip -q flash-attention-2.8.3.zip
    cd flash-attention-2.8.3
    MAX_JOBS=16 pip install . --no-build-isolation 2>&1 | tail -20
fi

python3 -c "import flash_attn; print(f'OK: {flash_attn.__version__}')" 2>&1

echo "HF_DOWNLOAD_DONE"