Update src/florence/model.py

src/florence/model.py

@@ -1,6 +1,6 @@
 # Importing necessary libraries
 import sys
-import subprocess
+# import subprocess
 from typing import Optional
 from PIL import Image
 import torch
@@ -14,17 +14,17 @@ from src.exception import CustomExceptionHandling
 
 
 # Install the required dependencies
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
-)
+# subprocess.run(
+#     "pip install flash-attn --no-build-isolation",
+#     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+#     shell=True,
+# )
 
 # Load model and processor from Hugging Face
 model_id = "microsoft/Florence-2-large-ft"
 try:
     model = (
-        AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, trust_remote_code=True)
+        AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="sdpa", torch_dtype=torch.float16, trust_remote_code=True)
         .to("cuda")
         .eval()
     )