Update handler.py

handler.py

@@ -1,55 +1,55 @@
-import requests
-
-import torch
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM 
-
-class EndpointHandler():
-    def __init__(self, path=""):
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-        model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
-        processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
-
-
-    def predict_image(self, url, prompt):
-        url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
-        image = Image.open(requests.get(url, stream=True).raw)
-
-        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
-
-        generated_ids = model.generate(
-            input_ids=inputs["input_ids"],
-            pixel_values=inputs["pixel_values"],
-            max_new_tokens=4096,
-            num_beams=3,
-            do_sample=False
-        )
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-
-        parsed_answer = processor.post_process_generation(generated_text, task="<OD>", image_size=(image.width, image.height))
-        return parsed_answer
-        
-
-    def __call__(self, event):
-        if "url" not in event:
-            return {
-                "statusCode": 400,
-                "body": json.dumps("Error: Please provide an 'url' parameter."),
-            }
-
-        if "prompt" not in event:
-            return {
-                "statusCode": 400,
-                "body": json.dumps("Error: Please provide an 'prompt' parameter."),
-            }    
-
-        url = event["url"]
-        prompt = event["prompt"]
-        parsed_answer = self.predict_image(self, url, prompt)
-
-        return {
-                "statusCode": 200,
-                "body": json.dumps(parsed_answer),
+import requests
+import json
+import torch
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM 
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+        model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
+        processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
+
+
+    def predict_image(self, url, prompt):
+        url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
+        image = Image.open(requests.get(url, stream=True).raw)
+
+        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
+
+        generated_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=4096,
+            num_beams=3,
+            do_sample=False
+        )
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+
+        parsed_answer = processor.post_process_generation(generated_text, task="<OD>", image_size=(image.width, image.height))
+        return parsed_answer
+        
+
+    def __call__(self, event):
+        if "url" not in event:
+            return {
+                "statusCode": 400,
+                "body": json.dumps("Error: Please provide an 'url' parameter."),
+            }
+
+        if "prompt" not in event:
+            return {
+                "statusCode": 400,
+                "body": json.dumps("Error: Please provide an 'prompt' parameter."),
+            }    
+
+        url = event["url"]
+        prompt = event["prompt"]
+        parsed_answer = self.predict_image(self, url, prompt)
+
+        return {
+                "statusCode": 200,
+                "body": json.dumps(parsed_answer),
             }