Update app.py

app.py

@@ -4,6 +4,8 @@ from PIL import Image
 import requests
 import gradio as gr
 import spaces
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 model_id = "yifeihu/TB-OCR-preview-0.1"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -20,7 +22,7 @@ processor = AutoProcessor.from_pretrained(model_id,
     num_crops=16
 )
 
-@spaces.GPU
+
 def phi_ocr(image):
     question = "Convert the text to markdown format."
     prompt_message = [{
@@ -40,6 +42,7 @@ def phi_ocr(image):
     response = response.split("<image_end>")[0]
     return response
 
+@spaces.GPU
 def process_image(input_image):
     return phi_ocr(input_image)