Update app.py

app.py

@@ -16,7 +16,7 @@ from transformers import AutoModelForCausalLM, AutoProcessor
 from PIL import Image
 import requests
 import os
-os.system('pip install -U flash-attn')
+# os.system('pip install -U flash-attn')
 model_id = "yifeihu/TB-OCR-preview-0.1"
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -29,7 +29,7 @@ model = AutoModelForCausalLM.from_pretrained(
   device_map=DEVICE, 
   trust_remote_code=True, 
   torch_dtype="auto",
-  _attn_implementation='flash_attention_2',
+  # _attn_implementation='flash_attention_2',
   #load_in_4bit=True # Optional: Load model in 4-bit mode to save memory
 )