Update handler.py

handler.py

@@ -5,25 +5,13 @@ from PIL import Image
 import os
 import base64
 
-run("pip install flash-attn --no-build-isolation", shell=True, check=True)
+#run("pip install flash-attn --no-build-isolation", shell=True, check=True)
 run("pip install --upgrade pip", shell=True, check=True)
 run("pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu124", shell=True, check=True)
 
-#run("pip install --upgrade accelerate transformers", shell=True, check=True)
-#run("pip -qqq install --force-reinstall https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-0.44.2.dev0-py3-none-manylinux_2_24_x86_64.whl", shell=True, check=True)
 
 
 
-
-try:
-    import flash_attn
-    print("FlashAttention is installed")
-    USE_FLASH_ATTENTION = True
-except ImportError:
-    print("FlashAttention is not installed")
-    USE_FLASH_ATTENTION = False
-
-
 from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
 
 model_id = "ibm-granite/granite-vision-3.2-2b"
@@ -82,8 +70,7 @@ class PromptSet:
 class EndpointHandler():
     def __init__(self, path=""):
         self.model=AutoModelForVision2Seq.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16,
-                                                          quantization_config=bnb_config,
-                                                          _attn_implementation="flash_attention_2" if USE_FLASH_ATTENTION else None,)
+                                                          quantization_config=bnb_config)
         self.processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
 
     def __call__(self, data):