import subprocess
import sys

def install_flash_attn():
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "/home/user/app/wheels/flash_attn-2.8.2+cu129torch2.8-cp310-cp310-linux_x86_64.whl", "--no-build-isolation"])
        print(f"Successfully installed flash_attn")
    except subprocess.CalledProcessError as e:
        print(f"Error installing flash_attn: {e}. This demo won't work properly.")

try:
    import flash_attn
    print(f"`flash_attn` has been installed.")
except ImportError:
    print(f"`flash_attn` is NOT installed. Trying to install from local wheel...")
    install_flash_attn()