Update handler.py

handler.py

@@ -5,7 +5,7 @@ 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)
 
@@ -14,12 +14,25 @@ run("pip install torch torchvision torchaudio --extra-index-url https://download
 
 
 
+
+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"
 
 bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
         llm_int8_skip_modules=["vision_tower", "lm_head"],
         llm_int8_enable_fp32_cpu_offload=True
 )
@@ -69,7 +82,8 @@ 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)
+                                                          quantization_config=bnb_config,
+                                                          _attn_implementation="flash_attention_2" if USE_FLASH_ATTENTION else None,)
         self.processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
 
     def __call__(self, data):