Update handler.py

handler.py

@@ -13,10 +13,10 @@ class EndpointHandler():
         self.processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
 
 
-    def predict_image(self, url, prompt):
+    def predict_image(self, url, task, prompt):
         image = Image.open(requests.get(url, stream=True).raw)
 
-        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.device, self.torch_dtype)
+        inputs = self.processor(text=task + prompt, images=image, return_tensors="pt").to(self.device, self.torch_dtype)
 
         generated_ids = self.model.generate(
             input_ids=inputs["input_ids"],
@@ -27,7 +27,7 @@ class EndpointHandler():
         )
         generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
 
-        parsed_answer = self.processor.post_process_generation(generated_text, task="<OD>", image_size=(image.width, image.height))
+        parsed_answer = self.processor.post_process_generation(generated_text, task=task, image_size=(image.width, image.height))
         return parsed_answer
         
 
@@ -40,8 +40,9 @@ class EndpointHandler():
 
         inputs = event["inputs"]
         url = inputs["url"]
+        task = inputs["task"]
         prompt = inputs["prompt"]
-        parsed_answer = self.predict_image(url, prompt)
+        parsed_answer = self.predict_image(url, task, prompt)
 
         return {
                 "statusCode": 200,