Upload handler.py

handler.py

@@ -1,4 +1,4 @@
-from transformers import Blip2Processor, Blip2Model
+from transformers import Blip2Processor, Blip2Model, Blip2ForConditionalGeneration
 from typing import Dict, List, Any
 from PIL import Image
 from transformers import pipeline
@@ -12,8 +12,8 @@ class EndpointHandler():
         """
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.processor = Blip2Processor.from_pretrained(path)
-        self.model = Blip2Model.from_pretrained(path, torch_dtype=torch.float16)
-        self.model.to(self.device)
+        self.generate_model = Blip2ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16)
+        self.generate_model.to(self.device)
         
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -27,6 +27,6 @@ class EndpointHandler():
         image_url = inputs['image_url']
         image = Image.open(requests.get(image_url, stream=True).raw)
         processed_image = self.processor(images=image, return_tensors="pt").to(self.device, torch.float16)
-        generated_ids = self.model.generate(**processed_image)
+        generated_ids = self.generate_model.generate(**processed_image)
         generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
         return image_url, generated_text