| import requests | |
| import json | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoProcessor, AutoModelForCausalLM | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| self.device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
| self.model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=self.torch_dtype, trust_remote_code=True).to(self.device) | |
| self.processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True) | |
| def predict_image(self, url, task, prompt): | |
| image = Image.open(requests.get(url, stream=True).raw) | |
| 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"], | |
| pixel_values=inputs["pixel_values"], | |
| max_new_tokens=4096, | |
| num_beams=3, | |
| do_sample=False | |
| ) | |
| generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0] | |
| parsed_answer = self.processor.post_process_generation(generated_text, task=task, image_size=(image.width, image.height)) | |
| return parsed_answer | |
| def __call__(self, event): | |
| if "inputs" not in event: | |
| return { | |
| "statusCode": 400, | |
| "body": json.dumps("Error: Please provide an 'inputs' parameter."), | |
| } | |
| inputs = event["inputs"] | |
| url = inputs["url"] | |
| task = inputs["task"] | |
| prompt = inputs["prompt"] | |
| parsed_answer = self.predict_image(url, task, prompt) | |
| return { | |
| "statusCode": 200, | |
| "body": json.dumps(parsed_answer), | |
| } |