Update handler.py

handler.py

@@ -7,8 +7,9 @@ import io
 import base64
 from io import BytesIO
 # from transformers import BlipForConditionalGeneration, BlipProcessor
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
-
+# from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from transformers import Blip2ForConditionalGeneration, AutoProcessor
+from peft import PeftModel, PeftConfig
 
 # -
 
@@ -18,10 +19,17 @@ class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
         print("####### Start Deploying #####")
-        self.processor = Blip2Processor.from_pretrained("ChirathD/Blip-2-test-1")
-        self.model = Blip2ForConditionalGeneration.from_pretrained("ChirathD/Blip-2-test-1") 
+        # self.processor = Blip2Processor.from_pretrained("ChirathD/Blip-2-test-1")
+        # self.model = Blip2ForConditionalGeneration.from_pretrained("ChirathD/Blip-2-test-1") 
         # self.model.eval()
         # self.model = self.model.to(device)
+
+        peft_model_id = "ChirathD/Blip-2-test-4"
+        config = PeftConfig.from_pretrained(peft_model_id)
+        
+        self.model = Blip2ForConditionalGeneration.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map="auto")
+        self.model = PeftModel.from_pretrained(self.model, peft_model_id)
+        self.processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")