flash_attn

app.py

@@ -15,6 +15,10 @@ import requests
 from outlines import models, generate, samplers
 from pydantic import BaseModel
 
+# Install Flash attention
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
 # Constants
 MAX_IMAGE_SIZE = (1024, 1024)
 TARGET_IMAGE_SIZE = 180_000

requirements.txt

@@ -3,5 +3,4 @@ torch
 accelerate
 git+https://github.com/huggingface/transformers.git@main
 spaces
-torchvision
-https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.9.post1/flash_attn-2.5.9.post1+cu118torch1.12cxx11abiFALSE-cp310-cp310-linux_x86_64.whl 
+torchvision