Update app.py

app.py

@@ -2,8 +2,16 @@ import gradio as gr
 import os
 import torch, torchvision, einops
 import spaces
+import flash_attention
 from transformers import AutoModelForCausalLM
 from huggingface_hub import login
+import subprocess
+
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 
 hf_token = os.getenv("HF_TOKEN")
 login(token=hf_token, add_to_git_credential=True)