Update app.py (#8)

- Update app.py (ea8e736e6853df72bf0f777bffbb4ff11260ebae)

app.py

@@ -1,15 +1,14 @@
-import subprocess
-# Installing flash_attn
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
 import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoProcessor
 import torch
 from PIL import Image
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 models = {
-    "microsoft/Phi-3.5-vision-instruct": AutoModelForCausalLM.from_pretrained("microsoft/Phi-3.5-vision-instruct", trust_remote_code=True, torch_dtype="auto").cuda().eval()
+    "microsoft/Phi-3.5-vision-instruct": AutoModelForCausalLM.from_pretrained("microsoft/Phi-3.5-vision-instruct", trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").cuda().eval()
+
 }
 
 processors = {