Upload handler.py

handler.py

@@ -15,8 +15,8 @@ class EndpointHandler():
         self.generate_model = Blip2ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16)
         self.generate_model.to(self.device)
 
-        self.feature_model =  Blip2Model.from_pretrained(path, torch_dtype=torch.float16)
-        self.feature_model.to(self.device)
+        #self.feature_model = Blip2Model.from_pretrained(path, torch_dtype=torch.float16)
+        #self.feature_model.to(self.device)
         
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -33,27 +33,28 @@ class EndpointHandler():
             prompt = inputs["prompt"]
         else:
             prompt = None
-        if "extract_feature" in inputs:
-            extract_feature = inputs["extract_feature"]
-        else:
-            extract_feature = False
+        #if "extract_feature" in inputs:
+        #    extract_feature = inputs["extract_feature"]
+        #else:
+        #    extract_feature = False
 
         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.generate_model.generate(**processed_image)
         generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
         result["image_caption"] = generated_text
-        if extract_feature:
-            caption_feature = self.feature_model(**processed_image)
-            result["caption_feature"] = caption_feature
+
+        #if extract_feature:
+        #    caption_feature = self.feature_model(**processed_image)
+        #    result["caption_feature"] = caption_feature
 
         if prompt:
             prompt_image_processed = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device, torch.float16)
             generated_ids = self.generate_model.generate(**prompt_image_processed)
             generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
             result["image_prompt"] = generated_text
-            if extract_feature:
-                prompt_feature = self.feature_model(**prompt_image_processed)
-                result["prompt_feature"] = prompt_feature
+            #if extract_feature:
+            #    prompt_feature = self.feature_model(**prompt_image_processed)
+            #    result["prompt_feature"] = prompt_feature
 
         return result