Upload handler.py

handler.py

@@ -1,4 +1,4 @@
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from transformers import Blip2Processor, Blip2Model
 from typing import Dict, List, Any
 from PIL import Image
 from transformers import pipeline
@@ -14,9 +14,12 @@ class EndpointHandler():
         # pseudo:
         # self.model= load_model(path)
         #self.processor = Blip2Processor.from_pretrained(path)
-        self.pipeline = pipeline(model = path)
+        #self.pipeline = pipeline(model = path)
+        self.path = path
         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)
+        
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
        data args:
@@ -33,4 +36,4 @@ class EndpointHandler():
         #generated_ids = self.pipeline(**inputs)
         #generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
 
-        return image_url
+        return image_url, self.path, self.device