update

app.py

@@ -61,6 +61,13 @@ from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file, get_dir
 
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
+
 result = subprocess.run(
             ["bash", "setup.sh"],
             capture_output=True,

setup.sh

@@ -1,7 +1,4 @@
 #!/bin/bash
-pip install --upgrade pip setuptools wheel
-pip install --no-build-isolation flash-attn==2.5.9.post1
-
 git clone https://github.com/NVIDIA/apex
 cd apex
 # if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...