Upload 5 files

handler.py

@@ -0,0 +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),
+            }

requirements.py

@@ -0,0 +1,6 @@
+starlette
+transformers
+pillow
+torch
+timm
+einops

tokenizer_config.json

@@ -0,0 +1,3 @@
+{
+    "model_max_length": 1024
+}