Update handler.py

handler.py

@@ -1,6 +1,7 @@
 from typing import Dict, Any
 from PIL import Image    
 import requests
+import flash-attn
 import torch
 import numpy as np
 from transformers import AutoProcessor, LlavaForConditionalGeneration, BitsAndBytesConfig
@@ -13,7 +14,7 @@ class EndpointHandler():
             model_id, 
             torch_dtype=torch.float16, 
             low_cpu_mem_usage=True,
-            load_in_4bit=True
+            use_flash_attention_2=True
         )
         self.processor = AutoProcessor.from_pretrained(model_id)