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")